mongrator 0.1.0__py3-none-any.whl

mongrator/__init__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+__all__ = ["main"]

mongrator/_templates/migration.py.tmpl

@@ -0,0 +1,40 @@
+"""Migration: {slug}
+
+Generated by mongrator on {timestamp}.
+
+Use the ops helpers for common reversible operations, or write plain PyMongo
+calls directly against the `db` argument for anything more complex.
+
+Example using ops helpers (auto-rollback supported):
+
+    from mongrator import ops
+
+    def up(db):
+        return [
+            ops.create_index("my_collection", {{"field": 1}}, unique=True),
+        ]
+
+Example using plain PyMongo (define down() for rollback):
+
+    def up(db):
+        db["my_collection"].update_many({{}}, {{"$set": {{"new_field": None}}}})
+
+    def down(db):
+        db["my_collection"].update_many({{}}, {{"$unset": {{"new_field": ""}}}})
+"""
+
+
+# Uncomment to use the ops helpers:
+# from mongrator import ops
+
+
+def up(db):
+    pass
+    # return [
+    #     ops.create_index("my_collection", {{"field": 1}}),
+    # ]
+
+
+# def down(db):
+#     """Optional: define rollback logic here if not using ops helpers."""
+#     pass

mongrator/cli.py

@@ -0,0 +1,256 @@
+"""Command-line interface for mongrator.
+
+Subcommands:
+    init     — create the migrations directory and a mongrator.toml stub
+    create   — generate a new timestamped migration file
+    status   — show applied/pending migration table
+    up       — apply pending migrations
+    down     — roll back applied migrations
+    validate — verify checksums of applied migrations
+"""
+
+import argparse
+import asyncio
+import sys
+from datetime import UTC, datetime
+from importlib.resources import files
+from pathlib import Path
+
+from .config import MigratorConfig
+from .exceptions import MigratorError
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="mongrator",
+        description="Lightweight MongoDB schema migration tool",
+    )
+    parser.add_argument(
+        "--config",
+        metavar="PATH",
+        default="mongrator.toml",
+        help="path to config file (default: mongrator.toml)",
+    )
+
+    sub = parser.add_subparsers(dest="command", metavar="COMMAND")
+    sub.required = True
+
+    # init
+    sub.add_parser("init", help="create migrations directory and config stub")
+
+    # create
+    p_create = sub.add_parser("create", help="generate a new migration file")
+    p_create.add_argument("name", help="short description, e.g. add_users_email_index")
+
+    # status
+    sub.add_parser("status", help="show applied/pending migration table")
+
+    # up
+    p_up = sub.add_parser("up", help="apply pending migrations")
+    p_up.add_argument("--target", metavar="ID", help="apply only up to this migration ID")
+    p_up.add_argument("--async", dest="use_async", action="store_true", help="use async runner")
+
+    # down
+    p_down = sub.add_parser("down", help="roll back applied migrations")
+    p_down.add_argument(
+        "--steps", type=int, default=1, metavar="N", help="number of migrations to roll back (default: 1)"
+    )
+    p_down.add_argument("--async", dest="use_async", action="store_true", help="use async runner")
+
+    # validate
+    sub.add_parser("validate", help="verify checksums of applied migration files")
+
+    return parser
+
+
+def _load_config(args: argparse.Namespace) -> MigratorConfig:
+    config_path = Path(args.config)
+    if config_path.exists():
+        return MigratorConfig.from_toml(config_path)
+    return MigratorConfig.from_env()
+
+
+# ---------------------------------------------------------------------------
+# Subcommand handlers
+# ---------------------------------------------------------------------------
+
+
+def _cmd_init(args: argparse.Namespace) -> int:
+    config_path = Path(args.config)
+    if not config_path.exists():
+        config_path.write_text(
+            '[mongrator]\nuri = "mongodb://localhost:27017"\ndatabase = "mydb"\n'
+            'migrations_dir = "migrations"\ncollection = "mongrator_migrations"\n'
+        )
+        print(f"Created {config_path}")
+
+    migrations_dir = Path("migrations")
+    migrations_dir.mkdir(exist_ok=True)
+    print(f"Created {migrations_dir}/")
+    return 0
+
+
+def _cmd_create(args: argparse.Namespace) -> int:
+    config = _load_config(args)
+    config.migrations_dir.mkdir(parents=True, exist_ok=True)
+
+    slug = args.name.strip().replace(" ", "_").lower()
+    timestamp = datetime.now(tz=UTC).strftime("%Y%m%d_%H%M%S")
+    filename = f"{timestamp}_{slug}.py"
+    dest = config.migrations_dir / filename
+
+    template_text = files("mongrator._templates").joinpath("migration.py.tmpl").read_text(encoding="utf-8")
+    content = template_text.format(slug=slug, timestamp=timestamp)
+    dest.write_text(content, encoding="utf-8")
+    print(f"Created {dest}")
+    return 0
+
+
+def _cmd_status(args: argparse.Namespace) -> int:
+    from .runner import SyncRunner
+
+    config = _load_config(args)
+    try:
+        import pymongo
+    except ImportError:
+        print("error: pymongo is required. Install with: pip install pymongo", file=sys.stderr)
+        return 1
+
+    client = pymongo.MongoClient(config.uri)
+    runner = SyncRunner(client, config)
+    statuses = runner.status()
+
+    if not statuses:
+        print("No migrations found.")
+        return 0
+
+    col_width = max(len(s.id) for s in statuses) + 2
+    print(f"{'Migration':<{col_width}} {'Status':<10} {'Applied At'}")
+    print("-" * (col_width + 30))
+    for s in statuses:
+        state = "applied" if s.applied else "pending"
+        if s.applied and not s.checksum_ok:
+            state = "MODIFIED"
+        applied_at = s.applied_at.isoformat() if s.applied_at else "-"
+        print(f"{s.id:<{col_width}} {state:<10} {applied_at}")
+    return 0
+
+
+def _cmd_up(args: argparse.Namespace) -> int:
+    config = _load_config(args)
+    if args.use_async:
+        return asyncio.run(_async_up(config, args.target))
+    import pymongo
+
+    from .runner import SyncRunner
+
+    client = pymongo.MongoClient(config.uri)
+    runner = SyncRunner(client, config)
+    applied = runner.up(target=args.target)
+    if applied:
+        for mid in applied:
+            print(f"  applied  {mid}")
+    else:
+        print("Nothing to apply.")
+    return 0
+
+
+async def _async_up(config: MigratorConfig, target: str | None) -> int:
+    from pymongo import AsyncMongoClient
+
+    from .runner import AsyncRunner
+
+    client = AsyncMongoClient(config.uri)
+    runner = AsyncRunner(client, config)
+    applied = await runner.up(target=target)
+    if applied:
+        for mid in applied:
+            print(f"  applied  {mid}")
+    else:
+        print("Nothing to apply.")
+    return 0
+
+
+def _cmd_down(args: argparse.Namespace) -> int:
+    config = _load_config(args)
+    if args.use_async:
+        return asyncio.run(_async_down(config, args.steps))
+    import pymongo
+
+    from .runner import SyncRunner
+
+    client = pymongo.MongoClient(config.uri)
+    runner = SyncRunner(client, config)
+    rolled_back = runner.down(steps=args.steps)
+    if rolled_back:
+        for mid in rolled_back:
+            print(f"  rolled back  {mid}")
+    else:
+        print("Nothing to roll back.")
+    return 0
+
+
+async def _async_down(config: MigratorConfig, steps: int) -> int:
+    from pymongo import AsyncMongoClient
+
+    from .runner import AsyncRunner
+
+    client = AsyncMongoClient(config.uri)
+    runner = AsyncRunner(client, config)
+    rolled_back = await runner.down(steps=steps)
+    if rolled_back:
+        for mid in rolled_back:
+            print(f"  rolled back  {mid}")
+    else:
+        print("Nothing to roll back.")
+    return 0
+
+
+def _cmd_validate(args: argparse.Namespace) -> int:
+    import pymongo
+
+    from .runner import SyncRunner
+
+    config = _load_config(args)
+    client = pymongo.MongoClient(config.uri)
+    runner = SyncRunner(client, config)
+    errors = runner.validate()
+
+    if not errors:
+        print("All applied migrations have valid checksums.")
+        return 0
+
+    eg = ExceptionGroup("Checksum mismatches detected", errors)
+    print(f"error: {eg}", file=sys.stderr)
+    for e in errors:
+        print(f"  {e}", file=sys.stderr)
+    return 1
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = _build_parser()
+    args = parser.parse_args()
+
+    dispatch = {
+        "init": _cmd_init,
+        "create": _cmd_create,
+        "status": _cmd_status,
+        "up": _cmd_up,
+        "down": _cmd_down,
+        "validate": _cmd_validate,
+    }
+
+    try:
+        rc = dispatch[args.command](args)
+    except MigratorError as e:
+        print(f"error: {e}", file=sys.stderr)
+        sys.exit(1)
+    except KeyboardInterrupt:
+        sys.exit(130)
+
+    sys.exit(rc)

mongrator/config.py

@@ -0,0 +1,61 @@
+import os
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Self
+
+from .exceptions import ConfigurationError
+
+_DEFAULT_COLLECTION = "mongrator_migrations"
+_DEFAULT_MIGRATIONS_DIR = Path("migrations")
+
+
+@dataclass(frozen=True)
+class MigratorConfig:
+    """Immutable configuration for a migrator instance."""
+
+    uri: str
+    database: str
+    migrations_dir: Path
+    collection: str = _DEFAULT_COLLECTION
+
+    @classmethod
+    def from_toml(cls, path: Path) -> Self:
+        """Load configuration from a TOML file (e.g. mongrator.toml)."""
+        try:
+            with open(path, "rb") as f:
+                data = tomllib.load(f)
+        except FileNotFoundError:
+            raise ConfigurationError(f"Config file not found: {path}")
+        except tomllib.TOMLDecodeError as e:
+            raise ConfigurationError(f"Invalid TOML in {path}: {e}")
+
+        try:
+            uri: str = data["uri"]
+            database: str = data["database"]
+        except KeyError as e:
+            raise ConfigurationError(f"Missing required config key: {e}")
+
+        migrations_dir = Path(data.get("migrations_dir", str(_DEFAULT_MIGRATIONS_DIR)))
+        collection: str = data.get("collection", _DEFAULT_COLLECTION)
+        return cls(uri=uri, database=database, migrations_dir=migrations_dir, collection=collection)
+
+    @classmethod
+    def from_env(cls) -> Self:
+        """Load configuration from environment variables.
+
+        Variables:
+            MONGRATOR_URI         — MongoDB connection URI (required)
+            MONGRATOR_DB          — database name (required)
+            MONGRATOR_MIGRATIONS_DIR — path to migrations directory (default: migrations)
+            MONGRATOR_COLLECTION  — tracking collection name (default: mongrator_migrations)
+        """
+        uri = os.environ.get("MONGRATOR_URI")
+        database = os.environ.get("MONGRATOR_DB")
+        if not uri:
+            raise ConfigurationError("MONGRATOR_URI environment variable is not set")
+        if not database:
+            raise ConfigurationError("MONGRATOR_DB environment variable is not set")
+        migrations_dir = Path(os.environ.get("MONGRATOR_MIGRATIONS_DIR", str(_DEFAULT_MIGRATIONS_DIR)))
+        collection = os.environ.get("MONGRATOR_COLLECTION", _DEFAULT_COLLECTION)
+        return cls(uri=uri, database=database, migrations_dir=migrations_dir, collection=collection)

mongrator/exceptions.py

@@ -0,0 +1,65 @@
+class MigratorError(Exception):
+    """Base class for all mongrator errors."""
+
+
+class ConfigurationError(MigratorError):
+    """Invalid or missing configuration."""
+
+
+class ChecksumMismatchError(MigratorError):
+    """Applied migration file has been modified since it was run."""
+
+    def __init__(self, migration_id: str, expected: str, actual: str) -> None:
+        self.migration_id = migration_id
+        self.expected = expected
+        self.actual = actual
+        super().__init__(
+            f"Checksum mismatch for '{migration_id}': "
+            f"expected {expected!r}, got {actual!r}. "
+            "The migration file has been modified after being applied."
+        )
+
+
+class DuplicateMigrationIdError(MigratorError):
+    """Two migration files share the same ID."""
+
+    def __init__(self, migration_id: str) -> None:
+        self.migration_id = migration_id
+        super().__init__(f"Duplicate migration ID: '{migration_id}'")
+
+
+class MigrationImportError(MigratorError):
+    """A migration file could not be imported."""
+
+    def __init__(self, path: str, cause: Exception) -> None:
+        self.path = path
+        self.cause = cause
+        super().__init__(f"Failed to import migration '{path}': {cause}")
+
+
+class InvalidMigrationFileError(MigratorError):
+    """A migration file is missing required callables or has an invalid structure."""
+
+    def __init__(self, path: str, reason: str) -> None:
+        self.path = path
+        self.reason = reason
+        super().__init__(f"Invalid migration file '{path}': {reason}")
+
+
+class NoDownMethodError(MigratorError):
+    """A migration has no rollback path."""
+
+    def __init__(self, migration_id: str) -> None:
+        self.migration_id = migration_id
+        super().__init__(
+            f"Migration '{migration_id}' has no rollback path. "
+            "Define a down() function or use ops.* helpers that support auto-rollback."
+        )
+
+
+class MigrationNotFoundError(MigratorError):
+    """A referenced migration ID does not exist."""
+
+    def __init__(self, migration_id: str) -> None:
+        self.migration_id = migration_id
+        super().__init__(f"Migration not found: '{migration_id}'")

mongrator/loader.py

@@ -0,0 +1,74 @@
+import hashlib
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+from .config import MigratorConfig
+from .exceptions import (
+    DuplicateMigrationIdError,
+    InvalidMigrationFileError,
+    MigrationImportError,
+)
+from .migration import Checksum, MigrationFile, MigrationId
+
+
+def _checksum(path: Path) -> Checksum:
+    return hashlib.sha256(path.read_bytes()).hexdigest()
+
+
+def _migration_id(path: Path) -> MigrationId:
+    return path.stem
+
+
+def load(config: MigratorConfig) -> list[MigrationFile]:
+    """Scan migrations_dir, import each .py file, and return an ordered list.
+
+    Files are sorted lexicographically by filename, which is chronological
+    when the recommended {timestamp}_{slug}.py naming convention is used.
+
+    Raises:
+        DuplicateMigrationIdError: if two files share the same stem.
+        MigrationImportError: if a file cannot be imported.
+        InvalidMigrationFileError: if a file does not define an up() callable.
+    """
+    migrations_dir = config.migrations_dir
+    if not migrations_dir.exists():
+        return []
+
+    paths = sorted(migrations_dir.glob("*.py"))
+    seen: dict[MigrationId, Path] = {}
+    results: list[MigrationFile] = []
+
+    for path in paths:
+        migration_id = _migration_id(path)
+
+        if migration_id in seen:
+            raise DuplicateMigrationIdError(migration_id)
+        seen[migration_id] = path
+
+        checksum = _checksum(path)
+        module = _import_file(path, migration_id)
+
+        if not callable(getattr(module, "up", None)):
+            raise InvalidMigrationFileError(str(path), "missing a callable up() function")
+
+        results.append(MigrationFile(id=migration_id, path=path, checksum=checksum, module=module))
+
+    return results
+
+
+def _import_file(path: Path, module_name: str) -> types.ModuleType:
+    # Use a namespaced module name to avoid collisions with installed packages.
+    qualified_name = f"mongrator._migrations.{module_name}"
+    spec = importlib.util.spec_from_file_location(qualified_name, path)
+    if spec is None or spec.loader is None:
+        raise MigrationImportError(str(path), RuntimeError("could not create module spec"))
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[qualified_name] = module
+    try:
+        spec.loader.exec_module(module)  # type: ignore[union-attr]
+    except Exception as e:
+        del sys.modules[qualified_name]
+        raise MigrationImportError(str(path), e) from e
+    return module

mongrator/migration.py

@@ -0,0 +1,57 @@
+import types
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Literal, TypedDict
+
+type MigrationId = str
+type Checksum = str
+
+
+class MigrationRecord(TypedDict):
+    """Document stored in the tracking collection for each applied migration."""
+
+    _id: MigrationId
+    applied_at: datetime
+    checksum: Checksum
+    direction: Literal["up", "down"]
+    duration_ms: int
+
+
+@dataclass
+class MigrationFile:
+    """Represents a discovered migration file on disk."""
+
+    id: MigrationId
+    path: Path
+    checksum: Checksum
+    module: types.ModuleType | None = field(default=None, repr=False)
+
+    @property
+    def up(self) -> Callable[..., Any] | None:
+        if self.module is None:
+            return None
+        return getattr(self.module, "up", None)
+
+    @property
+    def down(self) -> Callable[..., Any] | None:
+        if self.module is None:
+            return None
+        return getattr(self.module, "down", None)
+
+    def has_up(self) -> bool:
+        return self.up is not None
+
+    def has_down(self) -> bool:
+        return self.down is not None
+
+
+@dataclass
+class MigrationStatus:
+    """Status of a single migration as reported by runner.status()."""
+
+    id: MigrationId
+    applied: bool
+    applied_at: datetime | None = None
+    checksum_ok: bool = True

mongrator/ops.py

@@ -0,0 +1,128 @@
+"""Declarative operation helpers for MongoDB migrations.
+
+Each helper returns an Operation whose apply() and revert() methods perform the
+forward and reverse changes respectively. When a migration's up() function returns
+a list[Operation] and no down() is defined, the runner auto-generates rollback by
+calling revert() on each operation in reverse order.
+
+Usage in a migration file::
+
+    from mongrator import ops
+
+    def up(db):
+        return [
+            ops.create_index("users", {"email": 1}, unique=True),
+            ops.rename_field("users", "name", "full_name"),
+        ]
+"""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from pymongo.database import Database
+
+
+@dataclass
+class Operation:
+    """An atomic, reversible database operation."""
+
+    description: str
+    _apply: Any = field(repr=False)
+    _revert: Any = field(repr=False)
+
+    def apply(self, db: Database) -> None:  # type: ignore[type-arg]
+        self._apply(db)
+
+    def revert(self, db: Database) -> None:  # type: ignore[type-arg]
+        self._revert(db)
+
+
+def create_index(
+    collection: str,
+    keys: dict[str, int],
+    **kwargs: Any,
+) -> Operation:
+    """Create an index. Reverts by dropping the index."""
+    index_name: str | None = kwargs.get("name")
+
+    def apply(db: Database) -> None:  # type: ignore[type-arg]
+        db[collection].create_index(list(keys.items()), **kwargs)
+
+    def revert(db: Database) -> None:  # type: ignore[type-arg]
+        name = index_name or "_".join(f"{k}_{v}" for k, v in keys.items())
+        db[collection].drop_index(name)
+
+    key_repr = ", ".join(f"{k}: {v}" for k, v in keys.items())
+    return Operation(
+        description=f"create_index({collection!r}, {{{key_repr}}})",
+        _apply=apply,
+        _revert=revert,
+    )
+
+
+def drop_index(collection: str, index_name: str) -> Operation:
+    """Drop an index by name. Not auto-reversible (index spec is unknown)."""
+
+    def apply(db: Database) -> None:  # type: ignore[type-arg]
+        db[collection].drop_index(index_name)
+
+    def revert(db: Database) -> None:  # type: ignore[type-arg]
+        raise NotImplementedError(
+            f"drop_index({collection!r}, {index_name!r}) cannot be auto-reverted. "
+            "Define a down() function to recreate the index."
+        )
+
+    return Operation(
+        description=f"drop_index({collection!r}, {index_name!r})",
+        _apply=apply,
+        _revert=revert,
+    )
+
+
+def rename_field(
+    collection: str,
+    old_name: str,
+    new_name: str,
+    filter: dict[str, Any] | None = None,
+) -> Operation:
+    """Rename a field across all (or filtered) documents. Reverts by renaming back."""
+    query = filter or {}
+
+    def apply(db: Database) -> None:  # type: ignore[type-arg]
+        db[collection].update_many(query, {"$rename": {old_name: new_name}})
+
+    def revert(db: Database) -> None:  # type: ignore[type-arg]
+        db[collection].update_many(query, {"$rename": {new_name: old_name}})
+
+    return Operation(
+        description=f"rename_field({collection!r}, {old_name!r} → {new_name!r})",
+        _apply=apply,
+        _revert=revert,
+    )
+
+
+def add_field(
+    collection: str,
+    field_name: str,
+    default_value: Any,
+    filter: dict[str, Any] | None = None,
+) -> Operation:
+    """Add a field with a default value to all (or filtered) documents.
+    Reverts by unsetting the field.
+    """
+    query = filter or {}
+
+    def apply(db: Database) -> None:  # type: ignore[type-arg]
+        db[collection].update_many(
+            {**query, field_name: {"$exists": False}},
+            {"$set": {field_name: default_value}},
+        )
+
+    def revert(db: Database) -> None:  # type: ignore[type-arg]
+        db[collection].update_many(query, {"$unset": {field_name: ""}})
+
+    return Operation(
+        description=f"add_field({collection!r}, {field_name!r}={default_value!r})",
+        _apply=apply,
+        _revert=revert,
+    )

mongrator/planner.py

@@ -0,0 +1,76 @@
+from dataclasses import dataclass
+
+from .exceptions import MigrationNotFoundError
+from .migration import MigrationFile, MigrationId
+
+
+@dataclass
+class MigrationPlan:
+    to_apply: list[MigrationFile]
+    to_skip: list[MigrationFile]
+
+
+def plan_up(
+    files: list[MigrationFile],
+    applied: set[MigrationId],
+    target: MigrationId | None = None,
+) -> MigrationPlan:
+    """Compute which migrations need to be applied in the forward direction.
+
+    Args:
+        files:   All known migration files, in chronological order.
+        applied: Set of migration IDs already recorded in the tracking collection.
+        target:  If given, apply only up to and including this migration ID.
+                 Raises MigrationNotFoundError if target is not in files.
+
+    Returns:
+        A MigrationPlan with to_apply (pending) and to_skip (already applied).
+    """
+    if target is not None:
+        ids = {f.id for f in files}
+        if target not in ids:
+            raise MigrationNotFoundError(target)
+
+    to_apply: list[MigrationFile] = []
+    to_skip: list[MigrationFile] = []
+
+    for f in files:
+        if f.id in applied:
+            to_skip.append(f)
+        else:
+            to_apply.append(f)
+        if target is not None and f.id == target:
+            break
+
+    return MigrationPlan(to_apply=to_apply, to_skip=to_skip)
+
+
+def plan_down(
+    files: list[MigrationFile],
+    applied: set[MigrationId],
+    steps: int = 1,
+) -> MigrationPlan:
+    """Compute which migrations to roll back.
+
+    Rolls back the most recently applied migrations, up to `steps` of them.
+    Migrations are identified by their position in the applied set intersected
+    with the ordered file list (file order is the canonical order).
+
+    Args:
+        files:   All known migration files, in chronological order.
+        applied: Set of migration IDs already recorded in the tracking collection.
+        steps:   Number of most-recent applied migrations to roll back.
+
+    Returns:
+        A MigrationPlan where to_apply contains the migrations to roll back
+        (in reverse order) and to_skip contains those left untouched.
+    """
+    if steps < 1:
+        raise ValueError(f"steps must be >= 1, got {steps}")
+
+    applied_in_order = [f for f in files if f.id in applied]
+    to_rollback = list(reversed(applied_in_order[-steps:]))
+    rollback_ids = {f.id for f in to_rollback}
+    to_skip = [f for f in files if f.id in applied and f.id not in rollback_ids]
+
+    return MigrationPlan(to_apply=to_rollback, to_skip=to_skip)

mongrator/runner.py

@@ -0,0 +1,218 @@
+"""Migration runners: Protocol definitions and sync/async implementations.
+
+SyncRunner wraps a pymongo MongoClient.
+AsyncRunner wraps a pymongo AsyncMongoClient.
+
+Both share the same non-IO logic via loader and planner.
+"""
+
+import time
+from typing import Any, Protocol, cast, runtime_checkable
+
+from pymongo import AsyncMongoClient, MongoClient
+
+from . import loader, planner
+from .config import MigratorConfig
+from .exceptions import ChecksumMismatchError, NoDownMethodError
+from .migration import MigrationFile, MigrationId, MigrationStatus
+from .ops import Operation
+from .state import AsyncMongoStateStore, SyncStateStore, make_record
+
+
+@runtime_checkable
+class MigrationRunner(Protocol):
+    def up(self, target: MigrationId | None = None) -> list[MigrationId]: ...
+    def down(self, steps: int = 1) -> list[MigrationId]: ...
+    def status(self) -> list[MigrationStatus]: ...
+    def validate(self) -> list[ChecksumMismatchError]: ...
+
+
+@runtime_checkable
+class AsyncMigrationRunner(Protocol):
+    async def up(self, target: MigrationId | None = None) -> list[MigrationId]: ...
+    async def down(self, steps: int = 1) -> list[MigrationId]: ...
+    async def status(self) -> list[MigrationStatus]: ...
+    async def validate(self) -> list[ChecksumMismatchError]: ...
+
+
+def _resolve_rollback(migration: MigrationFile) -> None:
+    """Raise NoDownMethodError if the migration cannot be rolled back."""
+    if migration.has_down():
+        return
+    # Check if up() returns a list of Operations with revert support
+    # (We cannot call up() here; this is checked at rollback time instead.)
+    raise NoDownMethodError(migration.id)
+
+
+def _run_up_migration(migration: MigrationFile, db: Any) -> None:
+    """Execute the up() callable, calling Operation.apply() for each op if needed."""
+    up_fn = migration.up
+    if up_fn is None:
+        return
+    result = up_fn(db)
+    if isinstance(result, list) and all(isinstance(op, Operation) for op in result):
+        # ops-based migration: up() returns ops, runner applies them.
+        # Raw pymongo migrations return None.
+        for op in cast(list[Operation], result):
+            op.apply(db)  # type: ignore[arg-type]
+
+
+def _run_down_migration(migration: MigrationFile, db: Any) -> None:
+    """Execute the down() callable, or auto-revert ops returned by up()."""
+    down_fn = migration.down
+    if down_fn is not None:
+        down_fn(db)
+        return
+    # Try auto-rollback via ops returned from up()
+    up_fn = migration.up
+    if up_fn is None:
+        raise NoDownMethodError(migration.id)
+    result = up_fn(db)
+    if isinstance(result, list) and all(isinstance(op, Operation) for op in result):
+        for op in reversed(cast(list[Operation], result)):
+            op.revert(db)  # type: ignore[arg-type]
+        return
+    raise NoDownMethodError(migration.id)
+
+
+class SyncRunner:
+    """Synchronous migration runner backed by pymongo."""
+
+    def __init__(self, client: "MongoClient", config: MigratorConfig) -> None:  # type: ignore[type-arg]
+        self._db = client[config.database]
+        self._store = SyncStateStore(self._db[config.collection])
+        self._config = config
+
+    def up(self, target: MigrationId | None = None) -> list[MigrationId]:
+        """Apply pending migrations, optionally up to `target`."""
+        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()
+            _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) -> list[MigrationId]:
+        """Roll back the most recently applied migrations."""
+        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()
+            _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."""
+        files = loader.load(self._config)
+        applied = self._store.get_applied()
+        statuses: list[MigrationStatus] = []
+        for f in files:
+            record = self._store.get_record(f.id)
+            checksum_ok = record is None or record["checksum"] == f.checksum
+            statuses.append(
+                MigrationStatus(
+                    id=f.id,
+                    applied=f.id in applied,
+                    applied_at=record["applied_at"] if record else None,
+                    checksum_ok=checksum_ok,
+                )
+            )
+        return statuses
+
+    def validate(self) -> list[ChecksumMismatchError]:
+        """Check that applied migration files match their recorded checksums."""
+        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 up()/down() functions always receive a synchronous pymongo Database
+    because ops helpers are synchronous. Only state tracking uses the async client.
+    """
+
+    def __init__(self, client: "AsyncMongoClient", config: MigratorConfig) -> None:  # type: ignore[type-arg]
+        # Sync DB passed to migration functions — ops helpers are synchronous pymongo.
+        self._db = MongoClient(config.uri)[config.database]
+        # Async store for non-blocking state tracking.
+        async_db = client[config.database]
+        self._store = AsyncMongoStateStore(async_db[config.collection])
+        self._config = config
+
+    async def up(self, target: MigrationId | None = None) -> list[MigrationId]:
+        """Apply pending migrations, optionally up to `target`."""
+        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()
+            _run_up_migration(migration, 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) -> list[MigrationId]:
+        """Roll back the most recently applied migrations."""
+        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()
+            _run_down_migration(migration, 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."""
+        files = loader.load(self._config)
+        applied = await self._store.get_applied()
+        statuses: list[MigrationStatus] = []
+        for f in files:
+            record = await self._store.get_record(f.id)
+            checksum_ok = record is None or record["checksum"] == f.checksum
+            statuses.append(
+                MigrationStatus(
+                    id=f.id,
+                    applied=f.id in applied,
+                    applied_at=record["applied_at"] if record else None,
+                    checksum_ok=checksum_ok,
+                )
+            )
+        return statuses
+
+    async def validate(self) -> list[ChecksumMismatchError]:
+        """Check that applied migration files match their recorded checksums."""
+        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

mongrator/state.py

@@ -0,0 +1,101 @@
+"""Migration state storage: Protocol definitions and sync/async implementations.
+
+The tracking collection stores one document per applied migration::
+
+    {
+        "_id": "20260408_143022_add_users_email_index",
+        "applied_at": ISODate("2026-04-08T14:30:22Z"),
+        "checksum": "e3b0c44298fc1c149afb...",
+        "direction": "up",
+        "duration_ms": 42
+    }
+"""
+
+from datetime import UTC, datetime
+from typing import Literal, Protocol, runtime_checkable
+
+from pymongo.asynchronous.collection import AsyncCollection
+from pymongo.collection import Collection
+
+from .migration import MigrationId, MigrationRecord
+
+
+@runtime_checkable
+class StateStore(Protocol):
+    """Synchronous migration state store."""
+
+    def get_applied(self) -> set[MigrationId]: ...
+
+    def record_applied(self, record: MigrationRecord) -> None: ...
+
+    def remove_record(self, migration_id: MigrationId) -> None: ...
+
+    def get_record(self, migration_id: MigrationId) -> MigrationRecord | None: ...
+
+
+@runtime_checkable
+class AsyncStateStore(Protocol):
+    """Asynchronous migration state store."""
+
+    async def get_applied(self) -> set[MigrationId]: ...
+
+    async def record_applied(self, record: MigrationRecord) -> None: ...
+
+    async def remove_record(self, migration_id: MigrationId) -> None: ...
+
+    async def get_record(self, migration_id: MigrationId) -> MigrationRecord | None: ...
+
+
+class SyncStateStore:
+    """StateStore backed by a synchronous pymongo collection."""
+
+    def __init__(self, collection: "Collection") -> None:  # type: ignore[type-arg]
+        self._col = collection
+
+    def get_applied(self) -> set[MigrationId]:
+        return {doc["_id"] for doc in self._col.find({"direction": "up"}, {"_id": 1})}
+
+    def record_applied(self, record: MigrationRecord) -> None:
+        self._col.replace_one({"_id": record["_id"]}, record, upsert=True)
+
+    def remove_record(self, migration_id: MigrationId) -> None:
+        self._col.delete_one({"_id": migration_id})
+
+    def get_record(self, migration_id: MigrationId) -> MigrationRecord | None:
+        return self._col.find_one({"_id": migration_id})  # type: ignore[return-value]
+
+
+class AsyncMongoStateStore:
+    """AsyncStateStore backed by a pymongo AsyncCollection."""
+
+    def __init__(self, collection: "AsyncCollection") -> None:  # type: ignore[type-arg]
+        self._col = collection
+
+    async def get_applied(self) -> set[MigrationId]:
+        cursor = self._col.find({"direction": "up"}, {"_id": 1})
+        return {doc["_id"] async for doc in cursor}
+
+    async def record_applied(self, record: MigrationRecord) -> None:
+        await self._col.replace_one({"_id": record["_id"]}, record, upsert=True)
+
+    async def remove_record(self, migration_id: MigrationId) -> None:
+        await self._col.delete_one({"_id": migration_id})
+
+    async def get_record(self, migration_id: MigrationId) -> MigrationRecord | None:
+        return await self._col.find_one({"_id": migration_id})  # type: ignore[return-value]
+
+
+def make_record(
+    migration_id: MigrationId,
+    checksum: str,
+    direction: Literal["up", "down"],
+    duration_ms: int,
+) -> MigrationRecord:
+    """Construct a MigrationRecord with the current UTC timestamp."""
+    return MigrationRecord(
+        _id=migration_id,
+        applied_at=datetime.now(tz=UTC),
+        checksum=checksum,
+        direction=direction,
+        duration_ms=duration_ms,
+    )

mongrator-0.1.0.dist-info/METADATA

@@ -0,0 +1,158 @@
+Metadata-Version: 2.3
+Name: mongrator
+Version: 0.1.0
+Summary: Lightweight MongoDB schema migration tool
+Author: Sasha Gerrand
+Author-email: Sasha Gerrand <mongrator@sgerrand.dev>
+Requires-Dist: pymongo>=4.10
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+
+# mongrator
+
+[![CI](https://github.com/sgerrand/pymongrator/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/sgerrand/pymongrator/actions/workflows/ci.yml)
+
+Lightweight MongoDB schema migration tool with synchronous and asynchronous PyMongo support.
+
+## Installation
+
+```sh
+pip install mongrator
+```
+
+## Quick start
+
+```sh
+# Create config and migrations directory
+mongrator init
+
+# Generate a new migration file
+mongrator create add_users_email_index
+
+# Check migration status
+mongrator status
+
+# Apply pending migrations
+mongrator up
+
+# Roll back the last migration
+mongrator down
+```
+
+## Configuration
+
+`mongrator init` creates a `mongrator.toml` stub:
+
+```toml
+uri = "mongodb://localhost:27017"
+database = "mydb"
+migrations_dir = "migrations"
+collection = "mongrator_migrations"  # optional
+```
+
+Alternatively, configure via environment variables:
+
+| Variable | Description | Required |
+|----------|-------------|----------|
+| `MONGRATOR_URI` | MongoDB connection URI | yes |
+| `MONGRATOR_DB` | Database name | yes |
+| `MONGRATOR_MIGRATIONS_DIR` | Path to migrations directory | no (default: `migrations`) |
+| `MONGRATOR_COLLECTION` | Tracking collection name | no (default: `mongrator_migrations`) |
+
+## Writing migrations
+
+Migration files are plain Python named `{timestamp}_{slug}.py` (e.g. `20260408_143022_add_users_email_index.py`). Each file must define an `up(db)` function. A `down(db)` function is optional but enables rollback.
+
+### Using the ops helpers (recommended)
+
+The `ops` helpers record their own inverses, so `down()` is generated automatically:
+
+```python
+from mongrator import ops
+
+def up(db):
+    return [
+        ops.create_index("users", {"email": 1}, unique=True),
+        ops.rename_field("users", "username", "handle"),
+        ops.add_field("users", "verified", default_value=False),
+    ]
+```
+
+### Using plain PyMongo
+
+For complex logic, write directly against the `db` argument and define `down()` manually:
+
+```python
+def up(db):
+    db["orders"].update_many(
+        {"status": {"$exists": False}},
+        {"$set": {"status": "pending"}},
+    )
+
+def down(db):
+    db["orders"].update_many({}, {"$unset": {"status": ""}})
+```
+
+### Available ops helpers
+
+| Helper | Reversible | Description |
+|--------|-----------|-------------|
+| `ops.create_index(collection, keys, **kwargs)` | yes | Create an index |
+| `ops.drop_index(collection, index_name)` | no | Drop an index by name |
+| `ops.rename_field(collection, old, new, filter=None)` | yes | Rename a field across documents |
+| `ops.add_field(collection, field, default_value, filter=None)` | yes | Add a field with a default value |
+
+## CLI reference
+
+```
+mongrator init                        create migrations dir and mongrator.toml
+mongrator create <name>               generate a new migration file
+mongrator status                      show applied/pending migrations
+mongrator up [--target ID]            apply pending migrations
+mongrator up --async [--target ID]    apply using async runner
+mongrator down [--steps N]            roll back N migrations (default: 1)
+mongrator down --async [--steps N]    roll back using async runner
+mongrator validate                    verify checksums of applied migrations
+mongrator --config PATH <command>     use an alternate config file
+```
+
+## Async usage
+
+Pass `--async` to `up` or `down` to use the async runner (backed by `pymongo.AsyncMongoClient`):
+
+```sh
+mongrator up --async
+```
+
+To use the runners programmatically:
+
+```python
+# Synchronous
+from pathlib import Path
+import pymongo
+from mongrator.config import MigratorConfig
+from mongrator.runner import SyncRunner
+
+config = MigratorConfig(uri="mongodb://localhost:27017", database="mydb", migrations_dir=Path("migrations"))
+runner = SyncRunner(pymongo.MongoClient(config.uri), config)
+runner.up()
+
+# Asynchronous
+from pymongo import AsyncMongoClient
+from mongrator.runner import AsyncRunner
+
+runner = AsyncRunner(AsyncMongoClient(config.uri), config)
+await runner.up()
+```
+
+## Migration tracking
+
+Applied migrations are recorded in the `mongrator_migrations` collection (configurable) within the target database. Each document stores:
+
+- `_id` — migration file stem
+- `applied_at` — UTC timestamp
+- `checksum` — SHA-256 of the migration file at time of application
+- `direction` — `"up"` or `"down"`
+- `duration_ms` — execution time in milliseconds
+
+Running `mongrator validate` compares current file checksums against recorded values and reports any modifications.

mongrator-0.1.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+mongrator/__init__.py,sha256=oKsCpoiQSSXx57DGs2VfYzeO1ukJQxdGnV-2hC3BI0M,42
+mongrator/_templates/migration.py.tmpl,sha256=GMVZV8hOPEMkQiATScQgDcvswRpWJ3JJCQgSFB6PUiM,950
+mongrator/cli.py,sha256=LhQNfR5hFPe1sJDIZz_6VxgugpEUXmZLlxaITJGw1UI,7637
+mongrator/config.py,sha256=sna2uNl_5GjGqbGfybfL8UpPbVDugT0LYfGf89YqBfg,2412
+mongrator/exceptions.py,sha256=5SzfkbThOIDBN3FaUg5mTzB7thaDmiK1SVkkGxDkDMI,2173
+mongrator/loader.py,sha256=zKBDBiMNfwffgJMs2lSPijsUiT3nQJZlxbIv6858iZY,2478
+mongrator/migration.py,sha256=W1axR4h8-RTwMfiwHejYdc8i-DJQqAE4ICizeXwBlX0,1385
+mongrator/ops.py,sha256=a3PsKoRlayBUwgESxLF5aYHsL_ea5k14qyCH1rSTXj8,3922
+mongrator/planner.py,sha256=t6xmIEWgGybpVqoeQ70S-Domsr0_ntyPruzc2YBI4pA,2521
+mongrator/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mongrator/runner.py,sha256=_XCkks-KCxWIgxGVbtRVyLWIswjITSA3aqNGN9KPm00,9271
+mongrator/state.py,sha256=qw7VBvNe8dhePm3k20g6jCaaZzy39iKht5S9m_T_ick,3408
+mongrator-0.1.0.dist-info/WHEEL,sha256=lh7MMMfiuFQLQaR9J7pNBODdWf-aa5UOeuuDAol3xps,79
+mongrator-0.1.0.dist-info/entry_points.txt,sha256=zu1IQipcRKmYmNXcxqY6_xSfWYBpjFlsXeBHa2yFsQo,46
+mongrator-0.1.0.dist-info/METADATA,sha256=LX_NsgZ91syMTeN6e5vISEDJ-2fQtYRGCdYNMgMuwFk,4762
+mongrator-0.1.0.dist-info/RECORD,,

mongrator-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.8.20
+Root-Is-Purelib: true
+Tag: py3-none-any

mongrator-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+mongrator = mongrator:main
+