mongrator.runner¶

class SyncRunner:
    """Synchronous migration runner backed by pymongo."""

    def __init__(self, client: "MongoClient", config: MigratorConfig) -> None:  # type: ignore[type-arg]
        self._client = client
        self._db = client[config.database]
        self._store = SyncStateStore(self._db[config.collection])
        self._lock = SyncMigrationLock(self._db[config.collection])
        self._config = config

    def plan_up(self, target: MigrationId | None = None) -> MigrationPlan:
        """Return the plan for applying pending migrations without executing."""
        files = loader.load(self._config)
        applied = self._store.get_applied()
        return planner.plan_up(files, applied, target)

    def plan_down(self, steps: int = 1) -> MigrationPlan:
        """Return the plan for rolling back migrations without executing."""
        files = loader.load(self._config)
        applied = self._store.get_applied()
        return planner.plan_down(files, applied, steps)

    def up(self, target: MigrationId | None = None, *, transactional: bool = False) -> list[MigrationId]:
        """Apply pending migrations, optionally up to `target`.

        Args:
            target: Stop after applying this migration ID.
            transactional: When True, wrap each migration in a MongoDB
                transaction. Requires a replica set; raises
                ``TransactionNotSupportedError`` otherwise.
        """
        if transactional:
            _check_transaction_support(self._client)
        with self._lock:
            files = loader.load(self._config)
            applied = self._store.get_applied()
            plan = planner.plan_up(files, applied, target)
            applied_ids: list[MigrationId] = []
            for migration in plan.to_apply:
                start = time.monotonic()
                if transactional:
                    with self._client.start_session() as session:
                        with session.start_transaction():
                            _run_up_migration(migration, _SessionBoundDatabase(self._db, session))
                else:
                    _run_up_migration(migration, self._db)
                duration_ms = int((time.monotonic() - start) * 1000)
                self._store.record_applied(make_record(migration.id, migration.checksum, "up", duration_ms))
                applied_ids.append(migration.id)
            return applied_ids

    def down(self, steps: int = 1, *, transactional: bool = False) -> list[MigrationId]:
        """Roll back the most recently applied migrations.

        Args:
            steps: Number of migrations to roll back.
            transactional: When True, wrap each migration in a MongoDB
                transaction. Requires a replica set; raises
                ``TransactionNotSupportedError`` otherwise.
        """
        if transactional:
            _check_transaction_support(self._client)
        with self._lock:
            files = loader.load(self._config)
            applied = self._store.get_applied()
            plan = planner.plan_down(files, applied, steps)
            rolled_back: list[MigrationId] = []
            for migration in plan.to_apply:
                start = time.monotonic()
                if transactional:
                    with self._client.start_session() as session:
                        with session.start_transaction():
                            _run_down_migration(migration, _SessionBoundDatabase(self._db, session))
                else:
                    _run_down_migration(migration, self._db)
                duration_ms = int((time.monotonic() - start) * 1000)
                self._store.record_applied(make_record(migration.id, migration.checksum, "down", duration_ms))
                rolled_back.append(migration.id)
            return rolled_back

    def status(self) -> list[MigrationStatus]:
        """Return the status of every known migration, including orphans.

        This is a read-only operation and intentionally does not acquire the
        advisory lock so it can run safely while migrations are in progress.
        """
        files = loader.load(self._config)
        applied = self._store.get_applied()
        file_ids = {f.id for f in files}
        statuses: list[MigrationStatus] = []
        for f in files:
            if f.id in applied:
                record = self._store.get_record(f.id)
                checksum_ok = record is None or record["checksum"] == f.checksum
                applied_at = record["applied_at"] if record else None
            else:
                checksum_ok = True
                applied_at = None
            statuses.append(
                MigrationStatus(
                    id=f.id,
                    applied=f.id in applied,
                    applied_at=applied_at,
                    checksum_ok=checksum_ok,
                )
            )
        for orphan_id in sorted(applied - file_ids):
            record = self._store.get_record(orphan_id)
            statuses.append(
                MigrationStatus(
                    id=orphan_id,
                    applied=True,
                    applied_at=record["applied_at"] if record else None,
                    orphaned=True,
                )
            )
        return statuses

    def validate(self) -> list[ChecksumMismatchError]:
        """Check that applied migration files match their recorded checksums.

        This is a read-only operation and intentionally does not acquire the
        advisory lock so it can run safely while migrations are in progress.
        """
        files = loader.load(self._config)
        applied = self._store.get_applied()
        errors: list[ChecksumMismatchError] = []
        for f in files:
            if f.id not in applied:
                continue
            record = self._store.get_record(f.id)
            if record and record["checksum"] != f.checksum:
                errors.append(ChecksumMismatchError(f.id, record["checksum"], f.checksum))
        return errors

class AsyncRunner:
    """Asynchronous migration runner backed by pymongo AsyncMongoClient.

    Migration functions are dispatched based on their type:

    - **Coroutine functions** (``async def up(db)``) receive the async database
      and are awaited. This allows migrations to use ``await`` for non-blocking
      I/O when running with ``--async``.
    - **Regular functions** that return a ``list[Operation]`` (ops-based migrations)
      continue to receive the sync database, since ops helpers are synchronous.
    - **Regular functions** that perform raw pymongo calls also receive the sync
      database, preserving backwards compatibility.

    State tracking always uses the async client for non-blocking operation.
    """

    def __init__(  # type: ignore[type-arg]
        self,
        client: "AsyncMongoClient",
        config: MigratorConfig,
        *,
        sync_client: "MongoClient | None" = None,
    ) -> None:
        # Sync DB passed to sync migration functions and ops helpers.
        if sync_client is None:
            sync_client = MongoClient(config.uri)
        self._sync_client = sync_client
        self._db = sync_client[config.database]
        # Async DB passed to coroutine migration functions.
        self._async_db = client[config.database]
        # Async store for non-blocking state tracking.
        self._store = AsyncMongoStateStore(self._async_db[config.collection])
        self._lock = AsyncMigrationLock(self._async_db[config.collection])
        self._config = config

    async def plan_up(self, target: MigrationId | None = None) -> MigrationPlan:
        """Return the plan for applying pending migrations without executing."""
        files = loader.load(self._config)
        applied = await self._store.get_applied()
        return planner.plan_up(files, applied, target)

    async def plan_down(self, steps: int = 1) -> MigrationPlan:
        """Return the plan for rolling back migrations without executing."""
        files = loader.load(self._config)
        applied = await self._store.get_applied()
        return planner.plan_down(files, applied, steps)

    async def up(self, target: MigrationId | None = None, *, transactional: bool = False) -> list[MigrationId]:
        """Apply pending migrations, optionally up to `target`.

        Coroutine migration functions receive the async database and are awaited;
        sync functions and ops-based migrations use the sync database.

        Args:
            target: Stop after applying this migration ID.
            transactional: When True, wrap each migration in a MongoDB
                transaction. Requires a replica set; raises
                ``TransactionNotSupportedError`` otherwise.
        """
        if transactional:
            _check_transaction_support(self._sync_client)
        async with self._lock:
            files = loader.load(self._config)
            applied = await self._store.get_applied()
            plan = planner.plan_up(files, applied, target)
            applied_ids: list[MigrationId] = []
            for migration in plan.to_apply:
                start = time.monotonic()
                if transactional:
                    with self._sync_client.start_session() as session:
                        with session.start_transaction():
                            await _async_run_up_migration(
                                migration, self._async_db, _SessionBoundDatabase(self._db, session)
                            )
                else:
                    await _async_run_up_migration(migration, self._async_db, self._db)
                duration_ms = int((time.monotonic() - start) * 1000)
                await self._store.record_applied(make_record(migration.id, migration.checksum, "up", duration_ms))
                applied_ids.append(migration.id)
            return applied_ids

    async def down(self, steps: int = 1, *, transactional: bool = False) -> list[MigrationId]:
        """Roll back the most recently applied migrations.

        Coroutine migration functions receive the async database and are awaited;
        sync functions and ops-based migrations use the sync database.

        Args:
            steps: Number of migrations to roll back.
            transactional: When True, wrap each migration in a MongoDB
                transaction. Requires a replica set; raises
                ``TransactionNotSupportedError`` otherwise.
        """
        if transactional:
            _check_transaction_support(self._sync_client)
        async with self._lock:
            files = loader.load(self._config)
            applied = await self._store.get_applied()
            plan = planner.plan_down(files, applied, steps)
            rolled_back: list[MigrationId] = []
            for migration in plan.to_apply:
                start = time.monotonic()
                if transactional:
                    with self._sync_client.start_session() as session:
                        with session.start_transaction():
                            await _async_run_down_migration(
                                migration, self._async_db, _SessionBoundDatabase(self._db, session)
                            )
                else:
                    await _async_run_down_migration(migration, self._async_db, self._db)
                duration_ms = int((time.monotonic() - start) * 1000)
                await self._store.record_applied(make_record(migration.id, migration.checksum, "down", duration_ms))
                rolled_back.append(migration.id)
            return rolled_back

    async def status(self) -> list[MigrationStatus]:
        """Return the status of every known migration, including orphans.

        This is a read-only operation and intentionally does not acquire the
        advisory lock so it can run safely while migrations are in progress.
        """
        files = loader.load(self._config)
        applied = await self._store.get_applied()
        file_ids = {f.id for f in files}
        statuses: list[MigrationStatus] = []
        for f in files:
            if f.id in applied:
                record = await self._store.get_record(f.id)
                checksum_ok = record is None or record["checksum"] == f.checksum
                applied_at = record["applied_at"] if record else None
            else:
                checksum_ok = True
                applied_at = None
            statuses.append(
                MigrationStatus(
                    id=f.id,
                    applied=f.id in applied,
                    applied_at=applied_at,
                    checksum_ok=checksum_ok,
                )
            )
        for orphan_id in sorted(applied - file_ids):
            record = await self._store.get_record(orphan_id)
            statuses.append(
                MigrationStatus(
                    id=orphan_id,
                    applied=True,
                    applied_at=record["applied_at"] if record else None,
                    orphaned=True,
                )
            )
        return statuses

    async def validate(self) -> list[ChecksumMismatchError]:
        """Check that applied migration files match their recorded checksums.

        This is a read-only operation and intentionally does not acquire the
        advisory lock so it can run safely while migrations are in progress.
        """
        files = loader.load(self._config)
        applied = await self._store.get_applied()
        errors: list[ChecksumMismatchError] = []
        for f in files:
            if f.id not in applied:
                continue
            record = await self._store.get_record(f.id)
            if record and record["checksum"] != f.checksum:
                errors.append(ChecksumMismatchError(f.id, record["checksum"], f.checksum))
        return errors