mongomig 0.1.0.dev0__py3-none-any.whl

mongomig/migrations/revision.py

@@ -0,0 +1,107 @@
+"""Revision ids, file names and rendering of new revision files."""
+
+from __future__ import annotations
+
+import json
+import re
+import secrets
+import unicodedata
+from datetime import UTC, datetime
+from importlib.resources import files
+from pathlib import Path
+from string import Template
+
+from mongomig.errors import ScriptError
+
+REVISION_ID_RE = re.compile(r"^[A-Za-z0-9_]{1,64}$")
+MAX_SLUG_LENGTH = 40
+
+
+def new_revision_id() -> str:
+    """12 random hex chars: collisions between developers are practically impossible."""
+    return secrets.token_hex(6)
+
+
+def validate_revision_id(rev_id: str) -> str:
+    if not REVISION_ID_RE.match(rev_id):
+        raise ScriptError(
+            f"Invalid revision id {rev_id!r}.",
+            suggestion="Use 1-64 letters, digits or underscores.",
+        )
+    return rev_id
+
+
+def slugify(message: str) -> str:
+    ascii_text = (
+        unicodedata.normalize("NFKD", message).encode("ascii", "ignore").decode("ascii").lower()
+    )
+    slug = re.sub(r"[^a-z0-9]+", "_", ascii_text).strip("_")
+    slug = slug[:MAX_SLUG_LENGTH].rstrip("_")
+    return slug or "revision"
+
+
+def revision_filename(rev_id: str, message: str, created: datetime) -> str:
+    """``YYYYMMDD_HHMM_<id>_<slug>.py`` — sorts chronologically in a file browser."""
+    return f"{created:%Y%m%d_%H%M}_{rev_id}_{slugify(message)}.py"
+
+
+def render_revision(
+    *,
+    rev_id: str,
+    message: str,
+    down_revision: str | tuple[str, ...] | None,
+    snapshot_hash: str | None,
+    created: datetime,
+) -> str:
+    template = Template(
+        files("mongomig").joinpath("templates/revision.py.tmpl").read_text(encoding="utf-8")
+    )
+    if isinstance(down_revision, tuple):
+        down_text = ", ".join(down_revision)
+    else:
+        down_text = down_revision or "<base>"
+    return template.substitute(
+        # The docstring must not be terminated early by the message itself.
+        message=message.replace("\\", "\\\\").replace('"""', "'''").strip() or "empty message",
+        revision=rev_id,
+        down_revision_text=down_text,
+        created=f"{created:%Y-%m-%d %H:%M:%S}",
+        revision_repr=_py_literal(rev_id),
+        down_revision_repr=_py_literal(down_revision),
+        snapshot_hash_repr=_py_literal(snapshot_hash),
+    )
+
+
+def _py_literal(value: str | tuple[str, ...] | None) -> str:
+    """Python literal using double quotes (black/ruff style) for generated files."""
+    if value is None:
+        return "None"
+    if isinstance(value, tuple):
+        return "(" + ", ".join(json.dumps(v) for v in value) + ")"
+    return json.dumps(value)
+
+
+def write_revision(
+    versions_dir: Path,
+    *,
+    message: str,
+    down_revision: str | tuple[str, ...] | None,
+    snapshot_hash: str | None,
+    rev_id: str | None = None,
+    now: datetime | None = None,
+) -> tuple[str, Path]:
+    rev_id = validate_revision_id(rev_id) if rev_id else new_revision_id()
+    created = now or datetime.now(UTC)
+    path = versions_dir / revision_filename(rev_id, message, created)
+    if path.exists():
+        raise ScriptError(f"Refusing to overwrite existing file {path}")
+    content = render_revision(
+        rev_id=rev_id,
+        message=message,
+        down_revision=down_revision,
+        snapshot_hash=snapshot_hash,
+        created=created,
+    )
+    versions_dir.mkdir(parents=True, exist_ok=True)
+    path.write_text(content, encoding="utf-8")
+    return rev_id, path

mongomig/migrations/script.py

@@ -0,0 +1,198 @@
+"""Read revision files.
+
+Metadata is read by parsing the file's AST instead of importing it, so offline commands
+(``history``, ``heads``, ``revision``) are fast and never execute migration code or require
+the application's dependencies. The module is only imported when a migration actually runs.
+"""
+
+from __future__ import annotations
+
+import ast
+import hashlib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from mongomig.errors import RevisionConflictError, ScriptError
+from mongomig.migrations.revision import REVISION_ID_RE
+
+SUPPORTED_FORMAT = 1
+_METADATA_NAMES = frozenset(
+    {
+        "revision",
+        "down_revision",
+        "branch_labels",
+        "depends_on",
+        "reversible",
+        "snapshot_hash",
+        "mongomig_format",
+    }
+)
+
+
+@dataclass(frozen=True)
+class Script:
+    revision: str
+    down_revisions: tuple[str, ...]
+    path: Path
+    message: str
+    branch_labels: tuple[str, ...] = ()
+    depends_on: tuple[str, ...] = ()
+    reversible: bool = True
+    snapshot_hash: str | None = None
+    format_version: int = SUPPORTED_FORMAT
+    checksum: str = ""
+    has_downgrade: bool = True
+
+    @property
+    def is_base(self) -> bool:
+        return not self.down_revisions
+
+    @property
+    def is_merge(self) -> bool:
+        return len(self.down_revisions) > 1
+
+
+def file_checksum(source: bytes) -> str:
+    """sha256 over the file with normalised newlines (CRLF checkouts must not look modified)."""
+    normalised = source.replace(b"\r\n", b"\n")
+    return "sha256:" + hashlib.sha256(normalised).hexdigest()
+
+
+def load_script(path: Path) -> Script:
+    try:
+        source = path.read_bytes()
+    except OSError as exc:
+        raise ScriptError(f"Cannot read {path}: {exc.strerror}") from None
+    try:
+        tree = ast.parse(source, filename=str(path))
+    except SyntaxError as exc:
+        raise ScriptError(
+            f"Syntax error in {path.name} line {exc.lineno}: {exc.msg}",
+            details={"path": str(path)},
+        ) from None
+
+    values, functions = _top_level(tree, path)
+
+    def fail(message: str, suggestion: str | None = None) -> ScriptError:
+        return ScriptError(
+            f"{path.name}: {message}", suggestion=suggestion, details={"path": str(path)}
+        )
+
+    revision = values.get("revision")
+    if not isinstance(revision, str) or not REVISION_ID_RE.match(revision):
+        raise fail("missing or invalid `revision` (expected a string id).")
+
+    fmt = values.get("mongomig_format", SUPPORTED_FORMAT)
+    if not isinstance(fmt, int) or fmt < 1:
+        raise fail("`mongomig_format` must be a positive integer.")
+    if fmt > SUPPORTED_FORMAT:
+        raise fail(
+            f"written for migration format {fmt}; this mongomig supports up to {SUPPORTED_FORMAT}.",
+            "Upgrade mongomig: pip install -U mongomig",
+        )
+
+    reversible = values.get("reversible", True)
+    if not isinstance(reversible, bool):
+        raise fail("`reversible` must be True or False.")
+
+    if "upgrade" not in functions:
+        raise fail("missing `def upgrade(ctx)`.")
+    has_downgrade = "downgrade" in functions
+    if not has_downgrade and reversible:
+        raise fail(
+            "missing `def downgrade(ctx)`.",
+            "Add a downgrade function, or set `reversible = False` if it cannot be undone.",
+        )
+
+    snapshot = values.get("snapshot_hash")
+    if snapshot is not None and not isinstance(snapshot, str):
+        raise fail("`snapshot_hash` must be a string or None.")
+
+    return Script(
+        revision=revision,
+        down_revisions=_id_tuple(values.get("down_revision"), "down_revision", fail),
+        path=path,
+        message=_message(tree),
+        branch_labels=_id_tuple(values.get("branch_labels"), "branch_labels", fail),
+        depends_on=_id_tuple(values.get("depends_on"), "depends_on", fail),
+        reversible=reversible,
+        snapshot_hash=snapshot,
+        format_version=fmt,
+        checksum=file_checksum(source),
+        has_downgrade=has_downgrade,
+    )
+
+
+def load_scripts(versions_dir: Path) -> list[Script]:
+    """Load every revision file in ``versions_dir`` (non-recursive; ``_*`` and ``.*`` skipped)."""
+    if not versions_dir.is_dir():
+        return []
+    scripts: list[Script] = []
+    seen: dict[str, Path] = {}
+    for path in sorted(versions_dir.glob("*.py")):
+        if path.name.startswith(("_", ".")):
+            continue
+        script = load_script(path)
+        if script.revision in seen:
+            raise RevisionConflictError(
+                f"Duplicate revision id {script.revision!r}.",
+                suggestion="Give one of the files a new id; revision ids must be unique.",
+                details={"files": [str(seen[script.revision]), str(path)]},
+            )
+        seen[script.revision] = path
+        scripts.append(script)
+    return scripts
+
+
+def _top_level(tree: ast.Module, path: Path) -> tuple[dict[str, Any], set[str]]:
+    """Literal metadata assignments and function names defined at module level."""
+    values: dict[str, Any] = {}
+    functions: set[str] = set()
+    for node in tree.body:
+        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
+            functions.add(node.name)
+            continue
+        if isinstance(node, ast.Assign):
+            targets, value = node.targets, node.value
+        elif isinstance(node, ast.AnnAssign) and node.value is not None:
+            targets, value = [node.target], node.value
+        else:
+            continue
+        for target in targets:
+            if isinstance(target, ast.Name) and target.id in _METADATA_NAMES:
+                values[target.id] = _literal(value, target.id, path)
+    return values, functions
+
+
+def _literal(node: ast.expr, name: str, path: Path) -> Any:
+    try:
+        return ast.literal_eval(node)
+    except ValueError:
+        raise ScriptError(
+            f"{path.name}: `{name}` must be a literal value (string, tuple, bool, None).",
+            details={"path": str(path)},
+        ) from None
+
+
+def _id_tuple(value: Any, name: str, fail: Any) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, str):
+        items: tuple[Any, ...] = (value,)
+    elif isinstance(value, list | tuple):
+        items = tuple(value)
+    else:
+        raise fail(f"`{name}` must be None, a string, or a tuple of strings.")
+    if not all(isinstance(item, str) and item for item in items):
+        raise fail(f"`{name}` must contain only non-empty strings.")
+    if len(set(items)) != len(items):
+        raise fail(f"`{name}` contains duplicates.")
+    return items
+
+
+def _message(tree: ast.Module) -> str:
+    doc = ast.get_docstring(tree)
+    if not doc:
+        return ""
+    return doc.strip().splitlines()[0].strip()

mongomig/migrations/tracker.py

@@ -0,0 +1,149 @@
+"""The ``__mongomig_migrations`` collection: which revisions have been applied."""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import socket
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import TYPE_CHECKING, Any, Literal
+
+from mongomig._version import __version__
+from mongomig.migrations.graph import RevisionGraph
+from mongomig.migrations.script import Script
+
+if TYPE_CHECKING:
+    from pymongo.collection import Collection
+    from pymongo.database import Database
+
+Status = Literal["applied", "failed", "running"]
+
+
+@dataclass(frozen=True)
+class AppliedRecord:
+    revision: str
+    status: Status
+    down_revisions: tuple[str, ...]
+    description: str
+    checksum: str | None
+    applied_at: datetime | None
+    execution_time_ms: int | None
+
+    @classmethod
+    def from_doc(cls, doc: dict[str, Any]) -> AppliedRecord:
+        down = doc.get("down_revision")
+        if down is None:
+            downs: tuple[str, ...] = ()
+        elif isinstance(down, str):
+            downs = (down,)
+        else:
+            downs = tuple(down)
+        return cls(
+            revision=str(doc["_id"]),
+            status=doc.get("status", "applied"),
+            down_revisions=downs,
+            description=doc.get("description", ""),
+            checksum=doc.get("checksum"),
+            applied_at=doc.get("applied_at"),
+            execution_time_ms=doc.get("execution_time_ms"),
+        )
+
+
+@dataclass(frozen=True)
+class CurrentState:
+    """Where the database is relative to the revision files."""
+
+    applied_heads: list[str]
+    pending: list[str]
+    unknown: list[str]  # applied in the DB but no revision file (DB ahead of this code)
+    failed: list[str]
+
+
+class MigrationTracker:
+    def __init__(self, db: Database[dict[str, Any]], collection_name: str) -> None:
+        self.db = db
+        self.collection_name = collection_name
+
+    @property
+    def collection(self) -> Collection[dict[str, Any]]:
+        return self.db[self.collection_name]
+
+    def exists(self) -> bool:
+        return self.collection_name in self.db.list_collection_names(
+            filter={"name": self.collection_name}
+        )
+
+    def ensure(self) -> None:
+        """Create the tracking collection and its indexes (idempotent, race-safe)."""
+        from pymongo.errors import CollectionInvalid
+
+        if not self.exists():
+            with contextlib.suppress(CollectionInvalid):  # another runner created it first
+                self.db.create_collection(self.collection_name)
+        self.collection.create_index("status", name="status_1")
+        self.collection.create_index("applied_at", name="applied_at_1")
+
+    def records(self) -> list[AppliedRecord]:
+        """All records (any status); read-only, safe when the collection doesn't exist."""
+        docs = self.collection.find({}, sort=[("applied_at", 1), ("_id", 1)])
+        return [AppliedRecord.from_doc(doc) for doc in docs]
+
+    def applied_ids(self) -> set[str]:
+        return {r.revision for r in self.records() if r.status == "applied"}
+
+    def record_applied(
+        self, script: Script, *, execution_time_ms: int, environment: str | None = None
+    ) -> None:
+        self.collection.replace_one(
+            {"_id": script.revision},
+            self._doc(script, "applied", execution_time_ms, environment),
+            upsert=True,
+        )
+
+    def record_failed(self, script: Script, *, error: str, environment: str | None = None) -> None:
+        doc = self._doc(script, "failed", None, environment)
+        doc["error"] = error
+        self.collection.replace_one({"_id": script.revision}, doc, upsert=True)
+
+    def remove(self, revision: str) -> None:
+        self.collection.delete_one({"_id": revision})
+
+    def _doc(
+        self,
+        script: Script,
+        status: Status,
+        execution_time_ms: int | None,
+        environment: str | None,
+    ) -> dict[str, Any]:
+        downs = script.down_revisions
+        return {
+            "_id": script.revision,
+            "revision": script.revision,
+            "down_revision": None if not downs else downs[0] if len(downs) == 1 else list(downs),
+            "description": script.message,
+            "status": status,
+            "checksum": script.checksum,
+            "applied_at": datetime.now(UTC),
+            "execution_time_ms": execution_time_ms,
+            "mongomig_version": __version__,
+            "meta": {
+                "environment": environment,
+                "hostname": socket.gethostname(),
+                "user": os.environ.get("USER") or os.environ.get("USERNAME"),
+            },
+        }
+
+
+def compute_state(graph: RevisionGraph, records: list[AppliedRecord]) -> CurrentState:
+    applied = {r.revision for r in records if r.status == "applied"}
+    failed = [r.revision for r in records if r.status == "failed"]
+    known_applied = {rev for rev in applied if rev in graph}
+    heads = [
+        rev
+        for rev in graph.topological_order()
+        if rev in known_applied and not (graph.children[rev] & known_applied)
+    ]
+    pending = [rev for rev in graph.topological_order() if rev not in applied]
+    unknown = sorted(applied - set(graph.scripts))
+    return CurrentState(applied_heads=heads, pending=pending, unknown=unknown, failed=failed)

mongomig/output/console.py

@@ -0,0 +1,88 @@
+"""Single place that decides how results, warnings and errors reach the terminal.
+
+Human mode renders with rich (colour only on a TTY). ``--json`` mode writes exactly one JSON
+document to stdout per command; warnings go to stderr so stdout stays machine-parseable.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TextIO
+
+if TYPE_CHECKING:
+    from rich.console import Console
+
+    from mongomig.errors import MongoMigError
+
+
+class Output:
+    def __init__(
+        self,
+        *,
+        json_mode: bool = False,
+        verbose: bool = False,
+        stdout: TextIO | None = None,
+        stderr: TextIO | None = None,
+    ) -> None:
+        self.json_mode = json_mode
+        self.verbose = verbose
+        self._stdout = stdout
+        self._stderr = stderr
+        self._console: Console | None = None
+        self._err_console: Console | None = None
+
+    @property
+    def console(self) -> Console:
+        if self._console is None:
+            from rich.console import Console
+
+            self._console = Console(
+                file=self._stdout or sys.stdout, highlight=False, soft_wrap=True
+            )
+        return self._console
+
+    @property
+    def err_console(self) -> Console:
+        if self._err_console is None:
+            from rich.console import Console
+
+            self._err_console = Console(
+                file=self._stderr or sys.stderr, highlight=False, soft_wrap=True
+            )
+        return self._err_console
+
+    def result(self, data: Any, render: Callable[[Console], None]) -> None:
+        """Emit a command result: ``data`` as JSON, or ``render`` for humans."""
+        if self.json_mode:
+            stream = self._stdout or sys.stdout
+            stream.write(json.dumps(data, indent=2, default=str) + "\n")
+        else:
+            render(self.console)
+
+    def warn(self, message: str) -> None:
+        if self.json_mode:
+            (self._stderr or sys.stderr).write(f"warning: {message}\n")
+        else:
+            self.err_console.print(f"[yellow]warning:[/yellow] {message}")
+
+    def info(self, message: str) -> None:
+        """Verbose-only progress messages."""
+        if self.verbose and not self.json_mode:
+            self.err_console.print(f"[dim]{message}[/dim]")
+
+    def error(self, err: MongoMigError) -> None:
+        if self.json_mode:
+            stream = self._stdout or sys.stdout
+            stream.write(json.dumps({"error": err.to_dict()}, indent=2, default=str) + "\n")
+            return
+        from rich.markup import escape
+
+        con = self.err_console
+        con.print(f"[bold red]error:[/bold red] {escape(err.message)}")
+        if self.verbose and err.details:
+            for key, value in err.details.items():
+                con.print(f"  [dim]{key}:[/dim] {escape(str(value))}")
+        if err.suggestion:
+            con.print(f"[cyan]hint:[/cyan] {escape(err.suggestion)}")

mongomig/schema/snapshot.py

@@ -0,0 +1,35 @@
+"""``schema_snapshot.json``: the committed "last known expected schema".
+
+M1 only needs the empty document and a stable hash; the schema contents arrive in M3.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+SNAPSHOT_FORMAT = 1
+
+
+def empty_snapshot() -> dict[str, Any]:
+    return {"mongomig_format": SNAPSHOT_FORMAT, "storage": None, "collections": {}}
+
+
+def canonical_json(data: Any) -> str:
+    """Deterministic serialisation so git diffs and hashes are stable."""
+    return json.dumps(data, indent=2, sort_keys=True, ensure_ascii=False) + "\n"
+
+
+def snapshot_hash(path: Path) -> str | None:
+    """Hash of the snapshot's *content* (not formatting); ``None`` if there is no snapshot."""
+    if not path.is_file():
+        return None
+    data = json.loads(path.read_text(encoding="utf-8"))
+    compact = json.dumps(data, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
+    return "sha256:" + hashlib.sha256(compact.encode("utf-8")).hexdigest()
+
+
+def write_snapshot(path: Path, data: dict[str, Any]) -> None:
+    path.write_text(canonical_json(data), encoding="utf-8")

mongomig/templates/env.py.tmpl

@@ -0,0 +1,18 @@
+"""MongoMig environment.
+
+This file is plain Python, executed by MongoMig with your project root on sys.path, so you
+can import your application's models here (just like Alembic's env.py).
+
+`target_metadata` tells MongoMig which collections/models it should manage. Model-based
+autogeneration is not wired up yet; until then leave it as None and write migrations by hand
+with `mongomig revision -m "..."`.
+
+Coming soon:
+
+    from mongomig import MongoMetadata
+    import app.models  # noqa: F401  (importing registers @collection models)
+
+    target_metadata = MongoMetadata.default(storage="python", by_alias=True)
+"""
+
+target_metadata = None

mongomig/templates/mongomig.yaml.tmpl

@@ -0,0 +1,25 @@
+# MongoMig configuration.
+# Values support ${VAR} and ${VAR:-default} environment-variable interpolation.
+# Environment overlays: mongomig.<env>.yaml is merged on top when running
+# `mongomig --env <env> ...` (or MONGOMIG_ENV=<env>).
+
+database:
+  # Never put credentials here — keep them in the environment.
+  uri: ${MONGODB_URI}
+  # Optional if the URI already contains the database name.
+  name: ${MONGODB_DATABASE:-app}
+  server_selection_timeout_ms: 5000
+
+migrations:
+  directory: {{migrations_dir}}
+  tracking_collection: __mongomig_migrations
+  lock_collection: __mongomig_lock
+
+execution:
+  batch_size: 1000
+  sleep_ms_between_batches: 0
+  max_retries: 3
+  lock_ttl_seconds: 300
+
+sampling:
+  size: 10000

mongomig/templates/revision.py.tmpl

@@ -0,0 +1,22 @@
+"""${message}
+
+Revision: ${revision}
+Revises: ${down_revision_text}
+Created: ${created} UTC
+"""
+
+revision = ${revision_repr}
+down_revision = ${down_revision_repr}
+branch_labels = None
+depends_on = None
+reversible = True
+snapshot_hash = ${snapshot_hash_repr}
+mongomig_format = 1
+
+
+def upgrade(ctx):
+    pass
+
+
+def downgrade(ctx):
+    pass