fithia2/libs/intraday/cache.py

"""Disk-based Parquet cache for intraday bar data.

Layout: {cache_dir}/{TICKER}/{YYYY-MM-DD}.parquet

Each file stores one trading day of 5-minute bars for one ticker.
~78 bars × 7 columns ≈ 2-3 KB per file (snappy compressed).
"""
from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq


_CACHE_METADATA = {
    b"intraday_cache_version": b"3",    # v3: SIP data (v2 was IEX ~2.5% tape)
    b"intraday_cache_source": b"api_v1_alpaca_intraday",
    b"intraday_cache_interval": b"5min",
}

_SCHEMA = pa.schema([
    pa.field("timestamp", pa.string()),
    pa.field("open", pa.float64()),
    pa.field("high", pa.float64()),
    pa.field("low", pa.float64()),
    pa.field("close", pa.float64()),
    pa.field("volume", pa.float64()),
    pa.field("vwap", pa.float64()),
])
_REQUIRED_COLUMNS = {"timestamp", "open", "high", "low", "close", "volume"}
_SCHEMA_WITH_METADATA = _SCHEMA.with_metadata(_CACHE_METADATA)

class IntradayCache:
    """Disk-based cache for 5-minute intraday bars using Parquet.

    Thread-safe for reads; uses atomic write (tmp → rename) for writes.
    """

    def __init__(self, cache_dir: str = "data/cache/intraday") -> None:
        self._root = Path(cache_dir)

    def _path(self, ticker: str, date: str) -> Path:
        return self._root / ticker.upper() / f"{date}.parquet"

    @staticmethod
    def _metadata_valid(path: Path) -> bool:
        try:
            meta = pq.read_metadata(str(path))
            if meta.num_rows <= 0:
                return False
            arrow_schema = meta.schema.to_arrow_schema()
            if not _REQUIRED_COLUMNS.issubset(set(arrow_schema.names)):
                return False
            schema_meta = arrow_schema.metadata or {}
            return all(schema_meta.get(k) == v for k, v in _CACHE_METADATA.items())
        except Exception:
            return False

    def has(self, ticker: str, date: str) -> bool:
        """Return True if cached bars exist for ticker on date."""
        p = self._path(ticker, date)
        if not p.exists():
            return False
        if not self._metadata_valid(p):
            p.unlink(missing_ok=True)
            return False
        return True

    def get(self, ticker: str, date: str) -> list[dict[str, Any]] | None:
        """Read cached bars. Returns None on cache miss or read error."""
        p = self._path(ticker, date)
        if not p.exists():
            return None
        if not self._metadata_valid(p):
            p.unlink(missing_ok=True)
            return None
        try:
            table = pq.read_table(str(p))
            if table.num_rows == 0:
                return None
            return table.to_pylist()
        except Exception:
            return None

    def put(self, ticker: str, date: str, bars: list[dict[str, Any]]) -> None:
        """Write bars to cache using atomic write (tmp → rename).

        Silently skips if bars is empty.
        """
        if not bars:
            return

        p = self._path(ticker, date)
        p.parent.mkdir(parents=True, exist_ok=True)

        # Normalize bars to match schema
        rows = []
        for b in bars:
            rows.append({
                "timestamp": str(b.get("timestamp", "")),
                "open": float(b.get("open", 0.0)),
                "high": float(b.get("high", 0.0)),
                "low": float(b.get("low", 0.0)),
                "close": float(b.get("close", 0.0)),
                "volume": float(b.get("volume", 0.0)),
                "vwap": float(b.get("vwap", 0.0)) if b.get("vwap") is not None else 0.0,
            })

        table = pa.Table.from_pylist(rows, schema=_SCHEMA_WITH_METADATA)
        tmp = p.with_suffix(".tmp")
        try:
            pq.write_table(table, str(tmp), compression="snappy")
            os.replace(str(tmp), str(p))  # atomic on POSIX
        except Exception:
            if tmp.exists():
                tmp.unlink(missing_ok=True)
            raise

    def evict(
        self,
        ticker: str | None = None,
        before_date: str | None = None,
    ) -> int:
        """Remove cached files. Returns count of files removed.

        Args:
            ticker: If set, only evict this ticker's files.
            before_date: If set, only evict files for dates < this (YYYY-MM-DD).
        """
        removed = 0
        if ticker:
            ticker_dir = self._root / ticker.upper()
            if not ticker_dir.exists():
                return 0
            dirs_to_scan = [ticker_dir]
        else:
            if not self._root.exists():
                return 0
            dirs_to_scan = [d for d in self._root.iterdir() if d.is_dir()]

        for d in dirs_to_scan:
            for f in d.glob("*.parquet"):
                if before_date and f.stem >= before_date:
                    continue
                f.unlink(missing_ok=True)
                removed += 1

        return removed

    def stats(self) -> dict[str, Any]:
        """Return cache statistics."""
        if not self._root.exists():
            return {"total_files": 0, "total_bytes": 0, "tickers": 0}

        total_files = 0
        total_bytes = 0
        tickers = set()
        dates: list[str] = []

        for ticker_dir in self._root.iterdir():
            if not ticker_dir.is_dir():
                continue
            tickers.add(ticker_dir.name)
            for f in ticker_dir.glob("*.parquet"):
                total_files += 1
                total_bytes += f.stat().st_size
                dates.append(f.stem)

        return {
            "total_files": total_files,
            "total_bytes": total_bytes,
            "total_mb": round(total_bytes / 1_048_576, 2),
            "tickers": len(tickers),
            "date_min": min(dates) if dates else None,
            "date_max": max(dates) if dates else None,
        }