Getting closer

alembic/env.py

@@ -4,18 +4,13 @@ import asyncio
 from logging.config import fileConfig
 
 from alembic import context
-from sqlalchemy import pool
+from sqlalchemy import pool, text
 from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import async_engine_from_config
 
 # Import your models' Base to enable autogenerate
 from alpine_bits_python.config_loader import load_config
-from alpine_bits_python.db import (
-    Base,
-    configure_schema,
-    get_database_schema,
-    get_database_url,
-)
+from alpine_bits_python.db import Base, get_database_schema, get_database_url
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
@@ -39,14 +34,16 @@ if db_url:
     config.set_main_option("sqlalchemy.url", db_url)
 
 # Get schema name from application config
-schema_name = get_database_schema(app_config)
-if schema_name:
-    # Configure schema for all tables before migrations
-    configure_schema(schema_name)
+SCHEMA = get_database_schema(app_config)
 
 # add your model's MetaData object here for 'autogenerate' support
 target_metadata = Base.metadata
 
+# Configure metadata to resolve unqualified table names in the schema
+# This is needed so ForeignKey("customers.id") can find "alpinebits.customers"
+if SCHEMA:
+    target_metadata.schema = SCHEMA
+
 
 def run_migrations_offline() -> None:
     """Run migrations in 'offline' mode.
@@ -65,11 +62,13 @@ def run_migrations_offline() -> None:
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
-        version_table_schema=schema_name,  # Store alembic_version in our schema
-        include_schemas=True,
     )
 
     with context.begin_transaction():
+        # Set search path for offline mode if schema is configured
+        if SCHEMA:
+            print(f"Setting search_path to {SCHEMA}, public")
+            context.execute(f"SET search_path TO {SCHEMA}, public")
         context.run_migrations()
 
 
@@ -78,37 +77,38 @@ def do_run_migrations(connection: Connection) -> None:
     context.configure(
         connection=connection,
         target_metadata=target_metadata,
-        version_table_schema=schema_name,  # Store alembic_version in our schema
-        include_schemas=True,  # Allow Alembic to work with non-default schemas
     )
 
     with context.begin_transaction():
+        # Create schema if it doesn't exist
+        if SCHEMA:
+            connection.execute(text(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA}"))
+            # Set search path to our schema
+            print(f"setting search path to schema {SCHEMA}, ")
+            connection.execute(text(f"SET search_path TO {SCHEMA}"))
         context.run_migrations()
 
 
 async def run_async_migrations() -> None:
-    """Run migrations in 'online' mode using async engine.
-
-    In this scenario we need to create an Engine
+    """In this scenario we need to create an Engine
     and associate a connection with the context.
+
     """
-    # Get the config section for sqlalchemy settings
-    configuration = config.get_section(config.config_ini_section, {})
-
-    # Add connect_args for PostgreSQL schema support if needed
-    if schema_name and "postgresql" in configuration.get("sqlalchemy.url", ""):
-        configuration["connect_args"] = {
-            "server_settings": {"search_path": f"{schema_name},public"}
-        }
-
-    # Create async engine
     connectable = async_engine_from_config(
-        configuration,
+        config.get_section(config.config_ini_section, {}),
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
     )
 
     async with connectable.connect() as connection:
+        if connection.dialect.name == "postgresql":
+            # set search path on the connection, which ensures that
+            # PostgreSQL will emit all CREATE / ALTER / DROP statements
+            # in terms of this schema by default
+
+            connection.execute(text(f"SET search_path TO {SCHEMA}"))
+            # in SQLAlchemy v2+ the search path change needs to be committed
+            connection.commit()
         await connection.run_sync(do_run_migrations)
 
     await connectable.dispose()