Fix bugs in migrations

alembic/versions/1fa46383a9d3_add_is_bot_field_to_puppets.py

@@ -21,4 +21,5 @@ def upgrade():
 
 
 def downgrade():
-    op.drop_column('puppet', 'is_bot')
+    with op.batch_alter_table("puppet") as batch_op:
+        batch_op.drop_column('is_bot')

alembic/versions/30eca60587f1_add_megagroup_field_to_portals.py

@@ -20,4 +20,5 @@ def upgrade():
 
 
 def downgrade():
-    op.drop_column('portal', 'megagroup')
+    with op.batch_alter_table("portal") as batch_op:
+        batch_op.drop_column('megagroup')

alembic/versions/6ca3d74d51e4_move_state_store_to_main_database.py

@@ -72,10 +72,17 @@ def upgrade():
                     sa.Column("avatar_url", sa.String(), nullable=True),
                     sa.PrimaryKeyConstraint("room_id", "user_id"))
 
+    try:
+        migrate_state_store()
+    except Exception as e:
+        print("Failed to migrate state store:", e)
+        print("Migrating the state store isn't required, but you can retry by alembic downgrading "
+              "to revision 2228d49c383f and upgrading again.")
+
+
+def migrate_state_store():
     conn = op.get_bind()
-    session = orm.sessionmaker(bind=conn)
-    session = orm.scoping.scoped_session(session)
-    Puppet.query = session.query_property()
+    session = orm.sessionmaker(bind=conn)()  # type: orm.Session
 
     try:
         with open("mx-state.json") as file:
@@ -99,7 +106,7 @@ def upgrade():
         if not match:
             continue
 
-        puppet = Puppet.query.get(match.group(1))
+        puppet = session.query(Puppet).get(match.group(1))
         if not puppet:
             continue
 

alembic/versions/7d47d84380b6_add_timestamp_to_telegramfile.py

@@ -22,4 +22,5 @@ def upgrade():
 
 
 def downgrade():
-    op.drop_column('telegram_file', 'timestamp')
+    with op.batch_alter_table("telegram_file") as batch_op:
+        batch_op.drop_column('timestamp')

alembic/versions/bcfefa1f1299_add_displayname_source_fields_for_.py

@@ -20,4 +20,5 @@ def upgrade():
 
 
 def downgrade():
-    op.drop_column('puppet', 'displayname_source')
+    with op.batch_alter_table("puppet") as batch_op:
+        batch_op.drop_column('displayname_source')