feat!: allow override of registry for declarative bases (#307)

* feat: allow easy override of `registry` for declarative bases

* feat: implement a metadata registry

* feat: remove the need for `AlembicCommand` subclass in Litestar

Author: cofin

advanced_alchemy/alembic/commands.py

@@ -38,6 +38,7 @@ def __init__(
         self,
         engine: Engine | AsyncEngine,
         version_table_name: str,
+        bind_key: str | None = None,
         file_: str | os.PathLike[str] | None = None,
         ini_section: str = "alembic",
         output_buffer: TextIO | None = None,
@@ -56,6 +57,7 @@ def __init__(
         Args:
             engine (sqlalchemy.engine.Engine | sqlalchemy.ext.asyncio.AsyncEngine): The SQLAlchemy engine instance.
             version_table_name (str): The name of the version table.
+            bind_key (str | None): The bind key for the metadata.
             file_ (str | os.PathLike[str] | None): The file path for the alembic configuration.
             ini_section (str): The ini section name.
             output_buffer (typing.TextIO | None): The output buffer for alembic commands.
@@ -70,6 +72,7 @@ def __init__(
             user_module_prefix (str | None): The prefix for user modules.
         """
         self.template_directory = template_directory
+        self.bind_key = bind_key
         self.version_table_name = version_table_name
         self.version_table_pk = engine.dialect.name != "spanner+spanner"
         self.version_table_schema = version_table_schema

advanced_alchemy/alembic/templates/asyncio/env.py

@@ -6,7 +6,7 @@
 from sqlalchemy import Column, pool
 from sqlalchemy.ext.asyncio import AsyncEngine, async_engine_from_config
 
-from advanced_alchemy.base import orm_registry
+from advanced_alchemy.base import metadata_registry
 from alembic import context
 from alembic.autogenerate import rewriter
 from alembic.operations import ops
@@ -23,16 +23,6 @@
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config: AlembicCommandConfig = context.config  # type: ignore  # noqa: PGH003
-
-
-# add your model's MetaData object here
-# for 'autogenerate' support
-target_metadata = orm_registry.metadata
-
-# other values from the config, defined by the needs of env.py,
-# can be acquired:
-# ... etc.
-
 writer = rewriter.Rewriter()
 
 
@@ -75,7 +65,7 @@ def run_migrations_offline() -> None:
     """
     context.configure(
         url=config.db_url,
-        target_metadata=target_metadata,
+        target_metadata=metadata_registry.get(config.bind_key),
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
         compare_type=config.compare_type,
@@ -94,7 +84,7 @@ def do_run_migrations(connection: Connection) -> None:
     """Run migrations."""
     context.configure(
         connection=connection,
-        target_metadata=target_metadata,
+        target_metadata=metadata_registry.get(config.bind_key),
         compare_type=config.compare_type,
         version_table=config.version_table_name,
         version_table_pk=config.version_table_pk,

advanced_alchemy/alembic/templates/sync/env.py

@@ -4,7 +4,7 @@
 
 from sqlalchemy import Column, Engine, engine_from_config, pool
 
-from advanced_alchemy.base import orm_registry
+from advanced_alchemy.base import metadata_registry
 from alembic import context
 from alembic.autogenerate import rewriter
 from alembic.operations import ops
@@ -21,16 +21,6 @@
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config: AlembicCommandConfig = context.config  # type: ignore  # noqa: PGH003
-
-
-# add your model's MetaData object here
-# for 'autogenerate' support
-target_metadata = orm_registry.metadata
-
-# other values from the config, defined by the needs of env.py,
-# can be acquired:
-# ... etc.
-
 writer = rewriter.Rewriter()
 
 
@@ -73,7 +63,7 @@ def run_migrations_offline() -> None:
     """
     context.configure(
         url=config.db_url,
-        target_metadata=target_metadata,
+        target_metadata=metadata_registry.get(config.bind_key),
         literal_binds=True,
         dialect_opts={"paramstyle": "named"},
         compare_type=config.compare_type,
@@ -92,7 +82,7 @@ def do_run_migrations(connection: Connection) -> None:
     """Run migrations."""
     context.configure(
         connection=connection,
-        target_metadata=target_metadata,
+        target_metadata=metadata_registry.get(config.bind_key),
         compare_type=config.compare_type,
         version_table=config.version_table_name,
         version_table_pk=config.version_table_pk,