versioning/src/versiond/ingest.py

"""Version commits, project detection, coalescing and rate limiting."""

from __future__ import annotations

import asyncio
import logging
import os
import stat
import time
from dataclasses import dataclass
from pathlib import Path, PurePath
from typing import Any, Callable

from .db import Database
from .filters import Filters
from .store import BlobStore, sha256

log = logging.getLogger(__name__)

PROJECT_MARKERS = (
    (".git", "git"),
    ("pyproject.toml", "python"),
    ("setup.py", "python"),
    ("package.json", "node"),
    ("go.mod", "go"),
    ("Cargo.toml", "rust"),
    ("pom.xml", "maven"),
    ("build.gradle", "gradle"),
    ("composer.json", "php"),
    ("Gemfile", "ruby"),
    ("mix.exs", "elixir"),
    ("CMakeLists.txt", "cmake"),
    ("Makefile", "make"),
)


class Rejected(Exception):
    def __init__(self, reason: str, detail: str = ""):
        super().__init__(reason)
        self.reason = reason
        self.detail = detail


def is_within(path: str, parent: str) -> bool:
    return path == parent or path.startswith(parent.rstrip("/") + "/")


class Repository:
    """Owns the rules for turning file content into stored versions."""

    def __init__(self, db: Database, blobs: BlobStore, filters: Filters):
        self.db = db
        self.blobs = blobs
        self.filters = filters
        self._roots: list[dict[str, Any]] = []
        self._project_cache: dict[str, tuple[str, str]] = {}
        self.reload_roots()

    def reload_roots(self) -> None:
        self._roots = sorted(self.db.roots(), key=lambda r: len(r["path"]), reverse=True)
        self._project_cache.clear()

    def root_for(self, path: str) -> dict[str, Any] | None:
        for root in self._roots:
            if is_within(path, root["path"]):
                return root
        return None

    def relative(self, path: str, root: dict[str, Any] | None = None) -> PurePath:
        """Path used for ignore rules: relative to its root, else to $HOME."""
        root = root if root is not None else self.root_for(path)
        if root is not None:
            return PurePath(path).relative_to(root["path"])
        home = str(Path.home())
        if is_within(path, home):
            return PurePath(path).relative_to(home)
        return PurePath(path.lstrip("/"))

    def check_path(self, path: str) -> None:
        reason = self.filters.path_reason(self.relative(path))
        if reason:
            raise Rejected(reason, path)

    def check_content(self, content: bytes) -> None:
        reason = self.filters.size_reason(len(content)) or self.filters.content_reason(content)
        if reason:
            raise Rejected(reason, f"{len(content)} bytes")

    def read_file(self, path: str) -> tuple[bytes, os.stat_result]:
        """Read a regular file without following a final symlink."""
        fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_CLOEXEC)
        try:
            st = os.fstat(fd)
            if not stat.S_ISREG(st.st_mode):
                raise Rejected("not-a-regular-file", path)
            reason = self.filters.size_reason(st.st_size)
            if reason:
                raise Rejected(reason, f"{st.st_size} bytes")
            with os.fdopen(fd, "rb", closefd=False) as fh:
                content = fh.read(self.filters.max_file_bytes + 1)
        finally:
            os.close(fd)
        return content, st

    # projects

    def _detect_project(self, path: str, root: dict[str, Any] | None) -> tuple[str, str]:
        directory = str(Path(path).parent)
        if directory in self._project_cache:
            return self._project_cache[directory]
        stop = root["path"] if root else directory
        current = Path(directory)
        result = (stop, "root" if root else "directory")
        while is_within(str(current), stop):
            found = next(
                (kind for marker, kind in PROJECT_MARKERS if (current / marker).exists()),
                None,
            )
            if found:
                result = (str(current), found)
                break
            if str(current) == stop or current.parent == current:
                break
            current = current.parent
        self._project_cache[directory] = result
        return result

    def _project_id(self, path: str, root: dict[str, Any] | None) -> int:
        project_path, kind = self._detect_project(path, root)
        row = self.db.one("SELECT id FROM projects WHERE path = ?", (project_path,))
        if row:
            return row["id"]
        cur = self.db.execute(
            "INSERT INTO projects(root_id, path, detected_by, created_at) VALUES (?, ?, ?, ?)",
            (root["id"] if root else None, project_path, kind, time.time()),
        )
        return cur.lastrowid

    def _ensure_file(self, path: str) -> dict[str, Any]:
        row = self.db.file_by_path(path)
        if row:
            return row
        root = self.root_for(path)
        project_id = self._project_id(path, root)
        self.db.execute(
            "INSERT INTO files(root_id, project_id, path, first_seen_at) VALUES (?, ?, ?, ?)",
            (root["id"] if root else None, project_id, path, time.time()),
        )
        return self.db.file_by_path(path)

    # commits

    def latest_digest(self, path: str) -> str | None:
        row = self.db.file_by_path(path)
        if not row:
            return None
        latest = self.db.latest_version(row["id"])
        return latest["blob_sha256"] if latest else None

    def commit(
        self,
        path: str,
        content: bytes,
        source: str,
        reason: str,
        st: os.stat_result | None = None,
        digest: str | None = None,
    ) -> dict[str, Any]:
        digest = digest or sha256(content)
        now = time.time()
        stat_fields = (
            (st.st_ino, st.st_mtime_ns, st.st_size) if st else (None, None, len(content))
        )
        with self.db.transaction():
            file = self._ensure_file(path)
            latest = self.db.latest_version(file["id"])
            if latest and latest["blob_sha256"] == digest:
                self.db.execute(
                    "UPDATE files SET exists_on_disk = 1, inode = coalesce(?, inode), "
                    "mtime_ns = coalesce(?, mtime_ns), size = ? WHERE id = ?",
                    (*stat_fields, file["id"]),
                )
                return {"status": "unchanged", "version_id": latest["id"], "path": path}
            if not self.db.one("SELECT 1 FROM blobs WHERE sha256 = ?", (digest,)):
                _, stored = self.blobs.put(content, digest)
                self.db.execute(
                    "INSERT INTO blobs(sha256, size, stored_size, created_at) VALUES (?, ?, ?, ?)",
                    (digest, len(content), stored, now),
                )
            cur = self.db.execute(
                "INSERT INTO versions(file_id, blob_sha256, captured_at, source, reason) "
                "VALUES (?, ?, ?, ?, ?)",
                (file["id"], digest, now, source, reason),
            )
            self.db.execute(
                "UPDATE files SET exists_on_disk = 1, inode = ?, mtime_ns = ?, size = ?, "
                "last_version_at = ? WHERE id = ?",
                (*stat_fields, now, file["id"]),
            )
            if file["project_id"]:
                self.db.execute(
                    "UPDATE projects SET last_activity_at = ? WHERE id = ?",
                    (now, file["project_id"]),
                )
        log.debug("committed %s (%s, %s)", path, source, reason)
        return {"status": "committed", "version_id": cur.lastrowid, "path": path}

    def mark_deleted(self, path: str) -> None:
        self.db.execute("UPDATE files SET exists_on_disk = 0 WHERE path = ?", (path,))

    def mark_deleted_tree(self, directory: str) -> None:
        self.db.execute(
            "UPDATE files SET exists_on_disk = 0 WHERE path LIKE ? ESCAPE '\\'",
            (_like_prefix(directory),),
        )

    def rename(self, old: str, new: str) -> bool:
        """Move a file's history to a new path. False if not applicable."""
        file = self.db.file_by_path(old)
        if not file or self.db.file_by_path(new):
            return False
        with self.db.transaction():
            self.db.execute("UPDATE files SET path = ? WHERE id = ?", (new, file["id"]))
            self.db.execute(
                "INSERT INTO renames(file_id, old_path, new_path, at) VALUES (?, ?, ?, ?)",
                (file["id"], old, new, time.time()),
            )
        return True

    def rename_tree(self, old_dir: str, new_dir: str) -> int:
        rows = self.db.all(
            "SELECT id, path FROM files WHERE path LIKE ? ESCAPE '\\'", (_like_prefix(old_dir),)
        )
        moved = 0
        with self.db.transaction():
            for row in rows:
                new_path = new_dir + row["path"][len(old_dir):]
                if self.db.file_by_path(new_path):
                    continue
                self.db.execute("UPDATE files SET path = ? WHERE id = ?", (new_path, row["id"]))
                self.db.execute(
                    "INSERT INTO renames(file_id, old_path, new_path, at) VALUES (?, ?, ?, ?)",
                    (row["id"], row["path"], new_path, time.time()),
                )
                moved += 1
        self._project_cache.clear()
        return moved


def _like_prefix(directory: str) -> str:
    escaped = directory.rstrip("/").replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
    return escaped + "/%"


class TokenBucket:
    def __init__(self, capacity: float, per_seconds: float, clock: Callable[[], float]):
        self.capacity = capacity
        self.rate = capacity / per_seconds
        self.tokens = capacity
        self.clock = clock
        self.stamp = clock()

    def _refill(self) -> None:
        now = self.clock()
        self.tokens = min(self.capacity, self.tokens + (now - self.stamp) * self.rate)
        self.stamp = now

    def take(self) -> bool:
        self._refill()
        if self.tokens >= 1:
            self.tokens -= 1
            return True
        return False

    def wait_time(self) -> float:
        self._refill()
        return max(0.0, (1 - self.tokens) / self.rate)


@dataclass
class _Pending:
    opened_at: float
    content: bytes | None = None
    digest: str | None = None
    source: str = ""
    reason: str = ""
    st: os.stat_result | None = None
    handle: asyncio.TimerHandle | None = None


class Coalescer:
    """Keeps the first and last state of a burst of changes to one path.

    The first change is committed immediately and opens a window. Changes
    inside the window replace each other; when the window closes (debounced,
    capped by max_hold) only the last one is committed. A per-path token
    bucket caps committed versions per hour; over the limit, changes keep
    coalescing until a token is free, so the newest state is never lost.
    """

    def __init__(
        self,
        repo: Repository,
        window: float,
        max_hold: float,
        per_hour: int,
        clock: Callable[[], float] = time.monotonic,
    ):
        self.repo = repo
        self.window = window
        self.max_hold = max_hold
        self.per_hour = per_hour
        self.clock = clock
        self.pending: dict[str, _Pending] = {}
        self.buckets: dict[str, TokenBucket] = {}

    def _bucket(self, path: str) -> TokenBucket:
        bucket = self.buckets.get(path)
        if bucket is None:
            bucket = self.buckets[path] = TokenBucket(self.per_hour, 3600, self.clock)
        return bucket

    def _schedule(self, path: str, p: _Pending, delay: float) -> None:
        if p.handle:
            p.handle.cancel()
        p.handle = asyncio.get_running_loop().call_later(delay, self._close, path)

    def submit(
        self,
        path: str,
        content: bytes,
        source: str,
        reason: str,
        st: os.stat_result | None = None,
    ) -> dict[str, Any]:
        digest = sha256(content)
        now = self.clock()
        p = self.pending.get(path)
        if p is None:
            if self.repo.latest_digest(path) == digest:
                return self.repo.commit(path, content, source, reason, st, digest)
            p = self.pending[path] = _Pending(opened_at=now)
            if self._bucket(path).take():
                result = self.repo.commit(path, content, source, reason, st, digest)
                self._schedule(path, p, self.window)
                return result
            p.content, p.digest, p.source, p.reason, p.st = content, digest, source, reason, st
            self._schedule(path, p, self._bucket(path).wait_time())
            return {"status": "coalesced", "version_id": None, "path": path}

        if p.content is None and self.repo.latest_digest(path) == digest:
            return {"status": "unchanged", "version_id": None, "path": path}
        if p.digest == digest:
            return {"status": "unchanged", "version_id": None, "path": path}
        p.content, p.digest, p.source, p.reason, p.st = content, digest, source, reason, st
        deadline = min(now + self.window, p.opened_at + self.max_hold)
        self._schedule(path, p, max(0.0, deadline - now))
        return {"status": "coalesced", "version_id": None, "path": path}

    def _close(self, path: str) -> None:
        p = self.pending.get(path)
        if p is None:
            return
        if p.content is None:
            del self.pending[path]
            return
        bucket = self._bucket(path)
        if not bucket.take():
            self._schedule(path, p, bucket.wait_time())
            return
        try:
            self.repo.commit(path, p.content, p.source, p.reason, p.st, p.digest)
        except Exception:
            log.exception("commit failed for %s", path)
        # Reopen an empty window so continuous editing keeps coalescing.
        self.pending[path] = fresh = _Pending(opened_at=self.clock())
        self._schedule(path, fresh, self.window)

    def flush(self) -> int:
        """Commit every waiting change immediately (used on shutdown)."""
        committed = 0
        for path, p in list(self.pending.items()):
            if p.handle:
                p.handle.cancel()
            if p.content is not None:
                try:
                    self.repo.commit(path, p.content, p.source, p.reason, p.st, p.digest)
                    committed += 1
                except Exception:
                    log.exception("flush commit failed for %s", path)
        self.pending.clear()
        return committed