notebook-snapshot-extension/snapshot/store.py

"""Snapshot store: sanitize, hash, persist, and index notebook versions.

Wraps a StorageAdapter and provides:
- Manual commits with name + description
- Gzipped JSON envelopes that embed the metadata, so the manifest is
  fully rebuildable from the version files alone.
- Atomic manifest writes (delegated to the StorageAdapter).
- Skip-on-duplicate: if the new content hash matches the most recent
  version, ``commit`` raises ``SnapshotUnchangedError`` instead of saving.
"""

import copy
import gzip
import hashlib
import json
import re
import time

from .storage import StorageAdapter


_VERSION_FILE_RE = re.compile(r"^v_.*\.ipynb\.gz$")


class SnapshotUnchangedError(Exception):
    """Raised when the new content hash matches the most recent version."""


class SnapshotStore:
    """High-level notebook version operations on a StorageAdapter."""

    def __init__(self, storage: StorageAdapter) -> None:
        self._storage = storage

    # ---- public API --------------------------------------------------------

    def commit(
        self,
        path: str,
        content: dict,
        name: str,
        description: str = "",
    ) -> dict:
        """Persist a new snapshot and append it to the manifest.

        If the new content hash matches the most recent version's hash,
        raises :class:`SnapshotUnchangedError` without writing anything.
        This avoids accidental duplicate snapshots when the user re-fires
        the commit command without making changes.

        Parameters
        ----------
        path: str
            The notebook's relative path (acts as a key prefix).
        content: dict
            The notebook JSON as sent by the frontend.
        name: str
            User-provided snapshot name.
        description: str, optional
            User-provided snapshot description.

        Returns
        -------
        dict
            The manifest entry (metadata only — no notebook body).

        Raises
        ------
        SnapshotUnchangedError
            If the cleaned content hash equals the most recent version's
            hash. The route layer translates this into a 200 response
            with ``{"skipped": True, "reason": "unchanged"}``.
        """
        cleaned = self._sanitize(content)
        hash_ = self._hash(cleaned)

        # Skip-on-duplicate check (compare to the most recent version only).
        manifest = self._read_manifest(path) or {"versions": []}
        existing = manifest.get("versions", [])
        if existing and existing[-1].get("hash") == hash_:
            raise SnapshotUnchangedError(
                f"内容未变更（与最近版本 {existing[-1]['id']} hash 相同），已跳过保存"
            )

        timestamp = int(time.time() * 1_000_000)
        version_id = f"v_{timestamp}_{hash_[:8]}"
        envelope = {
            "id": version_id,
            "timestamp": timestamp,
            "name": name,
            "description": description,
            "hash": hash_,
            "size": 0,  # filled in after first encoding pass
            "notebook": cleaned,
        }
        file_bytes = self._encode_envelope(envelope)
        envelope["size"] = len(file_bytes)
        file_bytes = self._encode_envelope(envelope)
        self._storage.put(self._version_key(path, version_id), file_bytes)
        entry = {
            k: envelope[k]
            for k in ("id", "timestamp", "name", "description", "hash", "size")
        }
        self._append_manifest(path, entry)
        return entry

    def list_versions(self, path: str) -> list[dict]:
        """Return the versions for ``path``, newest first.

        Rebuilds the manifest from the version files if the manifest is
        missing or unparseable.
        """
        manifest = self._read_manifest(path)
        if manifest is None:
            manifest = self._rebuild_manifest(path)
        return sorted(
            manifest.get("versions", []),
            key=lambda e: e.get("timestamp", 0),
            reverse=True,
        )

    def get_notebook(self, path: str, version_id: str) -> dict:
        """Return the cleaned notebook body for ``version_id``.

        Raises
        ------
        KeyError
            If no such version exists.
        """
        key = self._version_key(path, version_id)
        try:
            raw = self._storage.get(key)
        except KeyError:
            raise KeyError(version_id) from None
        envelope = self._decode_envelope(raw)
        return envelope["notebook"]

    # ---- internals ---------------------------------------------------------

    def _prefix(self, path: str) -> str:
        return f"{path}/"

    def _manifest_key(self, path: str) -> str:
        return f"{self._prefix(path)}manifest.json"

    def _version_key(self, path: str, version_id: str) -> str:
        return f"{self._prefix(path)}{version_id}.ipynb.gz"

    def _sanitize(self, content: dict) -> dict:
        """Strip all code-cell outputs and execution counts."""
        cleaned = copy.deepcopy(content)
        for cell in cleaned.get("cells", []):
            if cell.get("cell_type") == "code":
                cell["outputs"] = []
                cell["execution_count"] = None
        return cleaned

    def _hash(self, cleaned: dict) -> str:
        encoded = json.dumps(
            cleaned, sort_keys=True, separators=(",", ":")
        ).encode()
        return hashlib.md5(encoded).hexdigest()

    def _encode_envelope(self, envelope: dict) -> bytes:
        return gzip.compress(json.dumps(envelope, ensure_ascii=False).encode())

    def _decode_envelope(self, raw: bytes) -> dict:
        return json.loads(gzip.decompress(raw).decode())

    def _read_manifest(self, path: str):
        """Return the parsed manifest, ``None`` if missing or corrupt.

        Returning ``None`` (instead of an empty stub) signals the caller to
        rebuild from version files, which also writes a fresh manifest.
        """
        try:
            raw = self._storage.get(self._manifest_key(path))
        except KeyError:
            return None
        try:
            return json.loads(raw.decode())
        except (ValueError, UnicodeDecodeError):
            return None

    def _write_manifest(self, path: str, manifest: dict) -> None:
        payload = json.dumps(manifest, ensure_ascii=False).encode()
        self._storage.put(self._manifest_key(path), payload)

    def _append_manifest(self, path: str, entry: dict) -> None:
        manifest = self._read_manifest(path) or {"versions": []}
        manifest.setdefault("versions", []).append(entry)
        self._write_manifest(path, manifest)

    def _rebuild_manifest(self, path: str) -> dict:
        prefix = self._prefix(path)
        versions: list[dict] = []
        for key in self._storage.list_keys(prefix=prefix):
            tail = key[len(prefix):]
            if not _VERSION_FILE_RE.match(tail):
                continue
            try:
                envelope = self._decode_envelope(self._storage.get(key))
            except (KeyError, ValueError, OSError):
                continue
            versions.append(
                {
                    "id": envelope.get("id"),
                    "timestamp": envelope.get("timestamp", 0),
                    "name": envelope.get("name", ""),
                    "description": envelope.get("description", ""),
                    "hash": envelope.get("hash", ""),
                    "size": envelope.get("size", 0),
                }
            )
        versions.sort(key=lambda e: e.get("timestamp", 0), reverse=True)
        manifest = {"versions": versions}
        self._write_manifest(path, manifest)
        return manifest