snake-orm 0.1.0b1__py3-none-any.whl

snakeorm/migration/asyncrunner.py

@@ -0,0 +1,168 @@
+"""Asynchronous migration runtime: same logic, the other execution seam.
+
+SQL generation has no "colour" (`up_sql()` does not execute), so diff, DDL, `realize()` and the
+operations are reused as they are; only who talks to the DB changes. The real difference: DATA
+migrations do NOT run here — `RunPython` receives a SYNCHRONOUS `SnakeSession`, which would block the
+event loop — and it stops dead rather than marking them applied with the data never executed.
+"""
+
+from __future__ import annotations
+
+from snakeorm.core.placement import DEFAULT_SCHEMA
+from snakeorm.dialects import SnakeDialect
+from snakeorm.drivers.asyncbase import AsyncDriver
+from snakeorm.core.exceptions import SnakeMigrationError
+from snakeorm.migration.operations import (
+    SnakeDataOperation,
+    SnakeMigrationOperation,
+    SnakeOperation,
+)
+from snakeorm.migration.realize import realize
+from snakeorm.migration.runner import (
+    Migration,
+    explain_rebuild_failure,
+    squash_already_done,
+    tracking_table_ddl,
+)
+from snakeorm.sql.refs import qualified
+
+_TRACKING_TABLE = "snake_migrations"
+_SCHEMA = DEFAULT_SCHEMA
+
+
+class AsyncMigrationRunner:
+    """Applies/reverts migrations over an `AsyncDriver`, with the same tracking as the sync one."""
+
+    def __init__(self, driver: AsyncDriver, dialect: SnakeDialect) -> None:
+        self._driver = driver
+        self._dialect = dialect
+
+    async def ensure_tracking_table(self) -> None:
+        """Creates the tracking table if it does not exist (idempotent)."""
+        await self._driver.execute(tracking_table_ddl(self._dialect), ())
+        await self._driver.commit()
+
+    async def applied_versions(self) -> set[str]:
+        """Returns the set of versions already applied."""
+        version = self._dialect.quote_ident("version")
+        rows = await self._driver.fetch_all(
+            f"SELECT {version} FROM {self._table_ref()}", ()
+        )
+        return {str(row[0]) for row in rows}
+
+    async def apply(self, migrations: list[Migration]) -> list[str]:
+        """Applies the pending migrations in order. Returns the versions applied just now.
+
+        Same semantics as the synchronous one: idempotent, atomic per migration with transactional
+        DDL, and a hard stop if a squash replaces a history that was applied HALFWAY.
+        """
+        await self.ensure_tracking_table()
+        already = await self.applied_versions()
+        replaced = {
+            version for migration in migrations for version in migration.replaces
+        }
+        newly: list[str] = []
+        for migration in migrations:
+            if migration.version in already or migration.version in replaced:
+                continue
+            if migration.replaces and squash_already_done(migration, already):
+                await self._record(migration.version)
+                await self._driver.commit()
+                newly.append(migration.version)
+                continue
+            if self._dialect.supports_transactional_ddl:
+                await self._apply_atomic(migration)
+            else:
+                await self._apply_stepwise(migration)
+            newly.append(migration.version)
+        return newly
+
+    async def rollback(self, migration: Migration) -> None:
+        """Reverts a migration (operations in reverse order) and deletes its record."""
+        for operation in reversed(realize(migration.operations, self._dialect)):
+            await self._backward(operation)
+        await self._unrecord(migration.version)
+        await self._driver.commit()
+
+    async def _apply_atomic(self, migration: Migration) -> None:
+        """Applies a migration wrapped in the engine's transaction: all or nothing."""
+        plan = realize(migration.operations, self._dialect)
+        try:
+            for operation in plan:
+                await self._forward(operation)
+            await self._record(migration.version)
+            await self._driver.commit()
+        except Exception as error:
+            await self._driver.rollback()
+            # The SAME sentence the synchronous runner gives, out of the same function: a rebuild
+            # that the deferred keys could not cover fails identically here, and this repository has
+            # already paid for one complaint told two ways.
+            explained = explain_rebuild_failure(plan, self._dialect, error)
+            if explained is not None:
+                raise explained from error
+            raise
+
+    async def _apply_stepwise(self, migration: Migration) -> None:
+        """Applies operation by operation (engine without transactional DDL); on error, reports how many."""
+        applied = 0
+        plan = realize(migration.operations, self._dialect)
+        try:
+            for operation in plan:
+                await self._forward(operation)
+                applied += 1
+            await self._record(migration.version)
+            await self._driver.commit()
+        except SnakeMigrationError:
+            raise
+        except Exception as error:
+            raise SnakeMigrationError(
+                f"Migration '{migration.version}' failed after applying {applied} of "
+                f"{len(plan)} operation(s): {error}. The engine does not support "
+                f"transactional DDL, so the database was left half-way: review it and "
+                f"fix it by hand."
+            ) from error
+
+    async def _forward(self, operation: SnakeMigrationOperation) -> None:
+        """Emits a schema operation's DDL; a data one cannot be executed here."""
+        _reject_data_operation(operation)
+        if isinstance(operation, SnakeOperation):
+            for sql in operation.up_sql(self._dialect):
+                await self._driver.execute(sql, ())
+
+    async def _backward(self, operation: SnakeMigrationOperation) -> None:
+        """Emits a schema operation's reverse DDL."""
+        _reject_data_operation(operation)
+        if isinstance(operation, SnakeOperation):
+            for sql in operation.down_sql(self._dialect):
+                await self._driver.execute(sql, ())
+
+    def _table_ref(self) -> str:
+        """Reference to the tracking table, qualified only if the engine has schemas."""
+        return qualified(_SCHEMA, _TRACKING_TABLE, self._dialect)
+
+    async def _record(self, version: str) -> None:
+        column = self._dialect.quote_ident("version")
+        await self._driver.execute(
+            f"INSERT INTO {self._table_ref()} ({column}) "
+            f"VALUES ({self._dialect.placeholder(1)})",
+            (version,),
+        )
+
+    async def _unrecord(self, version: str) -> None:
+        column = self._dialect.quote_ident("version")
+        await self._driver.execute(
+            f"DELETE FROM {self._table_ref()} WHERE {column} = {self._dialect.placeholder(1)}",
+            (version,),
+        )
+
+
+def _reject_data_operation(operation: SnakeMigrationOperation) -> None:
+    """Stops dead at a DATA migration: skipping it would mark it applied with the data never run,
+    and that never raises an error, only wrong rows months later."""
+    if isinstance(operation, SnakeDataOperation):
+        raise SnakeMigrationError(
+            "The async runner cannot execute a data operation (`RunPython`): its "
+            "`forward`/`backward` take a SYNCHRONOUS `SnakeSession`, and running it here would "
+            "block the event loop for the whole migration. Apply that migration with "
+            "`MigrationRunner` (synchronous) over a synchronous driver; the schema ones do go here."
+        )

snakeorm/migration/autodetect.py

@@ -0,0 +1,96 @@
+"""Django-style autogen: rebuilds the state by replaying migrations and diffs it.
+
+`replay` applies the operations of every migration onto an empty SchemaState -> the state the
+schema SHOULD have according to the history. `autodetect` diffs that against the current
+metadata (the code-first source of truth) and returns the operations for the new migration.
+Neither snapshots nor reflection: the migration history IS the record of the state.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from snakeorm.metadata import SnakeRoutineInfo, SnakeTableInfo, SnakeTriggerInfo
+from snakeorm.migration.diff import diff_routines, diff_schema, diff_triggers
+from snakeorm.migration.operations import SnakeOperation
+from snakeorm.migration.runner import Migration
+from snakeorm.migration.state import SchemaState
+from snakeorm.registry import SnakeRegistry, registry
+
+
+def current_schema(
+    reg: SnakeRegistry = registry,
+    *,
+    database: str | None = None,
+    include_unmanaged: bool = False,
+) -> list[SnakeTableInfo]:
+    """The current metadata (code-first source of truth): the tables of the registered models.
+
+    Excludes MIRROR models (`@snake_db_first`): they are not a source of truth, and autogen must not
+    touch them. With `database`, only those of THAT connection (which avoids creating every table in
+    every DB). `include_unmanaged=True` returns them too: drift detection uses it (comparing the
+    mirror against the DB), but MIGRATIONS never touch them (the default behaviour).
+    """
+    return [
+        table
+        for model in reg.models()
+        if (table := reg.table_of(model)) is not None
+        and (include_unmanaged or table.is_managed)
+        and (database is None or table.database == database)
+        # A polymorphic CHILD contributes no table: it shares its base's (which already carries
+        # the whole hierarchy after `snake_link()`). Including it would duplicate the `CREATE TABLE`
+        # and produce phantom drift on every `makemigrations`.
+        and not table.is_polymorphic_child
+    ]
+
+
+def current_routines(reg: SnakeRegistry = registry) -> list[SnakeRoutineInfo]:
+    """The desired routines (code-first source of truth): those declared with `snake_function`."""
+    return list(reg.routines())
+
+
+def current_triggers(reg: SnakeRegistry = registry) -> list[SnakeTriggerInfo]:
+    """The DESIRED triggers: those declared with `snake_trigger(...)`."""
+    return list(reg.triggers())
+
+
+def replay(migrations: Iterable[Migration]) -> SchemaState:
+    """Rebuilds the schema state by applying the migrations' operations in order."""
+    state = SchemaState()
+    for migration in migrations:
+        for operation in migration.operations:
+            operation.apply_to_state(state)
+    return state
+
+
+def autodetect(
+    migrations: Iterable[Migration],
+    current: Iterable[SnakeTableInfo],
+    routines: Iterable[SnakeRoutineInfo] | None = None,
+    triggers: Iterable[SnakeTriggerInfo] | None = None,
+) -> list[SnakeOperation]:
+    """Diffs the state replayed from the history against the current metadata -> new operations.
+
+    Resolves FK targets by model name through the global registry. Global order: tables -> FKs ->
+    views (topological) -> functions -> triggers (each one depends on the ones before it).
+    """
+    state = replay(migrations)
+    desired_routines = current_routines() if routines is None else list(routines)
+    # The state's triggers go in with the tables, and NOT because they are diffed here — that is
+    # `diff_triggers`, further down. A `RebuildTable` on an engine without `ALTER TABLE ADD
+    # CONSTRAINT` drops the table it remakes, which takes its triggers with it, and the operation
+    # has to carry them because `SnakeTableInfo` has no field for one. This is the only caller that
+    # has them, so this is where they are handed over.
+    schema_ops = diff_schema(
+        state.tables(),
+        current,
+        registry.table_by_name,
+        registry.table_by_qualified,
+        triggers=state.triggers(),
+    )
+    routine_ops = diff_routines(state.routines(), desired_routines)
+    # TRIGGERS go last: a Postgres one executes a FUNCTION, which has to exist beforehand.
+    trigger_ops = diff_triggers(
+        state.triggers(), current_triggers() if triggers is None else list(triggers)
+    )
+    return [*schema_ops, *routine_ops, *trigger_ops]