From 17aefd02da44775949da6d5d38e6c264053d0a42 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Fri, 22 Jun 2018 21:02:37 +0300
Subject: [PATCH] Make alembic result consistent with definitions in db.py and
 add bot_id to bot_chat table

---
 alembic/env.py                                | 13 ++++++
 .../6e8b438b5e84_add_bot_id_to_botchat.py     | 46 +++++++++++++++++++
 ...7d84380b6_add_timestamp_to_telegramfile.py |  2 +-
 mautrix_telegram/db.py                        |  7 +--
 4 files changed, 64 insertions(+), 4 deletions(-)
 create mode 100644 alembic/versions/6e8b438b5e84_add_bot_id_to_botchat.py

diff --git a/alembic/env.py b/alembic/env.py
index 58ebcf51..5da5e604 100644
--- a/alembic/env.py
+++ b/alembic/env.py
@@ -4,10 +4,12 @@ from logging.config import fileConfig
 
 import sys
 from os.path import abspath, dirname
+
 sys.path.insert(0, dirname(dirname(abspath(__file__))))
 
 from mautrix_telegram.base import Base
 from mautrix_telegram.config import Config
+from alchemysession import AlchemySessionContainer
 import mautrix_telegram.db
 
 # this is the Alembic Config object, which provides
@@ -20,6 +22,15 @@ mxtg_config.load()
 config.set_main_option("sqlalchemy.url",
                        mxtg_config.get("appservice.database", "sqlite:///mautrix-telegram.db"))
 
+
+class FakeDB:
+    @staticmethod
+    def query_property():
+        return None
+
+
+AlchemySessionContainer.create_table_classes(FakeDB(), "telethon_", Base)
+
 # Interpret the config file for Python logging.
 # This line sets up loggers basically.
 fileConfig(config.config_file_name)
@@ -30,6 +41,7 @@ fileConfig(config.config_file_name)
 # target_metadata = mymodel.Base.metadata
 target_metadata = Base.metadata
 
+
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
 # my_important_option = config.get_main_option("my_important_option")
@@ -77,6 +89,7 @@ def run_migrations_online():
         with context.begin_transaction():
             context.run_migrations()
 
+
 if context.is_offline_mode():
     run_migrations_offline()
 else:
diff --git a/alembic/versions/6e8b438b5e84_add_bot_id_to_botchat.py b/alembic/versions/6e8b438b5e84_add_bot_id_to_botchat.py
new file mode 100644
index 00000000..dcd73f24
--- /dev/null
+++ b/alembic/versions/6e8b438b5e84_add_bot_id_to_botchat.py
@@ -0,0 +1,46 @@
+"""Add bot_id to BotChat
+
+Revision ID: 6e8b438b5e84
+Revises: 2228d49c383f
+Create Date: 2018-06-22 16:32:31.922480
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = "6e8b438b5e84"
+down_revision = "2228d49c383f"
+branch_labels = None
+depends_on = None
+
+convention = {
+    "pk": "pk_%(table_name)s"
+}
+
+metadata = sa.MetaData(naming_convention=convention)
+
+
+def upgrade():
+    conn = op.get_bind()
+    res = conn.execute("SELECT id, `type` FROM bot_chat")
+    results = res.fetchall()
+    op.drop_table("bot_chat")
+    bot_chat = op.create_table('bot_chat',
+                               sa.Column('bot_id', sa.Integer),
+                               sa.Column('chat_id', sa.Integer),
+                               sa.Column('type', sa.String, nullable=False),
+                               sa.PrimaryKeyConstraint('bot_id', 'chat_id'))
+    op.bulk_insert(bot_chat, [{"bot_id": 0, "chat_id": r[0], "type": r[1]} for r in results])
+
+
+def downgrade():
+    conn = op.get_bind()
+    res = conn.execute("SELECT chat_id, `type` FROM bot_chat WHERE bot_id=0")
+    results = res.fetchall()
+    op.drop_table("bot_chat")
+    bot_chat = op.create_table('bot_chat',
+                               sa.Column('id', sa.Integer),
+                               sa.Column('type', sa.String, nullable=False),
+                               sa.PrimaryKeyConstraint('id'))
+    op.bulk_insert(bot_chat, [{"id": r[0], "type": r[1]} for r in results])
diff --git a/alembic/versions/7d47d84380b6_add_timestamp_to_telegramfile.py b/alembic/versions/7d47d84380b6_add_timestamp_to_telegramfile.py
index 9c7176ac..0a0d3cd9 100644
--- a/alembic/versions/7d47d84380b6_add_timestamp_to_telegramfile.py
+++ b/alembic/versions/7d47d84380b6_add_timestamp_to_telegramfile.py
@@ -17,7 +17,7 @@ depends_on = None
 
 def upgrade():
     op.add_column('telegram_file',
-                  sa.Column('timestamp', sa.BigInteger(), nullable=False, default=0,
+                  sa.Column('timestamp', sa.BigInteger(), nullable=True, default=0,
                             server_default="0"))
 
 
diff --git a/mautrix_telegram/db.py b/mautrix_telegram/db.py
index e56c95ed..1ce3b604 100644
--- a/mautrix_telegram/db.py
+++ b/mautrix_telegram/db.py
@@ -28,7 +28,7 @@ class Portal(Base):
     # Telegram chat information
     tgid = Column(Integer, primary_key=True)
     tg_receiver = Column(Integer, primary_key=True)
-    peer_type = Column(String)
+    peer_type = Column(String, nullable=False)
     megagroup = Column(Boolean)
 
     # Matrix portal information
@@ -74,7 +74,7 @@ class User(Base):
     mxid = Column(String, primary_key=True)
     tgid = Column(Integer, nullable=True, unique=True)
     tg_username = Column(String, nullable=True)
-    saved_contacts = Column(Integer, default=0)
+    saved_contacts = Column(Integer, default=0, nullable=False)
     contacts = relationship("Contact", uselist=True,
                             cascade="save-update, merge, delete, delete-orphan")
     portals = relationship("Portal", secondary="user_portal")
@@ -104,7 +104,8 @@ class Puppet(Base):
 class BotChat(Base):
     query = None
     __tablename__ = "bot_chat"
-    id = Column(Integer, primary_key=True)
+    bot_id = Column(Integer, primary_key=True, default=0)
+    chat_id = Column(Integer, primary_key=True)
     type = Column(String, nullable=False)