From 0f05a7993632f4c1ee560ef460a8450b7fdecf12 Mon Sep 17 00:00:00 2001 From: Kevin Glisson Date: Mon, 2 Jun 2025 13:41:27 -0700 Subject: [PATCH 1/2] Fixing our env.py to work with sqlalchemy 2.0 --- src/dispatch/database/revisions/core/env.py | 17 +++++------------ src/dispatch/database/revisions/tenant/env.py | 10 +++------- 2 files changed, 8 insertions(+), 19 deletions(-) diff --git a/src/dispatch/database/revisions/core/env.py b/src/dispatch/database/revisions/core/env.py index 353c7f2e8560..457d52aa2079 100644 --- a/src/dispatch/database/revisions/core/env.py +++ b/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, inspect, 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, ) diff --git a/src/dispatch/database/revisions/tenant/env.py b/src/dispatch/database/revisions/tenant/env.py index 937dbfc06636..fda4effe7689 100644 --- a/src/dispatch/database/revisions/tenant/env.py +++ b/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,16 +43,13 @@ 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 @@ -60,12 +57,11 @@ def process_revision_directives(context, revision, directives): 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() context.configure( connection=connection, target_metadata=target_metadata, - include_schemas=True, include_object=include_object, process_revision_directives=process_revision_directives, ) From 3e34edeb1bc67d2366b850634761bca0013e17b7 Mon Sep 17 00:00:00 2001 From: Kevin Glisson Date: Mon, 2 Jun 2025 13:57:00 -0700 Subject: [PATCH 2/2] Fixing ruff --- src/dispatch/database/revisions/core/env.py | 2 +- src/dispatch/database/revisions/tenant/env.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/dispatch/database/revisions/core/env.py b/src/dispatch/database/revisions/core/env.py index 457d52aa2079..0ae70a6c01f2 100644 --- a/src/dispatch/database/revisions/core/env.py +++ b/src/dispatch/database/revisions/core/env.py @@ -1,5 +1,5 @@ from alembic import context -from sqlalchemy import create_engine, inspect, text +from sqlalchemy import create_engine, text from dispatch.logging import logging from dispatch.config import SQLALCHEMY_DATABASE_URI diff --git a/src/dispatch/database/revisions/tenant/env.py b/src/dispatch/database/revisions/tenant/env.py index fda4effe7689..da533ee9f1ce 100644 --- a/src/dispatch/database/revisions/tenant/env.py +++ b/src/dispatch/database/revisions/tenant/env.py @@ -59,6 +59,7 @@ def process_revision_directives(context, revision, directives): connection.execute(set_search_path) connection.commit() + print(target_metadata) context.configure( connection=connection, target_metadata=target_metadata,