mongomig 0.1.0.dev0__py3-none-any.whl

mongomig/config/loader.py

@@ -0,0 +1,178 @@
+"""Locate, read, interpolate and validate ``mongomig.yaml`` (+ environment overlays)."""
+
+from __future__ import annotations
+
+import os
+import re
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+from mongomig.config.models import DEFAULT_CONFIG_FILENAME, LoadedConfig, MongoMigConfig
+from mongomig.errors import ConfigError
+
+CONFIG_ENV_VAR = "MONGOMIG_CONFIG"
+ENVIRONMENT_ENV_VAR = "MONGOMIG_ENV"
+
+# ${VAR} or ${VAR:-default}
+_VAR_RE = re.compile(r"\$\{(?P<name>[A-Za-z_][A-Za-z0-9_]*)(?::-(?P<default>[^}]*))?\}")
+_ENV_NAME_RE = re.compile(r"^[A-Za-z0-9_-]+$")
+
+
+def find_config(start: Path | None = None) -> Path | None:
+    """Walk up from ``start`` (default: cwd) looking for ``mongomig.yaml``."""
+    current = (start or Path.cwd()).resolve()
+    for directory in (current, *current.parents):
+        candidate = directory / DEFAULT_CONFIG_FILENAME
+        if candidate.is_file():
+            return candidate
+    return None
+
+
+def load_config(
+    config_path: Path | None = None,
+    *,
+    environment: str | None = None,
+    env: Mapping[str, str] | None = None,
+    start_dir: Path | None = None,
+) -> LoadedConfig:
+    """Load configuration.
+
+    Resolution order for the file: explicit ``config_path`` → ``$MONGOMIG_CONFIG`` → search
+    upward from ``start_dir``. The environment overlay (``mongomig.<environment>.yaml``) comes
+    from ``environment`` or ``$MONGOMIG_ENV`` and is deep-merged over the base file.
+    """
+    env = os.environ if env is None else env
+
+    path = _resolve_config_path(config_path, env, start_dir)
+    environment = environment or env.get(ENVIRONMENT_ENV_VAR) or None
+
+    raw = _read_yaml(path)
+    warnings = _credential_warnings(raw, path)
+
+    if environment:
+        if not _ENV_NAME_RE.match(environment):
+            raise ConfigError(
+                f"Invalid environment name {environment!r}.",
+                suggestion="Use letters, digits, '-' or '_' (e.g. --env production).",
+            )
+        overlay_path = path.with_name(f"{path.stem}.{environment}{path.suffix}")
+        if not overlay_path.is_file():
+            raise ConfigError(
+                f"No configuration found for environment {environment!r}.",
+                suggestion=f"Create {overlay_path.name} next to {path.name}.",
+                details={"expected_path": str(overlay_path)},
+            )
+        overlay = _read_yaml(overlay_path)
+        warnings += _credential_warnings(overlay, overlay_path)
+        raw = deep_merge(raw, overlay)
+
+    missing: set[str] = set()
+    interpolated = interpolate(raw, env, missing)
+
+    from pydantic import ValidationError
+
+    try:
+        settings = MongoMigConfig.model_validate(interpolated)
+    except ValidationError as exc:
+        problems = "; ".join(
+            f"{'.'.join(str(p) for p in err['loc']) or '<root>'}: {err['msg']}"
+            for err in exc.errors()
+        )
+        raise ConfigError(
+            f"Invalid configuration in {path.name}: {problems}",
+            suggestion="Check the file against the documented mongomig.yaml format.",
+            details={"path": str(path)},
+        ) from None
+
+    return LoadedConfig(
+        settings=settings,
+        config_path=path,
+        environment=environment,
+        missing_env_vars=frozenset(missing),
+        warnings=tuple(warnings),
+    )
+
+
+def interpolate(value: Any, env: Mapping[str, str], missing: set[str]) -> Any:
+    """Recursively expand ``${VAR}`` / ``${VAR:-default}`` in string values.
+
+    Unset variables without a default are left as-is and recorded in ``missing`` so commands
+    that don't need them (e.g. ``history``) keep working offline.
+    """
+    if isinstance(value, str):
+
+        def _sub(match: re.Match[str]) -> str:
+            name, default = match.group("name"), match.group("default")
+            if name in env:
+                return env[name]
+            if default is not None:
+                return default
+            missing.add(name)
+            return match.group(0)
+
+        return _VAR_RE.sub(_sub, value)
+    if isinstance(value, dict):
+        return {k: interpolate(v, env, missing) for k, v in value.items()}
+    if isinstance(value, list):
+        return [interpolate(v, env, missing) for v in value]
+    return value
+
+
+def deep_merge(base: Mapping[str, Any], overlay: Mapping[str, Any]) -> dict[str, Any]:
+    """Merge mappings recursively; overlay wins. Lists and scalars are replaced, not merged."""
+    merged = dict(base)
+    for key, value in overlay.items():
+        if isinstance(value, Mapping) and isinstance(merged.get(key), Mapping):
+            merged[key] = deep_merge(merged[key], value)
+        else:
+            merged[key] = value
+    return merged
+
+
+def _resolve_config_path(
+    config_path: Path | None, env: Mapping[str, str], start_dir: Path | None
+) -> Path:
+    if config_path is None and env.get(CONFIG_ENV_VAR):
+        config_path = Path(env[CONFIG_ENV_VAR])
+    if config_path is not None:
+        path = config_path.expanduser().resolve()
+        if not path.is_file():
+            raise ConfigError(f"Config file not found: {path}")
+        return path
+    found = find_config(start_dir)
+    if found is None:
+        raise ConfigError(
+            f"No {DEFAULT_CONFIG_FILENAME} found in this directory or any parent.",
+            suggestion="Run `mongomig init` to create one, or pass --config PATH.",
+        )
+    return found
+
+
+def _read_yaml(path: Path) -> dict[str, Any]:
+    import yaml
+
+    try:
+        data = yaml.safe_load(path.read_text(encoding="utf-8"))
+    except yaml.YAMLError as exc:
+        raise ConfigError(f"{path.name} is not valid YAML: {exc}") from None
+    except OSError as exc:
+        raise ConfigError(f"Cannot read {path}: {exc.strerror}") from None
+    if data is None:
+        return {}
+    if not isinstance(data, dict):
+        raise ConfigError(f"{path.name} must contain a YAML mapping at the top level.")
+    return data
+
+
+def _credential_warnings(raw: Mapping[str, Any], path: Path) -> list[str]:
+    from mongomig.database.redact import uri_has_inline_password
+
+    database = raw.get("database")
+    uri = database.get("uri") if isinstance(database, Mapping) else None
+    if isinstance(uri, str) and uri_has_inline_password(uri):
+        return [
+            f"{path.name} contains a password in database.uri. "
+            "Use an environment variable instead, e.g. uri: ${MONGODB_URI}"
+        ]
+    return []

mongomig/config/models.py

@@ -0,0 +1,87 @@
+"""Typed model of ``mongomig.yaml``."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field, PositiveInt, field_validator
+
+DEFAULT_CONFIG_FILENAME = "mongomig.yaml"
+DEFAULT_MIGRATIONS_DIR = "migrations"
+VERSIONS_DIRNAME = "versions"
+ENV_FILENAME = "env.py"
+SNAPSHOT_FILENAME = "schema_snapshot.json"
+
+
+class _Section(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+
+class DatabaseConfig(_Section):
+    uri: str
+    # Optional: falls back to the database named in the URI path.
+    name: str | None = None
+    server_selection_timeout_ms: PositiveInt = 5000
+
+    @field_validator("name")
+    @classmethod
+    def _empty_name_is_none(cls, value: str | None) -> str | None:
+        return value or None
+
+
+class MigrationsConfig(_Section):
+    directory: str = DEFAULT_MIGRATIONS_DIR
+    tracking_collection: str = "__mongomig_migrations"
+    lock_collection: str = "__mongomig_lock"
+
+
+class ExecutionConfig(_Section):
+    batch_size: PositiveInt = 1000
+    sleep_ms_between_batches: int = Field(default=0, ge=0)
+    max_retries: int = Field(default=3, ge=0)
+    lock_ttl_seconds: PositiveInt = 300
+
+
+class SamplingConfig(_Section):
+    size: PositiveInt = 10000
+
+
+class MongoMigConfig(_Section):
+    database: DatabaseConfig
+    migrations: MigrationsConfig = MigrationsConfig()
+    execution: ExecutionConfig = ExecutionConfig()
+    sampling: SamplingConfig = SamplingConfig()
+
+
+class LoadedConfig(BaseModel):
+    """A parsed config plus where it came from. Paths are absolute."""
+
+    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
+
+    settings: MongoMigConfig
+    config_path: Path
+    environment: str | None = None
+    # Environment variables referenced by the config but not set. Commands that need them
+    # (anything touching the database) fail with a clear error; offline commands still work.
+    missing_env_vars: frozenset[str] = frozenset()
+    warnings: tuple[str, ...] = ()
+
+    @property
+    def root_dir(self) -> Path:
+        return self.config_path.parent
+
+    @property
+    def migrations_dir(self) -> Path:
+        return (self.root_dir / self.settings.migrations.directory).resolve()
+
+    @property
+    def versions_dir(self) -> Path:
+        return self.migrations_dir / VERSIONS_DIRNAME
+
+    @property
+    def env_py_path(self) -> Path:
+        return self.migrations_dir / ENV_FILENAME
+
+    @property
+    def snapshot_path(self) -> Path:
+        return self.migrations_dir / SNAPSHOT_FILENAME

mongomig/database/client.py

@@ -0,0 +1,82 @@
+"""The only module that constructs PyMongo clients."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from mongomig.config.models import LoadedConfig
+from mongomig.database.redact import describe_hosts, redact_text
+from mongomig.errors import ConfigError, DatabaseError
+
+if TYPE_CHECKING:
+    from pymongo import MongoClient
+    from pymongo.database import Database
+
+
+def create_client(config: LoadedConfig) -> MongoClient[dict[str, Any]]:
+    db_cfg = config.settings.database
+    unresolved = sorted(v for v in config.missing_env_vars if f"${{{v}" in db_cfg.uri)
+    if unresolved:
+        raise ConfigError(
+            f"Environment variable {unresolved[0]} is not set (needed by database.uri).",
+            suggestion=f"export {unresolved[0]}=mongodb://... or add it to your deployment env.",
+        )
+
+    from pymongo import MongoClient
+    from pymongo.errors import ConfigurationError, InvalidURI
+
+    try:
+        return MongoClient(
+            db_cfg.uri,
+            appname="mongomig",
+            serverSelectionTimeoutMS=db_cfg.server_selection_timeout_ms,
+            tz_aware=True,
+            connect=False,
+        )
+    except (InvalidURI, ConfigurationError) as exc:
+        raise ConfigError(
+            f"Invalid MongoDB URI: {redact_text(str(exc))}",
+            suggestion="Check database.uri (format: mongodb://host:27017/dbname).",
+        ) from None
+
+
+def get_database(
+    client: MongoClient[dict[str, Any]], config: LoadedConfig
+) -> Database[dict[str, Any]]:
+    name = config.settings.database.name
+    if name and "${" in name:
+        raise ConfigError(
+            f"database.name references an unset environment variable: {name}",
+        )
+    if name:
+        return client[name]
+
+    from pymongo.errors import ConfigurationError
+
+    try:
+        return client.get_default_database()
+    except ConfigurationError:
+        raise ConfigError(
+            "No database name configured.",
+            suggestion="Set database.name in mongomig.yaml or include it in the URI path.",
+        ) from None
+
+
+def ping(client: MongoClient[dict[str, Any]], config: LoadedConfig) -> None:
+    """Fail fast with a readable, credential-free error when MongoDB is unreachable."""
+    from pymongo.errors import OperationFailure, PyMongoError
+
+    hosts = describe_hosts(config.settings.database.uri)
+    try:
+        client.admin.command("ping")
+    except OperationFailure as exc:
+        raise DatabaseError(
+            f"MongoDB rejected the connection to {hosts}: {redact_text(str(exc.details or exc))}",
+            suggestion="Check credentials and that the user has access to this database.",
+        ) from None
+    except PyMongoError as exc:
+        raise DatabaseError(
+            f"Cannot connect to MongoDB at {hosts}: {type(exc).__name__}",
+            suggestion="Is MongoDB running and reachable? For local dev: `docker compose up -d`.",
+            details={"driver_error": redact_text(str(exc))[:500]},
+        ) from None

mongomig/database/redact.py

@@ -0,0 +1,70 @@
+"""Keep credentials out of anything MongoMig prints or logs."""
+
+from __future__ import annotations
+
+import re
+from urllib.parse import parse_qsl, urlencode
+
+_SCHEME_RE = re.compile(r"^(?P<scheme>mongodb(?:\+srv)?://)(?P<rest>.*)$", re.IGNORECASE)
+_SENSITIVE_PARAM_RE = re.compile(r"(pass|secret|token|key|credential)", re.IGNORECASE)
+# user:password@ in a URI; used to spot URIs embedded in free text (e.g. driver error messages).
+_INLINE_URI_RE = re.compile(r"mongodb(?:\+srv)?://[^\s'\"]+", re.IGNORECASE)
+
+REDACTED = "***"
+
+
+def redact_uri(uri: str) -> str:
+    """Return ``uri`` with password and sensitive query parameters masked.
+
+    ``mongodb://user:pw@host/db?authSource=admin`` -> ``mongodb://user:***@host/db?authSource=admin``
+    """
+    match = _SCHEME_RE.match(uri.strip())
+    if not match:
+        return REDACTED
+    scheme, rest = match.group("scheme"), match.group("rest")
+
+    # Userinfo ends at the last '@' before the first '/' (passwords may contain '@' if unescaped).
+    path_start = rest.find("/")
+    authority = rest if path_start == -1 else rest[:path_start]
+    tail = "" if path_start == -1 else rest[path_start:]
+
+    if "@" in authority:
+        userinfo, hosts = authority.rsplit("@", 1)
+        user = userinfo.split(":", 1)[0]
+        authority = f"{user}:{REDACTED}@{hosts}" if ":" in userinfo else f"{user}@{hosts}"
+
+    if "?" in tail:
+        path, query = tail.split("?", 1)
+        params = [
+            (k, REDACTED if _SENSITIVE_PARAM_RE.search(k) else v)
+            for k, v in parse_qsl(query, keep_blank_values=True)
+        ]
+        tail = f"{path}?{urlencode(params, safe='*')}"
+
+    return f"{scheme}{authority}{tail}"
+
+
+def redact_text(text: str) -> str:
+    """Mask any MongoDB URI that appears inside arbitrary text."""
+    return _INLINE_URI_RE.sub(lambda m: redact_uri(m.group(0)), text)
+
+
+def uri_has_inline_password(uri: str) -> bool:
+    match = _SCHEME_RE.match(uri.strip())
+    if not match:
+        return False
+    rest = match.group("rest")
+    authority = rest.split("/", 1)[0]
+    if "@" not in authority:
+        return False
+    userinfo = authority.rsplit("@", 1)[0]
+    return ":" in userinfo and userinfo.split(":", 1)[1] != ""
+
+
+def describe_hosts(uri: str) -> str:
+    """Hosts portion of a URI, safe for messages like 'cannot connect to host:27017'."""
+    match = _SCHEME_RE.match(uri.strip())
+    if not match:
+        return "<unknown host>"
+    authority = match.group("rest").split("/", 1)[0].split("?", 1)[0]
+    return authority.rsplit("@", 1)[-1] or "<unknown host>"

mongomig/errors.py

@@ -0,0 +1,108 @@
+"""Exception hierarchy. Every error maps to a documented CLI exit code."""
+
+from __future__ import annotations
+
+from enum import IntEnum
+from typing import Any
+
+
+class ExitCode(IntEnum):
+    SUCCESS = 0
+    VALIDATION_FAILURE = 1
+    EXECUTION_FAILURE = 2
+    CONFIG_ERROR = 3
+    CONFLICT = 4
+    LOCK_FAILURE = 5
+    CHECKSUM_MISMATCH = 6
+
+
+class MongoMigError(Exception):
+    """Base class for all MongoMig errors.
+
+    ``suggestion`` is a human-readable recovery hint shown by the CLI.
+    ``details`` holds structured context (revision, collection, path, ...) for ``--json`` output.
+    """
+
+    exit_code: ExitCode = ExitCode.EXECUTION_FAILURE
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        suggestion: str | None = None,
+        details: dict[str, Any] | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.suggestion = suggestion
+        self.details = details or {}
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "type": type(self).__name__,
+            "message": self.message,
+            "suggestion": self.suggestion,
+            "exit_code": int(self.exit_code),
+            "details": self.details,
+        }
+
+
+# --- exit code 1: validation ---------------------------------------------------------------
+
+
+class ValidationError(MongoMigError):
+    exit_code = ExitCode.VALIDATION_FAILURE
+
+
+class ScriptError(ValidationError):
+    """A revision file is malformed."""
+
+
+class RevisionNotFoundError(ValidationError):
+    pass
+
+
+class AmbiguousRevisionError(ValidationError):
+    pass
+
+
+# --- exit code 2: execution ----------------------------------------------------------------
+
+
+class ExecutionError(MongoMigError):
+    exit_code = ExitCode.EXECUTION_FAILURE
+
+
+class DatabaseError(ExecutionError):
+    """Connecting to or talking to MongoDB failed."""
+
+
+# --- exit code 3: configuration ------------------------------------------------------------
+
+
+class ConfigError(MongoMigError):
+    exit_code = ExitCode.CONFIG_ERROR
+
+
+# --- exit code 4: conflicts in the revision graph ------------------------------------------
+
+
+class RevisionConflictError(MongoMigError):
+    """Duplicate revision ids, cycles, missing parents."""
+
+    exit_code = ExitCode.CONFLICT
+
+
+class MultipleHeadsError(RevisionConflictError):
+    pass
+
+
+# --- exit code 5 / 6 -----------------------------------------------------------------------
+
+
+class LockError(MongoMigError):
+    exit_code = ExitCode.LOCK_FAILURE
+
+
+class ChecksumMismatchError(MongoMigError):
+    exit_code = ExitCode.CHECKSUM_MISMATCH

mongomig/migrations/graph.py

@@ -0,0 +1,161 @@
+"""The revision DAG built from ``down_revision`` (and ``depends_on``) links."""
+
+from __future__ import annotations
+
+import heapq
+from collections.abc import Iterable
+
+from mongomig.errors import (
+    AmbiguousRevisionError,
+    MultipleHeadsError,
+    RevisionConflictError,
+    RevisionNotFoundError,
+)
+from mongomig.migrations.script import Script
+
+MIN_PREFIX_LENGTH = 4
+
+
+class RevisionGraph:
+    def __init__(self, scripts: Iterable[Script]) -> None:
+        self.scripts: dict[str, Script] = {}
+        for script in scripts:
+            if script.revision in self.scripts:
+                raise RevisionConflictError(f"Duplicate revision id {script.revision!r}.")
+            self.scripts[script.revision] = script
+
+        self.children: dict[str, set[str]] = {rev: set() for rev in self.scripts}
+        for script in self.scripts.values():
+            for parent in script.down_revisions:
+                if parent not in self.scripts:
+                    raise RevisionConflictError(
+                        f"Revision {script.revision} ({script.path.name}) points to "
+                        f"down_revision {parent!r}, which does not exist.",
+                        suggestion="Restore the missing revision file or fix down_revision.",
+                        details={"path": str(script.path)},
+                    )
+                self.children[parent].add(script.revision)
+            for dep in script.depends_on:
+                if self._lookup_label(dep) is None and dep not in self.scripts:
+                    raise RevisionConflictError(
+                        f"Revision {script.revision} depends_on {dep!r}, which does not exist.",
+                        details={"path": str(script.path)},
+                    )
+
+        self._order = self._topological_order()
+
+    def __len__(self) -> int:
+        return len(self.scripts)
+
+    def __contains__(self, rev: object) -> bool:
+        return rev in self.scripts
+
+    # --- structure ------------------------------------------------------------------------
+
+    def heads(self) -> list[str]:
+        """Revisions nothing builds on, in chronological order."""
+        return [rev for rev in self._order if not self.children[rev]]
+
+    def bases(self) -> list[str]:
+        return [rev for rev in self._order if self.scripts[rev].is_base]
+
+    def topological_order(self) -> list[str]:
+        """Parents before children; ties broken by file name (i.e. creation time)."""
+        return list(self._order)
+
+    def single_head(self) -> str | None:
+        heads = self.heads()
+        if len(heads) > 1:
+            raise MultipleHeadsError(
+                f"Multiple heads: {', '.join(heads)}.",
+                suggestion="Pass --head <revision> to choose a parent, or merge the branches.",
+                details={"heads": heads},
+            )
+        return heads[0] if heads else None
+
+    def ancestors(self, rev: str) -> set[str]:
+        """All revisions ``rev`` builds on (excluding itself)."""
+        seen: set[str] = set()
+        stack = list(self._dependencies(rev))
+        while stack:
+            current = stack.pop()
+            if current not in seen:
+                seen.add(current)
+                stack.extend(self._dependencies(current))
+        return seen
+
+    # --- lookup ---------------------------------------------------------------------------
+
+    def resolve(self, ref: str) -> str:
+        """Resolve a full id, a unique prefix (>= 4 chars), a branch label, or ``head``."""
+        ref = ref.strip()
+        if ref == "head":
+            head = self.single_head()
+            if head is None:
+                raise RevisionNotFoundError("There are no revisions yet.")
+            return head
+        if ref in self.scripts:
+            return ref
+        labelled = self._lookup_label(ref)
+        if labelled is not None:
+            return labelled
+        if len(ref) >= MIN_PREFIX_LENGTH:
+            matches = sorted(rev for rev in self.scripts if rev.startswith(ref))
+            if len(matches) == 1:
+                return matches[0]
+            if len(matches) > 1:
+                raise AmbiguousRevisionError(
+                    f"Revision prefix {ref!r} is ambiguous: {', '.join(matches)}.",
+                    suggestion="Use more characters of the revision id.",
+                )
+        hint = (
+            f"Prefixes need at least {MIN_PREFIX_LENGTH} characters."
+            if len(ref) < MIN_PREFIX_LENGTH
+            else "Run `mongomig history` to list revisions."
+        )
+        raise RevisionNotFoundError(f"Unknown revision {ref!r}.", suggestion=hint)
+
+    # --- internals ------------------------------------------------------------------------
+
+    def _dependencies(self, rev: str) -> tuple[str, ...]:
+        script = self.scripts[rev]
+        deps = [self._lookup_label(d) or d for d in script.depends_on]
+        return (*script.down_revisions, *deps)
+
+    def _lookup_label(self, label: str) -> str | None:
+        found = [rev for rev, s in self.scripts.items() if label in s.branch_labels]
+        if len(found) > 1:
+            raise RevisionConflictError(
+                f"Branch label {label!r} is used by several revisions: {', '.join(found)}."
+            )
+        return found[0] if found else None
+
+    def _topological_order(self) -> list[str]:
+        indegree = {rev: 0 for rev in self.scripts}
+        dependents: dict[str, list[str]] = {rev: [] for rev in self.scripts}
+        for rev in self.scripts:
+            for dep in set(self._dependencies(rev)):
+                indegree[rev] += 1
+                dependents[dep].append(rev)
+
+        def key(rev: str) -> tuple[str, str]:
+            return (self.scripts[rev].path.name, rev)
+
+        ready = [key(rev) for rev, n in indegree.items() if n == 0]
+        heapq.heapify(ready)
+        order: list[str] = []
+        while ready:
+            _, rev = heapq.heappop(ready)
+            order.append(rev)
+            for child in dependents[rev]:
+                indegree[child] -= 1
+                if indegree[child] == 0:
+                    heapq.heappush(ready, key(child))
+
+        if len(order) != len(self.scripts):
+            cyclic = sorted(rev for rev, n in indegree.items() if n > 0)
+            raise RevisionConflictError(
+                f"Revision graph contains a cycle involving: {', '.join(cyclic)}.",
+                suggestion="Fix down_revision/depends_on so revisions don't reference each other.",
+            )
+        return order