diff --git a/docs/usage/migrations.rst b/docs/usage/migrations.rst index f2efef6c0..803a1c985 100644 --- a/docs/usage/migrations.rst +++ b/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 diff --git a/sqlspec/migrations/runner.py b/sqlspec/migrations/runner.py index b3c9af134..e543ced16 100644 --- a/sqlspec/migrations/runner.py +++ b/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(): diff --git a/tests/integration/migrations/test_upgrade_downgrade_versions.py b/tests/integration/migrations/test_upgrade_downgrade_versions.py index b3a679819..419111838 100644 --- a/tests/integration/migrations/test_upgrade_downgrade_versions.py +++ b/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 diff --git a/tests/unit/migrations/test_migration.py b/tests/unit/migrations/test_migration.py index 3bcb4ac13..58302176a 100644 --- a/tests/unit/migrations/test_migration.py +++ b/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 diff --git a/tests/unit/migrations/test_migration_execution.py b/tests/unit/migrations/test_migration_execution.py index 09de8e765..573f24ef0 100644 --- a/tests/unit/migrations/test_migration_execution.py +++ b/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: diff --git a/tests/unit/migrations/test_migration_runner.py b/tests/unit/migrations/test_migration_runner.py index 377ab9b61..91556c0f4 100644 --- a/tests/unit/migrations/test_migration_runner.py +++ b/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: