fix(migrations): resolve extension SQL query names

docs/usage/migrations.rst

@@ -295,6 +295,19 @@ Both forms record the extension under ``extension_config`` and opt it into
 ``get_migration_commands()`` -- mutating ``extension_config`` directly after the
 configuration is built does not re-run discovery.
 
+SQL files inside a registered extension directory use their filename-local
+version in named-query directives. For example,
+``migrations/0001_create_queue.sql`` declares ``migrate-0001-up`` and
+``migrate-0001-down``. SQLSpec still records that migration as
+``ext_litestar_queues_0001`` when the registered extension name is
+``litestar_queues``.
+
+An extension migration stored in the application's main migration directory
+instead carries the prefix in its filename, such as
+``ext_litestar_queues_0001_create_queue.sql``, and therefore declares
+``migrate-ext_litestar_queues_0001-up`` and
+``migrate-ext_litestar_queues_0001-down``.
+
 .. note::
 
    Extension migrations are versioned under an ``ext_{name}_`` prefix, and that

sqlspec/migrations/runner.py

@@ -87,6 +87,7 @@ def __init__(
         self.extension_migrations = extension_migrations or {}
         self.runtime = runtime
         self.loader = SQLFileLoader(runtime=runtime)
+        self._extension_sql_loaders: dict[str, SQLFileLoader] = {}
         self.project_root: Path | None = None
         self.context = context
         self.extension_configs = extension_configs or {}
@@ -618,6 +619,22 @@ def _finalize_migration_sql(self, sql_statements: Any) -> "list[str] | None":
             return cast("list[str]", sql_statements)
         return None
 
+    def _migration_sql_loader(self, file_path: Path, version: "str | None") -> SQLFileLoader:
+        """Return the isolated core SQL loader for an extension migration."""
+        extension = parse_extension_stem(version) if version else None
+        if file_path.suffix != ".sql" or extension is None:
+            return self.loader
+
+        extension_name, _ = extension
+        if self.extension_migrations.get(extension_name) != file_path.parent:
+            return self.loader
+
+        loader = self._extension_sql_loaders.get(extension_name)
+        if loader is None:
+            loader = SQLFileLoader(runtime=self.runtime)
+            self._extension_sql_loaders[extension_name] = loader
+        return loader
+
 
 class SyncMigrationRunner(BaseMigrationRunner):
     """Synchronous migration runner with pure sync methods."""
@@ -643,15 +660,16 @@ def load_migration(self, file_path: Path, version: "str | None" = None) -> "Load
         metadata = self._load_metadata(file_path, version)
         context_to_use = self._migration_context(file_path)
 
-        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use, self.loader)
+        sql_loader = self._migration_sql_loader(file_path, metadata["version"])
+        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use, sql_loader)
         loader.validate_migration_file(file_path)
 
         has_upgrade, has_downgrade = True, False
 
         if file_path.suffix == ".sql":
-            version = metadata["version"]
-            up_query, down_query = f"migrate-{version}-up", f"migrate-{version}-down"
-            has_upgrade, has_downgrade = self.loader.has_query(up_query), self.loader.has_query(down_query)
+            partial = cast("LoadedMigrationMetadata", {"loader": loader, "file_path": file_path})
+            has_upgrade = bool(self._migration_sql(partial, "up"))
+            has_downgrade = bool(self._migration_sql(partial, "down"))
         else:
             try:
                 partial = cast("LoadedMigrationMetadata", {"loader": loader, "file_path": file_path})
@@ -827,6 +845,15 @@ def load_all_migrations(self) -> "dict[str, SQL]":
 
         for version, file_path in migrations:
             if file_path.suffix == ".sql":
+                if self._migration_sql_loader(file_path, version) is not self.loader:
+                    migration = self.load_migration(file_path, version)
+                    up_sql = self._migration_sql(migration, "up")
+                    down_sql = self._migration_sql(migration, "down")
+                    if up_sql:
+                        all_queries[f"migrate-{version}-up"] = SQL(up_sql[0])
+                    if down_sql:
+                        all_queries[f"migrate-{version}-down"] = SQL(down_sql[0])
+                    continue
                 if not self.loader.has_query(f"migrate-{version}-up"):
                     self.loader.load_sql(file_path)
                 for query_name in self.loader.list_queries():
@@ -877,15 +904,16 @@ async def load_migration(self, file_path: Path, version: "str | None" = None) ->
         metadata = self._load_metadata(file_path, version)
         context_to_use = self._migration_context(file_path)
 
-        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use, self.loader)
+        sql_loader = self._migration_sql_loader(file_path, metadata["version"])
+        loader = get_migration_loader(file_path, self.migrations_path, self.project_root, context_to_use, sql_loader)
         loader.validate_migration_file(file_path)
 
         has_upgrade, has_downgrade = True, False
 
         if file_path.suffix == ".sql":
-            version = metadata["version"]
-            up_query, down_query = f"migrate-{version}-up", f"migrate-{version}-down"
-            has_upgrade, has_downgrade = self.loader.has_query(up_query), self.loader.has_query(down_query)
+            partial = cast("LoadedMigrationMetadata", {"loader": loader, "file_path": file_path})
+            has_upgrade = bool(await self._migration_sql(partial, "up"))
+            has_downgrade = bool(await self._migration_sql(partial, "down"))
         else:
             try:
                 partial = cast("LoadedMigrationMetadata", {"loader": loader, "file_path": file_path})
@@ -1062,6 +1090,15 @@ async def load_all_migrations(self) -> "dict[str, SQL]":
 
         for version, file_path in migrations:
             if file_path.suffix == ".sql":
+                if self._migration_sql_loader(file_path, version) is not self.loader:
+                    migration = await self.load_migration(file_path, version)
+                    up_sql = await self._migration_sql(migration, "up")
+                    down_sql = await self._migration_sql(migration, "down")
+                    if up_sql:
+                        all_queries[f"migrate-{version}-up"] = SQL(up_sql[0])
+                    if down_sql:
+                        all_queries[f"migrate-{version}-down"] = SQL(down_sql[0])
+                    continue
                 if not self.loader.has_query(f"migrate-{version}-up"):
                     await async_(self.loader.load_sql)(file_path)
                 for query_name in self.loader.list_queries():

tests/integration/migrations/test_upgrade_downgrade_versions.py

@@ -414,8 +414,8 @@ def test_downgrade_dry_run_shows_pending_downgrades(sqlite_config: SqliteConfig,
 def test_upgrade_applies_third_party_extension_migrations(tmp_path: Path) -> None:
     """A package outside sqlspec.extensions applies migrations through the public API.
 
-    Covers both extension migration layouts: a file inside the extension's own directory,
-    and an ``ext_``-prefixed file sitting in the main migrations directory.
+    Covers Python and SQL files inside the extension's own directory and an
+    ``ext_``-prefixed SQL file sitting in the main migrations directory.
     """
     migrations_dir = tmp_path / "migrations"
     migrations_dir.mkdir()
@@ -442,11 +442,19 @@ def down(context: "object | None" = None) -> "list[str]":
     """Return the downgrade statements."""
     return ["DROP TABLE queue_tasks"]
 ''')
-    (migrations_dir / "ext_litestar_queues_0002_add_index.sql").write_text(
-        """-- name: migrate-ext_litestar_queues_0002-up
+    (vendor_dir / "0002_create_queue_metadata.sql").write_text(
+        """-- name: migrate-0002-up
+CREATE TABLE queue_metadata (id INTEGER PRIMARY KEY);
+
+-- name: migrate-0002-down
+DROP TABLE queue_metadata;
+"""
+    )
+    (migrations_dir / "ext_litestar_queues_0003_add_index.sql").write_text(
+        """-- name: migrate-ext_litestar_queues_0003-up
 CREATE INDEX queue_tasks_id_idx ON queue_tasks (id);
 
--- name: migrate-ext_litestar_queues_0002-down
+-- name: migrate-ext_litestar_queues_0003-down
 DROP INDEX queue_tasks_id_idx;
 """
     )
@@ -466,8 +474,10 @@ def down(context: "object | None" = None) -> "list[str]":
             tables = session.execute("SELECT name FROM sqlite_master WHERE type = 'table'").get_data()
 
         versions = [row["version_num"] for row in applied]
-        assert versions == ["0001", "ext_litestar_queues_0001", "ext_litestar_queues_0002"]
-        assert "queue_tasks" in {row["name"] for row in tables}
+        assert versions == ["0001", "ext_litestar_queues_0001", "ext_litestar_queues_0002", "ext_litestar_queues_0003"]
+        table_names = {row["name"] for row in tables}
+        assert "queue_tasks" in table_names
+        assert "queue_metadata" in table_names
 
         commands.downgrade(revision="base")
         assert commands.current() is None

tests/unit/migrations/test_migration.py

@@ -556,18 +556,8 @@ def test_load_migration_metadata_no_downgrade(tmp_path: Path) -> None:
 """
     migration_file.write_text(migration_content)
 
-    runner = MockMigrationRunner(migrations_path)
-
-    runner.loader.clear_cache = Mock()
-    runner.loader.load_sql = Mock()
-    runner.loader.has_query = Mock(side_effect=lambda query: query.endswith("-up"))
-
-    with patch("sqlspec.migrations.runner.get_migration_loader") as mock_get_loader:
-        mock_loader = Mock()
-        mock_loader.validate_migration_file = Mock()
-        mock_get_loader.return_value = mock_loader
-
-        metadata = runner.load_migration(migration_file)
+    runner = SyncMigrationRunner(migrations_path)
+    metadata = runner.load_migration(migration_file)
 
     assert metadata["has_upgrade"] is True
     assert metadata["has_downgrade"] is False

tests/unit/migrations/test_migration_execution.py

@@ -18,6 +18,7 @@
 import pytest
 
 from sqlspec.driver import ExecutionResult
+from sqlspec.migrations import MigrationLoadError
 from sqlspec.migrations.base import AppliedMigrationRecord, BaseMigrationTracker, LoadedMigrationMetadata
 from sqlspec.migrations.runner import SyncMigrationRunner
 
@@ -389,8 +390,6 @@ def test_multiple_migrations_execution_order(temp_workspace_with_migrations: Pat
 
 def test_migration_with_no_downgrade(temp_workspace_with_migrations: Path) -> None:
     """Test migration execution when no downgrade is available."""
-    from unittest.mock import AsyncMock
-
     migrations_dir = temp_workspace_with_migrations / "migrations"
 
     migration_file = migrations_dir / "0001_irreversible.sql"
@@ -403,38 +402,16 @@ def test_migration_with_no_downgrade(temp_workspace_with_migrations: Path) -> No
 
     runner = MockMigrationRunner(migrations_dir)
     mock_driver = Mock()
+    migration = runner.load_migration(migration_file)
 
-    with patch("sqlspec.migrations.runner.get_migration_loader") as mock_get_loader:
-        mock_loader = Mock()
-        mock_loader.validate_migration_file = Mock()
-        mock_get_loader.return_value = mock_loader
-
-        with (
-            patch.object(type(runner.loader), "clear_cache"),
-            patch.object(type(runner.loader), "load_sql"),
-            patch.object(type(runner.loader), "has_query", side_effect=lambda q: q.endswith("-up")),
-        ):
-            migration = runner.load_migration(migration_file)
-
-            assert migration["has_upgrade"] is True
-            assert migration["has_downgrade"] is False
-
-    with patch.object(migration["loader"], "get_up_sql", new_callable=AsyncMock) as mock_get_up_sql:
-        mock_get_up_sql.return_value = [
-            "CREATE TABLE irreversible_data AS SELECT DISTINCT column1, column2 FROM legacy_table;"
-        ]
-
-        result = runner.execute_upgrade(mock_driver, migration)
-        assert result is not None
+    assert migration["has_upgrade"] is True
+    assert migration["has_downgrade"] is False
 
-    with (
-        patch.object(migration["loader"], "get_down_sql", new_callable=AsyncMock) as mock_get_down_sql,
-        patch("sqlspec.migrations.runner.logger"),
-    ):
-        mock_get_down_sql.return_value = []
-        result = runner.execute_downgrade(mock_driver, migration)
+    result = runner.execute_upgrade(mock_driver, migration)
+    assert result is not None
 
-        assert result is not None
+    result = runner.execute_downgrade(mock_driver, migration)
+    assert result is not None
 
 
 def test_migration_state_recording() -> None:
@@ -572,27 +549,9 @@ def test_missing_upgrade_migration(temp_workspace_with_migrations: Path) -> None
     migration_file.write_text(migration_content)
 
     runner = MockMigrationRunner(migrations_dir)
-    mock_driver = Mock()
-
-    with patch("sqlspec.migrations.runner.get_migration_loader") as mock_get_loader:
-        mock_loader = Mock()
-        mock_loader.validate_migration_file = Mock()
-        mock_get_loader.return_value = mock_loader
-
-        with (
-            patch.object(type(runner.loader), "clear_cache"),
-            patch.object(type(runner.loader), "load_sql"),
-            patch.object(type(runner.loader), "has_query", side_effect=lambda q: q.endswith("-down")),
-        ):
-            migration = runner.load_migration(migration_file)
-
-            assert migration["has_upgrade"] is False
-            assert migration["has_downgrade"] is True
-
-    with pytest.raises(ValueError) as exc_info:
-        runner.execute_upgrade(mock_driver, migration)
 
-    assert "has no upgrade query" in str(exc_info.value)
+    with pytest.raises(MigrationLoadError, match="missing required 'up' query"):
+        runner.load_migration(migration_file)
 
 
 def test_corrupted_migration_file(temp_workspace_with_migrations: Path) -> None:

tests/unit/migrations/test_migration_runner.py

@@ -535,13 +535,7 @@ def test_load_migration_metadata_prefers_sql_description(tmp_path: Path) -> None
     )
 
     runner = create_migration_runner_with_metadata(tmp_path)
-
-    with (
-        patch.object(type(runner.loader), "clear_cache"),
-        patch.object(type(runner.loader), "load_sql"),
-        patch.object(type(runner.loader), "has_query", return_value=True),
-    ):
-        metadata = runner.load_migration(migration_file)
+    metadata = runner.load_migration(migration_file)
 
     assert metadata["description"] == "Custom summary"
 
@@ -568,6 +562,82 @@ def test_load_migration_metadata_prefers_python_docstring(tmp_path: Path) -> Non
     assert metadata["description"] == "Add feature"
 
 
+def test_sync_loads_sql_from_registered_extension_directory(tmp_path: Path) -> None:
+    migrations_dir = tmp_path / "migrations"
+    migrations_dir.mkdir()
+    extension_dir = tmp_path / "widgets"
+    extension_dir.mkdir()
+    migration_file = extension_dir / "0001_create_widgets.sql"
+    _write_basic_sql(migration_file, "0001", "CREATE TABLE widgets (id INTEGER PRIMARY KEY);")
+    runner = SyncMigrationRunner(migrations_dir, {"widgets": extension_dir})
+
+    version, discovered_file = runner.get_migration_files()[0]
+    migration = runner.load_migration(discovered_file, version)
+
+    assert migration["version"] == "ext_widgets_0001"
+    assert migration["has_upgrade"] is True
+    assert migration["has_downgrade"] is True
+    assert runner._migration_sql(migration, "up") == [  # pyright: ignore[reportPrivateUsage]
+        "CREATE TABLE widgets (id INTEGER PRIMARY KEY);"
+    ]
+
+
+@pytest.mark.anyio
+async def test_async_loads_sql_from_registered_extension_directory(tmp_path: Path) -> None:
+    migrations_dir = tmp_path / "migrations"
+    migrations_dir.mkdir()
+    extension_dir = tmp_path / "widgets"
+    extension_dir.mkdir()
+    migration_file = extension_dir / "0001_create_widgets.sql"
+    _write_basic_sql(migration_file, "0001", "CREATE TABLE widgets (id INTEGER PRIMARY KEY);")
+    runner = AsyncMigrationRunner(migrations_dir, {"widgets": extension_dir})
+
+    version, discovered_file = (await runner.get_migration_files())[0]
+    migration = await runner.load_migration(discovered_file, version)
+
+    assert migration["version"] == "ext_widgets_0001"
+    assert migration["has_upgrade"] is True
+    assert migration["has_downgrade"] is True
+    assert await runner._migration_sql(migration, "up") == [  # pyright: ignore[reportPrivateUsage]
+        "CREATE TABLE widgets (id INTEGER PRIMARY KEY);"
+    ]
+
+
+def test_sync_load_all_migrations_isolates_registered_extension_queries(tmp_path: Path) -> None:
+    migrations_dir = tmp_path / "migrations"
+    migrations_dir.mkdir()
+    alpha_dir = tmp_path / "alpha"
+    alpha_dir.mkdir()
+    beta_dir = tmp_path / "beta"
+    beta_dir.mkdir()
+    _write_basic_sql(alpha_dir / "0001_create_alpha.sql", "0001", "SELECT 'alpha';")
+    _write_basic_sql(beta_dir / "0001_create_beta.sql", "0001", "SELECT 'beta';")
+    runner = SyncMigrationRunner(migrations_dir, {"alpha": alpha_dir, "beta": beta_dir})
+
+    queries = runner.load_all_migrations()
+
+    assert queries["migrate-ext_alpha_0001-up"].raw_sql == "SELECT 'alpha';"
+    assert queries["migrate-ext_beta_0001-up"].raw_sql == "SELECT 'beta';"
+
+
+@pytest.mark.anyio
+async def test_async_load_all_migrations_isolates_registered_extension_queries(tmp_path: Path) -> None:
+    migrations_dir = tmp_path / "migrations"
+    migrations_dir.mkdir()
+    alpha_dir = tmp_path / "alpha"
+    alpha_dir.mkdir()
+    beta_dir = tmp_path / "beta"
+    beta_dir.mkdir()
+    _write_basic_sql(alpha_dir / "0001_create_alpha.sql", "0001", "SELECT 'alpha';")
+    _write_basic_sql(beta_dir / "0001_create_beta.sql", "0001", "SELECT 'beta';")
+    runner = AsyncMigrationRunner(migrations_dir, {"alpha": alpha_dir, "beta": beta_dir})
+
+    queries = await runner.load_all_migrations()
+
+    assert queries["migrate-ext_alpha_0001-up"].raw_sql == "SELECT 'alpha';"
+    assert queries["migrate-ext_beta_0001-up"].raw_sql == "SELECT 'beta';"
+
+
 def test_sync_load_all_migrations_skips_load_sql_when_query_is_loaded(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None: