Fixing our env.py to work with sqlalchemy 2.0 (#6030)

Author: kevgliss

src/dispatch/database/revisions/core/env.py

@@ -1,5 +1,5 @@
 from alembic import context
-from sqlalchemy import engine_from_config, pool, text
+from sqlalchemy import create_engine, text
 
 from dispatch.logging import logging
 from dispatch.config import SQLALCHEMY_DATABASE_URI
@@ -23,10 +23,8 @@
 
 def include_object(object, name, type_, reflected, compare_to):
     if type_ == "table":
-        if object.schema == CORE_SCHEMA_NAME:
-            return True
-    else:
-        return True
+        return object.schema == CORE_SCHEMA_NAME
+    return True
 
 
 def run_migrations_online():
@@ -36,28 +34,23 @@ def run_migrations_online():
     and associate a connection with the context.
 
     """
-
-    # don't create empty revisions
     def process_revision_directives(context, revision, directives):
         script = directives[0]
         if script.upgrade_ops.is_empty():
             directives[:] = []
             log.info("No changes found skipping revision creation.")
 
-    connectable = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy.", poolclass=pool.NullPool
-    )
+    connectable = create_engine(SQLALCHEMY_DATABASE_URI)
 
     log.info("Migrating dispatch core schema...")
     # migrate common tables
     with connectable.connect() as connection:
         set_search_path = text(f'set search_path to "{CORE_SCHEMA_NAME}"')
         connection.execute(set_search_path)
-        connection.dialect.default_schema_name = CORE_SCHEMA_NAME
+        connection.commit()
         context.configure(
             connection=connection,
             target_metadata=target_metadata,
-            include_schemas=True,
             include_object=include_object,
             process_revision_directives=process_revision_directives,
         )

src/dispatch/database/revisions/tenant/env.py

@@ -1,5 +1,5 @@
 from alembic import context
-from sqlalchemy import engine_from_config, pool, inspect, text
+from sqlalchemy import create_engine, inspect, text
 
 
 from dispatch.logging import logging
@@ -43,29 +43,26 @@ def run_migrations_online():
     and associate a connection with the context.
 
     """
-
     def process_revision_directives(context, revision, directives):
         script = directives[0]
         if script.upgrade_ops.is_empty():
             directives[:] = []
             log.info("No changes found skipping revision creation.")
 
-    connectable = engine_from_config(
-        config.get_section(config.config_ini_section), prefix="sqlalchemy.", poolclass=pool.NullPool
-    )
+    connectable = create_engine(SQLALCHEMY_DATABASE_URI)
 
     with connectable.connect() as connection:
         # get the schema names
         for schema in get_tenant_schemas(connection):
             log.info(f"Migrating {schema}...")
             set_search_path = text(f'set search_path to "{schema}"')
             connection.execute(set_search_path)
-            connection.dialect.default_schema_name = schema
+            connection.commit()
 
+            print(target_metadata)
             context.configure(
                 connection=connection,
                 target_metadata=target_metadata,
-                include_schemas=True,
                 include_object=include_object,
                 process_revision_directives=process_revision_directives,
             )