Source code for numeraire_dataset.zones.raw

"""Raw zone: the immutable download cache, one directory per ``(source, vintage)`` with a sidecar.

Rows and schema are serialized as returned by the connector and never edited in place — a re-pull
with different content is a new ``vintage`` directory, not an overwrite. Parquet and ``_meta.json``
are staged together before the complete directory is atomically published. The sidecar records
what the pull is, its digest contract, and whether it is itself point-in-time. The cache root is
user-configurable and defaults outside the repo (``paths.data_home``); **licensed data lives only
here, never in git**.
"""

from __future__ import annotations

import json
import stat
from dataclasses import asdict, dataclass
from pathlib import Path

from numeraire_dataset.paths import data_home
from numeraire_dataset.zones.steps import CONTENT_DIGEST_VERSION

# A pull is either a dated point-in-time vintage, or a latest-snapshot convenience pull (revised).
PIT_STATUSES = ("vintage", "snapshot")




@dataclass(frozen=True)
class RawMeta:
    """The ``_meta.json`` sidecar for one raw pull."""

    source: str
    vintage: str
    pulled_at: str
    content_digest: str
    row_count: int
    query_hash: str = ""
    pit_status: str = "vintage"
    content_digest_version: int = CONTENT_DIGEST_VERSION

    def __post_init__(self) -> None:
        if self.pit_status not in PIT_STATUSES:
            raise ValueError(f"pit_status must be one of {PIT_STATUSES}; got {self.pit_status!r}")
        if not isinstance(self.content_digest_version, int) or self.content_digest_version < 1:
            raise ValueError("content_digest_version must be a positive integer")



def _cache_component(value: str, *, label: str) -> str:
    """Reject path traversal before a cache component reaches any writer or cleanup path."""
    if (
        not value
        or value in {".", ".."}
        or any(char in value for char in ("/", "\\", "\0", "@", "#"))
    ):
        raise ValueError(f"{label} must be a non-empty path component without reserved separators")
    return value


def _raw_path(source: str, vintage: str, *, home: str | Path | None = None) -> Path:
    """Resolve a cache path without creating a vintage directory."""
    source = _cache_component(source, label="source")
    vintage = _cache_component(vintage, label="vintage")
    root = data_home(home).resolve()
    candidate = root / "raw" / source / vintage
    if not candidate.resolve(strict=False).is_relative_to(root):
        raise RuntimeError("raw cache path resolves outside the configured data home")
    return candidate


def _validate_cache_file(path: Path, *, label: str) -> Path:
    """Require a final cache artifact to be a regular file, never a symlink."""
    if path.is_symlink():
        raise RuntimeError(f"{label} is a symlink; inspect the immutable cache manually")
    mode = path.lstat().st_mode
    if not stat.S_ISREG(mode):
        raise RuntimeError(f"{label} is not a regular file; inspect the immutable cache manually")
    return path




def raw_dir(source: str, vintage: str, *, home: str | Path | None = None) -> Path:
    """Return and create the cache directory for a ``(source, vintage)`` pull."""
    directory = _raw_path(source, vintage, home=home)
    directory.mkdir(parents=True, exist_ok=True)
    return directory



def _write_meta_at(meta: RawMeta, directory: Path) -> Path:
    """Write ``meta`` into an unpublished staging directory."""
    directory.mkdir(parents=True, exist_ok=True)
    path = directory / "_meta.json"
    path.write_text(json.dumps(asdict(meta), sort_keys=True, indent=2), encoding="utf-8")
    return path




def write_meta(meta: RawMeta, *, home: str | Path | None = None) -> Path:
    """Create a new ``_meta.json`` sidecar without overwriting any cache artifact.

    Complete raw vintages are immutable. This public low-level helper therefore only writes into a
    new or empty directory; WRDS pulls use the private staging writer before atomically publishing
    the whole directory.
    """
    directory = _raw_path(meta.source, meta.vintage, home=home)
    directory.mkdir(parents=True, exist_ok=True)
    path = directory / "_meta.json"
    if any(directory.iterdir()):
        raise RuntimeError(
            f"raw cache {meta.source}@{meta.vintage} is immutable or incomplete; "
            "write metadata under a new vintage"
        )
    # Exclusive creation closes the check/write race. A failed partial sidecar is not a complete
    # cache (there is no parquet here) and the WRDS pull path can safely discard it.
    with path.open("x", encoding="utf-8") as stream:
        stream.write(json.dumps(asdict(meta), sort_keys=True, indent=2))
    return path





def read_meta(source: str, vintage: str, *, home: str | Path | None = None) -> RawMeta:
    """Read the ``_meta.json`` for a ``(source, vintage)`` pull."""
    path = _raw_path(source, vintage, home=home) / "_meta.json"
    _validate_cache_file(path, label="raw metadata sidecar")
    payload = json.loads(path.read_text(encoding="utf-8"))
    # Sidecars written before schema-aware digests had no explicit version and are v1.
    payload.setdefault("content_digest_version", 1)
    return RawMeta(**payload)