fix: Support Alembic migration autogeneration (#40)

isaacvando · web-flow · commit e44e22a3a80c · 2026-04-08T14:01:39.000-05:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file. The format
 
 ## Unreleased
 
+## [0.3.0] - 2026-04-07
+
+### Fixed
+
+- Fixed support for auto-generating migrations using Alembic.
+
 ## [0.2.0] - 2026-04-01
 
 ### Added
@@ -22,5 +28,6 @@ All notable changes to this project will be documented in this file. The format
 - CI workflow for lint, typing, unit, and integration checks.
 - Example scripts for quickstart, facets, autocomplete, MLT, hybrid RRF, and RAG retrieval.
 
+[0.3.0]: https://github.com/paradedb/sqlalchemy-paradedb/releases/tag/v0.3.0
 [0.2.0]: https://github.com/paradedb/sqlalchemy-paradedb/releases/tag/v0.2.0
 [0.1.0]: https://github.com/paradedb/sqlalchemy-paradedb/releases/tag/v0.1.0
diff --git a/paradedb/sqlalchemy/alembic.py b/paradedb/sqlalchemy/alembic.py
@@ -66,6 +66,9 @@ def create_bm25_index(
             )
         )
 
+    def reverse(self) -> MigrateOperation:
+        return DropBM25IndexOp(index_name=self.index_name, if_exists=True, schema=self.table_schema)
+
 
 @Operations.implementation_for(CreateBM25IndexOp)
 def _create_bm25_index_impl(operations: Operations, operation: CreateBM25IndexOp) -> None:
@@ -97,16 +100,62 @@ def _render_create_bm25_index_op(autogen_context, op: CreateBM25IndexOp) -> str:
 
 @Operations.register_operation("drop_bm25_index")
 class DropBM25IndexOp(MigrateOperation):
-    def __init__(self, index_name: str, if_exists: bool = True, schema: str | None = None) -> None:
+    def __init__(
+        self,
+        index_name: str,
+        if_exists: bool = True,
+        schema: str | None = None,
+        *,
+        table_name: str | None = None,
+        expressions: list[str] | None = None,
+        key_field: str | None = None,
+        where: str | None = None,
+    ) -> None:
         self.index_name = index_name
         self.if_exists = if_exists
         self.schema = schema
+        self.table_name = table_name
+        self.expressions = expressions
+        self.key_field = key_field
+        self.where = where
 
     @classmethod
     def drop_bm25_index(
-        cls, operations: Operations, index_name: str, if_exists: bool = True, schema: str | None = None
+        cls,
+        operations: Operations,
+        index_name: str,
+        if_exists: bool = True,
+        schema: str | None = None,
+        *,
+        table_name: str | None = None,
+        expressions: list[str] | None = None,
+        key_field: str | None = None,
+        where: str | None = None,
     ) -> MigrateOperation:
-        return operations.invoke(cls(index_name=index_name, if_exists=if_exists, schema=schema))
+        return operations.invoke(
+            cls(
+                index_name=index_name,
+                if_exists=if_exists,
+                schema=schema,
+                table_name=table_name,
+                expressions=expressions,
+                key_field=key_field,
+                where=where,
+            )
+        )
+
+    def reverse(self) -> MigrateOperation:
+        if self.table_name is None or self.expressions is None or self.key_field is None:
+            raise NotImplementedError("DropBM25IndexOp requires recreate metadata for Alembic downgrade generation")
+
+        return CreateBM25IndexOp(
+            index_name=self.index_name,
+            table_name=self.table_name,
+            expressions=self.expressions,
+            key_field=self.key_field,
+            table_schema=self.schema,
+            where=self.where,
+        )
 
 
 @Operations.implementation_for(DropBM25IndexOp)
@@ -120,6 +169,14 @@ def _render_drop_bm25_index_op(autogen_context, op: DropBM25IndexOp) -> str:
     parts = [repr(op.index_name), f"if_exists={op.if_exists!r}"]
     if op.schema is not None:
         parts.append(f"schema={op.schema!r}")
+    if op.table_name is not None:
+        parts.append(f"table_name={op.table_name!r}")
+    if op.expressions is not None:
+        parts.append(f"expressions={op.expressions!r}")
+    if op.key_field is not None:
+        parts.append(f"key_field={op.key_field!r}")
+    if op.where is not None:
+        parts.append(f"where={op.where!r}")
     return f"op.drop_bm25_index({', '.join(parts)})"
 
 
@@ -204,12 +261,28 @@ def _render_bm25_expression(expr: ClauseElement) -> str:
     return str(expr.compile(dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True}))  # type: ignore[no-untyped-call]
 
 
-def _strip_relation_qualifiers(expr: str, table_name: str) -> str:
-    # SQLAlchemy may render column refs as `table.col` in metadata compilation;
-    # CREATE INDEX field lists should be table-local expressions.
-    stripped = expr.replace(f'"{table_name}".', "")
-    stripped = stripped.replace(f"{table_name}.", "")
-    return stripped
+def _strip_relation_qualifiers(expr: str, table_name: str, schema_name: str | None) -> str:
+    # SQLAlchemy may render column refs as `table.col` or `schema.col` in metadata
+    # compilation; CREATE INDEX field lists should be table-local expressions.
+    qualifier_patterns: list[str] = []
+    for name in (schema_name, table_name):
+        if not name:
+            continue
+        qualifier_patterns.append(re.escape(_quote_ident(name)))
+        qualifier_patterns.append(rf"(?<![\w\"]){re.escape(name)}(?![\w\"])")
+
+    if not qualifier_patterns:
+        return expr
+
+    qualifier_re = re.compile(rf"('(?:''|[^'])*')|(?P<qualifier>(?:{'|'.join(qualifier_patterns)}))\.")
+
+    def _strip_match(match: re.Match[str]) -> str:
+        literal = match.group(1)
+        if literal is not None:
+            return literal
+        return ""
+
+    return qualifier_re.sub(_strip_match, expr)
 
 
 def _normalize_bm25_expression(expr: str) -> str:
@@ -312,8 +385,9 @@ def _render_where_from_index(index) -> str | None:
                 )
             ),
             index.table.name,
+            index.table.schema,
         )
-    return _strip_relation_qualifiers(str(where_clause), index.table.name)
+    return _strip_relation_qualifiers(str(where_clause), index.table.name, index.table.schema)
 
 
 def _suppress_standard_bm25_ops(upgrade_ops, bm25_names: set[str]) -> None:
@@ -363,15 +437,27 @@ def _compare_bm25_indexes(autogen_context, upgrade_ops, schemas) -> PriorityDisp
     # Emit drop ops for indexes present in DB but absent from MetaData.
     for key in db_bm25:
         if key not in meta_bm25:
-            upgrade_ops.ops.append(DropBM25IndexOp(index_name=key[1], if_exists=True, schema=key[0]))
+            db = db_bm25[key]
+            upgrade_ops.ops.append(
+                DropBM25IndexOp(
+                    index_name=key[1],
+                    if_exists=True,
+                    schema=key[0],
+                    table_name=db["table_name"],
+                    expressions=db["expressions"],
+                    key_field=db["key_field"],
+                    where=db.get("where"),
+                )
+            )
 
     # Emit create ops for indexes present in MetaData but absent from DB.
     # Also re-create indexes whose expression list, key_field, or WHERE clause differs from the DB.
     for key, index in meta_bm25.items():
         with_opts = index.dialect_options["postgresql"].get("with") or {}
         key_field = with_opts.get("key_field", "")
         expressions = [
-            _strip_relation_qualifiers(_render_bm25_expression(expr), index.table.name) for expr in index.expressions
+            _strip_relation_qualifiers(_render_bm25_expression(expr), index.table.name, index.table.schema)
+            for expr in index.expressions
         ]
         meta_where = _render_where_from_index(index)
         create_op = CreateBM25IndexOp(
@@ -393,7 +479,17 @@ def _compare_bm25_indexes(autogen_context, upgrade_ops, schemas) -> PriorityDisp
             key_field_changed = db["key_field"] != key_field
             where_changed = _normalize_where(db.get("where")) != _normalize_where(meta_where)
             if expressions_changed or key_field_changed or where_changed:
-                upgrade_ops.ops.append(DropBM25IndexOp(index_name=key[1], if_exists=True, schema=key[0]))
+                upgrade_ops.ops.append(
+                    DropBM25IndexOp(
+                        index_name=key[1],
+                        if_exists=True,
+                        schema=key[0],
+                        table_name=db["table_name"],
+                        expressions=db["expressions"],
+                        key_field=db["key_field"],
+                        where=db.get("where"),
+                    )
+                )
                 upgrade_ops.ops.append(create_op)
 
     return PriorityDispatchResult.CONTINUE
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "sqlalchemy-paradedb"
-version = "0.2.0"
+version = "0.3.0"
 description = "Typed SQLAlchemy helpers for ParadeDB"
 readme = "README.md"
 requires-python = ">=3.10"
diff --git a/tests/integration/alembic/env.py b/tests/integration/alembic/env.py
@@ -0,0 +1,77 @@
+"""
+Alembic configuration used to test generating and running migrations in test_alembic_commands_integration.py
+"""
+
+from __future__ import annotations
+
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+
+import paradedb.sqlalchemy.alembic  # noqa: F401
+
+config = context.config
+target_metadata = config.attributes["target_metadata"]
+target_table_keys = set(target_metadata.tables)
+target_schemas = {table.schema for table in target_metadata.tables.values() if table.schema is not None}
+version_table_schema = next(iter(target_schemas), None)
+
+
+def include_name(name, type_, parent_names):
+    if type_ == "schema":
+        return name in target_schemas
+    if type_ == "table":
+        schema_name = parent_names.get("schema_name")
+        if schema_name not in target_schemas:
+            return False
+        key = ".".join(part for part in (schema_name, name) if part)
+        return key in target_table_keys
+    return True
+
+
+def include_object(object_, name, type_, reflected, compare_to):
+    if type_ == "table":
+        table = object_ if hasattr(object_, "schema") else compare_to
+        if table is None:
+            return False
+        key = ".".join(part for part in (table.schema, table.name) if part)
+        return key in target_table_keys
+    return True
+
+
+def _configure_context(connection) -> None:
+    context.configure(
+        connection=connection,
+        target_metadata=target_metadata,
+        include_schemas=True,
+        include_name=include_name,
+        include_object=include_object,
+        version_table_schema=version_table_schema,
+    )
+
+
+def run_migrations_offline() -> None:
+    raise RuntimeError("offline migrations are not used in this test")
+
+
+def run_migrations_online() -> None:
+    connectable = config.attributes.get("connection")
+    if connectable is None:
+        connectable = engine_from_config(
+            config.get_section(config.config_ini_section, {}),
+            prefix="sqlalchemy.",
+            poolclass=pool.NullPool,
+        )
+
+    if hasattr(connectable, "connect"):
+        with connectable.connect() as connection:
+            _configure_context(connection)
+            with context.begin_transaction():
+                context.run_migrations()
+        return
+
+    _configure_context(connectable)
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+run_migrations_online()
diff --git a/tests/integration/alembic/script.py.mako b/tests/integration/alembic/script.py.mako
@@ -0,0 +1,25 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from __future__ import annotations
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import os
+import shutil
 from collections.abc import Iterator
+from pathlib import Path
 from typing import Any
 
 import pytest
+from alembic.config import Config
 from sqlalchemy import Boolean, DateTime, Integer, String, Text, create_engine, text
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.dialects.postgresql import JSONB
@@ -40,6 +43,7 @@ class MockItem(Base):
 
 
 PARADEDB_SCAN_PROVIDERS = {"ParadeDB Base Scan", "ParadeDB Aggregate Scan", "ParadeDB Join Scan"}
+ALEMBIC_TEMPLATE_DIR = Path(__file__).with_name("alembic")
 
 
 @pytest.fixture(scope="session")
@@ -165,3 +169,19 @@ def mock_session(engine: Engine, paradedb_ready: None) -> Iterator[Session]:
     """Session fixture for tests using the mock_items table."""
     with Session(engine) as session:
         yield session
+
+
+@pytest.fixture()
+def alembic_config_factory(tmp_path: Path, db_url: str):
+    def factory(metadata) -> Config:
+        script_location = tmp_path / "alembic"
+        shutil.copytree(ALEMBIC_TEMPLATE_DIR, script_location)
+        (script_location / "versions").mkdir()
+
+        config = Config()
+        config.set_main_option("script_location", str(script_location))
+        config.set_main_option("sqlalchemy.url", db_url)
+        config.attributes["target_metadata"] = metadata
+        return config
+
+    return factory
diff --git a/tests/integration/test_alembic_commands_integration.py b/tests/integration/test_alembic_commands_integration.py
diff --git a/tests/unit/test_alembic_unit.py b/tests/unit/test_alembic_unit.py
diff --git a/uv.lock b/uv.lock
