From 09108f6a730596a4936afb7e9d2152ce2ea47270 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sat, 11 Jul 2020 16:22:27 +0300
Subject: [PATCH] Add device_id to client db

---
 alembic/env.py                                |  4 ++-
 .../4b93300852aa_add_device_id_to_clients.py  | 32 +++++++++++++++++++
 maubot/client.py                              |  6 +++-
 maubot/db.py                                  |  3 +-
 maubot/lib/store_proxy.py                     | 10 +++---
 requirements.txt                              |  2 +-
 6 files changed, 47 insertions(+), 10 deletions(-)
 create mode 100644 alembic/versions/4b93300852aa_add_device_id_to_clients.py

diff --git a/alembic/env.py b/alembic/env.py
index 51f01c2..9946810 100644
--- a/alembic/env.py
+++ b/alembic/env.py
@@ -56,6 +56,7 @@ def run_migrations_offline():
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
+        render_as_batch=True,
     )
 
     with context.begin_transaction():
@@ -77,7 +78,8 @@ def run_migrations_online():
 
     with connectable.connect() as connection:
         context.configure(
-            connection=connection, target_metadata=target_metadata
+            connection=connection, target_metadata=target_metadata,
+            render_as_batch=True,
         )
 
         with context.begin_transaction():
diff --git a/alembic/versions/4b93300852aa_add_device_id_to_clients.py b/alembic/versions/4b93300852aa_add_device_id_to_clients.py
new file mode 100644
index 0000000..efc71cd
--- /dev/null
+++ b/alembic/versions/4b93300852aa_add_device_id_to_clients.py
@@ -0,0 +1,32 @@
+"""Add device_id to clients
+
+Revision ID: 4b93300852aa
+Revises: fccd1f95544d
+Create Date: 2020-07-11 15:49:38.831459
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '4b93300852aa'
+down_revision = 'fccd1f95544d'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('client', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('device_id', sa.String(length=255), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('client', schema=None) as batch_op:
+        batch_op.drop_column('device_id')
+
+    # ### end Alembic commands ###
diff --git a/maubot/client.py b/maubot/client.py
index 1c2ccc9..52d88dc 100644
--- a/maubot/client.py
+++ b/maubot/client.py
@@ -60,7 +60,7 @@ class Client:
         self.remote_avatar_url = None
         self.client = MaubotMatrixClient(mxid=self.id, base_url=self.homeserver,
                                          token=self.access_token, client_session=self.http_client,
-                                         log=self.log, loop=self.loop,
+                                         log=self.log, loop=self.loop, device_id=self.device_id,
                                          store=ClientStoreProxy(self.db_instance))
         self.client.ignore_initial_sync = True
         self.client.ignore_first_sync = True
@@ -273,6 +273,10 @@ class Client:
     def access_token(self) -> str:
         return self.db_instance.access_token
 
+    @property
+    def device_id(self) -> str:
+        return self.db_instance.device_id
+
     @property
     def enabled(self) -> bool:
         return self.db_instance.enabled
diff --git a/maubot/db.py b/maubot/db.py
index a9928d2..1a6b9fb 100644
--- a/maubot/db.py
+++ b/maubot/db.py
@@ -21,7 +21,7 @@ from sqlalchemy import Column, String, Boolean, ForeignKey, Text
 from sqlalchemy.engine.base import Engine
 import sqlalchemy as sql
 
-from mautrix.types import UserID, FilterID, SyncToken, ContentURI
+from mautrix.types import UserID, FilterID, DeviceID, SyncToken, ContentURI
 from mautrix.util.db import Base
 
 from .config import Config
@@ -53,6 +53,7 @@ class DBClient(Base):
     id: UserID = Column(String(255), primary_key=True)
     homeserver: str = Column(String(255), nullable=False)
     access_token: str = Column(Text, nullable=False)
+    device_id: DeviceID = Column(String(255), nullable=True)
     enabled: bool = Column(Boolean, nullable=False, default=False)
 
     next_batch: SyncToken = Column(String(255), nullable=False, default="")
diff --git a/maubot/lib/store_proxy.py b/maubot/lib/store_proxy.py
index 83bec7c..e9e3696 100644
--- a/maubot/lib/store_proxy.py
+++ b/maubot/lib/store_proxy.py
@@ -21,10 +21,8 @@ class ClientStoreProxy(ClientStore):
     def __init__(self, db_instance) -> None:
         self.db_instance = db_instance
 
-    @property
-    def next_batch(self) -> SyncToken:
-        return self.db_instance.next_batch
+    async def put_next_batch(self, next_batch: SyncToken) -> None:
+        self.db_instance.edit(next_batch=next_batch)
 
-    @next_batch.setter
-    def next_batch(self, value: SyncToken) -> None:
-        self.db_instance.edit(next_batch=value)
+    async def get_next_batch(self) -> SyncToken:
+        return self.db_instance.next_batch
diff --git a/requirements.txt b/requirements.txt
index fceea7b..e37ad38 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-mautrix==0.5.7
+mautrix==0.6.0.beta7
 aiohttp>=3,<4
 SQLAlchemy>=1,<2
 alembic>=1,<2