From c1d4e8e48289bcd20f3cc12cabe22bc3b30d442f Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Tue, 31 Mar 2020 22:19:43 +0300
Subject: [PATCH] Update mautrix-python to use SQLAlchemy for matrix-nio state
 storage

---
 ...d_add_matrix_nio_state_store_to_main_db.py | 71 +++++++++++++++++++
 mautrix_telegram/db/__init__.py               |  7 ++
 requirements.txt                              |  2 +-
 3 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 alembic/versions/dff56c93da8d_add_matrix_nio_state_store_to_main_db.py

diff --git a/alembic/versions/dff56c93da8d_add_matrix_nio_state_store_to_main_db.py b/alembic/versions/dff56c93da8d_add_matrix_nio_state_store_to_main_db.py
new file mode 100644
index 00000000..6a16f2dc
--- /dev/null
+++ b/alembic/versions/dff56c93da8d_add_matrix_nio_state_store_to_main_db.py
@@ -0,0 +1,71 @@
+"""Add matrix-nio state store to main db
+
+Revision ID: dff56c93da8d
+Revises: d3c922a6acd2
+Create Date: 2020-03-31 22:04:04.014048
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'dff56c93da8d'
+down_revision = 'd3c922a6acd2'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('nio_account',
+    sa.Column('user_id', sa.String(length=255), nullable=False),
+    sa.Column('device_id', sa.String(length=255), nullable=False),
+    sa.Column('shared', sa.Boolean(), nullable=False),
+    sa.Column('sync_token', sa.Text(), nullable=False),
+    sa.Column('account', sa.LargeBinary(), nullable=False),
+    sa.PrimaryKeyConstraint('user_id', 'device_id')
+    )
+    op.create_table('nio_device_key',
+    sa.Column('user_id', sa.String(length=255), nullable=False),
+    sa.Column('device_id', sa.String(length=255), nullable=False),
+    sa.Column('display_name', sa.String(length=255), nullable=False),
+    sa.Column('deleted', sa.Boolean(), nullable=False),
+    sa.Column('keys', sa.PickleType(), nullable=False),
+    sa.PrimaryKeyConstraint('user_id', 'device_id')
+    )
+    op.create_table('nio_megolm_inbound_session',
+    sa.Column('session_id', sa.String(length=255), nullable=False),
+    sa.Column('sender_key', sa.String(length=255), nullable=False),
+    sa.Column('fp_key', sa.String(length=255), nullable=False),
+    sa.Column('room_id', sa.String(length=255), nullable=False),
+    sa.Column('session', sa.LargeBinary(), nullable=False),
+    sa.Column('forwarded_chains', sa.PickleType(), nullable=False),
+    sa.PrimaryKeyConstraint('session_id')
+    )
+    op.create_table('nio_olm_session',
+    sa.Column('session_id', sa.String(length=255), nullable=False),
+    sa.Column('sender_key', sa.String(length=255), nullable=False),
+    sa.Column('session', sa.LargeBinary(), nullable=False),
+    sa.Column('created_at', sa.DateTime(), nullable=False),
+    sa.Column('last_used', sa.DateTime(), nullable=False),
+    sa.PrimaryKeyConstraint('session_id')
+    )
+    op.create_table('nio_outgoing_key_request',
+    sa.Column('request_id', sa.String(length=255), nullable=False),
+    sa.Column('session_id', sa.String(length=255), nullable=False),
+    sa.Column('room_id', sa.String(length=255), nullable=False),
+    sa.Column('algorithm', sa.String(length=255), nullable=False),
+    sa.PrimaryKeyConstraint('request_id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('nio_outgoing_key_request')
+    op.drop_table('nio_olm_session')
+    op.drop_table('nio_megolm_inbound_session')
+    op.drop_table('nio_device_key')
+    op.drop_table('nio_account')
+    # ### end Alembic commands ###
diff --git a/mautrix_telegram/db/__init__.py b/mautrix_telegram/db/__init__.py
index 92a824f0..0725d6ae 100644
--- a/mautrix_telegram/db/__init__.py
+++ b/mautrix_telegram/db/__init__.py
@@ -24,6 +24,11 @@ from .puppet import Puppet
 from .telegram_file import TelegramFile
 from .user import User, UserPortal, Contact
 
+try:
+    from mautrix.bridge.db.nio_state_store import init as init_nio_db
+except ImportError:
+    init_nio_db = None
+
 
 def init(db_engine: Engine) -> None:
     for table in (Portal, Message, User, Contact, UserPortal, Puppet, TelegramFile, UserProfile,
@@ -32,3 +37,5 @@ def init(db_engine: Engine) -> None:
         table.t = table.__table__
         table.c = table.t.c
         table.column_names = table.c.keys()
+    if init_nio_db:
+        init_nio_db(db_engine)
diff --git a/requirements.txt b/requirements.txt
index a64521f8..feaf4bb8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,6 +4,6 @@ ruamel.yaml>=0.15.35,<0.17
 python-magic>=0.4,<0.5
 commonmark>=0.8,<0.10
 aiohttp>=3,<4
-mautrix==0.5.0.beta9
+mautrix==0.5.0.beta10
 telethon>=1.10,<1.12
 telethon-session-sqlalchemy>=0.2.14,<0.3