pysmo.tools.archive

@runtime_checkable
class RawFetcher(Protocol):
    """Callable `(*, station, starttime, endtime) -> bytes` returning a raw, unparsed fetch response.

    A `Protocol` with a keyword-only `__call__`, not a plain `Callable[...]`
    type alias, specifically because the functions this slot is meant to be
    filled with directly — [`fetch_sac`][pysmo.tools.web.fetch_sac] and
    [`fetch_geocsvseismogram`][pysmo.tools.web.fetch_geocsvseismogram] — are
    themselves keyword-only. A plain positional `Callable` type cannot
    express that, and calling one positionally raises `TypeError` regardless
    of what a type checker allows.
    """

    def __call__(
        self, *, station: Station, starttime: pd.Timestamp, endtime: pd.Timestamp
    ) -> bytes:
        """Fetch raw bytes for a station and absolute time window."""
        ...

@define(kw_only=True)
class SqliteArchiveFetcher:
    """Caches raw fetch responses in one local SQLite database.

    Format-agnostic: stores whatever bytes `fetch_raw` returns
    (zlib-compressed), keyed by station identity and time window, and hands
    the decompressed bytes to `parse` on both a cache hit and a miss. A hit
    never calls `fetch_raw` again — a side effect of this is bit-identical
    replay of a previously fetched window, rather than only detecting drift
    after the fact.

    Warning: Local disk only
        SQLite's own documentation states that WAL mode does not work over a
        network filesystem, and recommends against concurrent multi-process
        access to a SQLite database over NFS at all. This class assumes the
        database file lives on local disk with correctly functioning file
        locking; it is not a safe choice for a cache shared over a network
        filesystem by more than one process at a time.
    """

    path: Path = field(converter=Path)
    """Location of the SQLite database file.

    The file itself is created on first use if it doesn't exist; its
    *parent directory* must already exist, checked at construction time.
    """

    fetch_raw: RawFetcher
    """Fetches a raw response for a station and absolute time window."""

    parse: RawParser
    """Parses a raw response (freshly fetched, or read back from cache) into a `Seismogram`."""

    wal: bool = False
    """Enable WAL mode. Only for a database confirmed to be on local disk — see the class docstring."""

    _conn: sqlite3.Connection | None = field(
        init=False, default=None, repr=False, eq=False
    )
    _lock: threading.Lock = field(
        init=False, factory=threading.Lock, repr=False, eq=False
    )
    """Guards the check-then-set on `_conn`: `check_same_thread=False` means
    this instance may legitimately be called from more than one thread, and
    without this lock two threads racing the first call could each open
    their own connection, silently leaking one."""

    def __attrs_post_init__(self) -> None:
        """Fail fast if `path`'s parent directory doesn't exist."""
        if not self.path.parent.is_dir():
            raise FileNotFoundError(
                f"Parent directory does not exist: {self.path.parent}"
            )

    def __getstate__(self) -> dict:
        """Drop the live connection; the lock is excluded entirely below (no `threading.Lock` is picklable, not even a fresh one)."""
        state = attrs_getstate(self, {"_conn": None})
        del state["_lock"]
        return state

    def __setstate__(self, state: dict) -> None:
        """Restore state without triggering any `on_setattr` hooks, then create a fresh lock."""
        attrs_setstate(self, state)
        object.__setattr__(self, "_lock", threading.Lock())

    def close(self) -> None:
        """Close the underlying connection, if one is open.

        Not required before the object is garbage-collected or the process
        exits — normal teardown closes the file descriptor regardless — but
        call it explicitly to release the connection sooner in a
        long-running process holding many such fetchers.
        """
        with self._lock:
            if self._conn is not None:
                self._conn.close()
                self._conn = None

    def _connect(self) -> sqlite3.Connection:
        with self._lock:
            if self._conn is None:
                conn = sqlite3.connect(self.path, timeout=30, check_same_thread=False)
                if self.wal:
                    conn.execute("PRAGMA journal_mode=WAL")
                conn.execute(
                    "CREATE TABLE IF NOT EXISTS cache "
                    "(key TEXT PRIMARY KEY, data BLOB NOT NULL)"
                )
                version = conn.execute("PRAGMA user_version").fetchone()[0]
                if version == 0:
                    conn.execute(f"PRAGMA user_version = {_ENCODING_VERSION}")
                elif version != _ENCODING_VERSION:
                    conn.close()
                    raise ValueError(
                        f"{self.path} was written with a different cache "
                        f"encoding (user_version={version}, expected "
                        f"{_ENCODING_VERSION})."
                    )
                self._conn = conn
            return self._conn

    def __call__(
        self, station: Station, starttime: pd.Timestamp, endtime: pd.Timestamp
    ) -> Seismogram:
        """Return a `Seismogram` for `station` and window, from cache if already fetched.

        Args:
            station: Station to fetch data for.
            starttime: Start of the requested window (UTC).
            endtime: End of the requested window (UTC).

        Returns:
            Parsed result — from the cache database on a hit, freshly
            fetched (and then stored) on a miss.
        """
        key = self._key(station, starttime, endtime)
        conn = self._connect()
        row = conn.execute("SELECT data FROM cache WHERE key = ?", (key,)).fetchone()
        if row is not None:
            return self.parse(zlib.decompress(row[0]))
        raw = self.fetch_raw(station=station, starttime=starttime, endtime=endtime)
        with conn:
            conn.execute(
                "INSERT OR IGNORE INTO cache (key, data) VALUES (?, ?)",
                (key, zlib.compress(raw)),
            )
        return self.parse(raw)

    @staticmethod
    def _key(station: Station, starttime: pd.Timestamp, endtime: pd.Timestamp) -> str:
        return (
            f"{station.network}.{station.name}.{station.location}."
            f"{station.channel}_{starttime.isoformat()}_{endtime.isoformat()}"
        )