server-26/drb-c2-core/app/internal/firestore.py

import asyncio
import time as _time
from contextvars import ContextVar
from typing import Optional, Any
import firebase_admin
from firebase_admin import credentials, firestore as fs
from google.cloud.firestore_v1.base_query import FieldFilter
from app.config import settings
from app.internal.logger import logger

# Re-exported so callers never need their own `firebase_admin.firestore` import
# just to delete a field. server-26#96/#114 review: `doc_set(..., merge=True)`
# merges nested maps by key but can never REMOVE one — writing `{"scenes": {}}`
# to clear a map is a no-op, not a delete. Use `doc_update(coll, id, {"field":
# fstore.DELETE_FIELD})` (or doc_set + merge, DELETE_FIELD works under both)
# whenever a re-extraction/reprocess path needs a stale nested field gone
# rather than merged over.
DELETE_FIELD = fs.DELETE_FIELD

# ---------------------------------------------------------------------------
# In-memory TTL cache for rarely-changing documents (systems, nodes config)
# ---------------------------------------------------------------------------
# Key: "collection/doc_id"  →  (expires_at_monotonic, data_or_None)
_doc_cache: dict[str, tuple[float, Optional[dict]]] = {}


def _init_firebase():
    if firebase_admin._apps:
        return firestore.client()

    if settings.gcp_credentials_path:
        cred = credentials.Certificate(settings.gcp_credentials_path)
    else:
        cred = credentials.ApplicationDefault()

    firebase_admin.initialize_app(cred)
    logger.info("Firebase initialised.")


_init_firebase()
db = fs.client(database_id=settings.firestore_database)


# ---------------------------------------------------------------------------
# Replay sandbox (app/internal/replay.py)
# ---------------------------------------------------------------------------
# While a replay run is executing, every read and write the pipeline makes to
# `calls` or `incidents` is redirected to that run's own subcollections under
# replay_runs/{run_id}/, so re-running the pipeline over past traffic can never
# touch a live call or incident. A subcollection keeps the same collection ID
# ("calls"/"incidents"), so the composite indexes prod queries depend on apply
# to it unchanged. Everything else (systems, nodes, config) is read from prod
# as-is. ContextVar for the same reason as app/internal/clock.py: the redirect
# follows the replay task and never a concurrent live upload.
SANDBOXED_COLLECTIONS = frozenset({"calls", "incidents"})
_sandbox_root: ContextVar[Optional[str]] = ContextVar("drb_fstore_sandbox", default=None)


def enter_sandbox(root: Optional[str]):
    """Redirect calls/incidents under `root` (e.g. "replay_runs/<id>") for this context."""
    return _sandbox_root.set(root)


def exit_sandbox(token) -> None:
    _sandbox_root.reset(token)


def in_sandbox() -> bool:
    """True inside a replay run. Anything that writes live state OTHER than
    calls/incidents (AI health alerts, pending-term queues) checks this and
    stands down — the redirect below only covers the two sandboxed collections."""
    return _sandbox_root.get() is not None


def _path(collection: str) -> str:
    root = _sandbox_root.get()
    if root and collection in SANDBOXED_COLLECTIONS:
        return f"{root}/{collection}"
    return collection


# ---------------------------------------------------------------------------
# Thin async wrappers — firebase-admin is synchronous, run in thread executor
# ---------------------------------------------------------------------------

async def doc_set(collection: str, doc_id: str, data: dict, merge: bool = True) -> None:
    ref = db.collection(_path(collection)).document(doc_id)
    await asyncio.to_thread(ref.set, data, merge=merge)


async def doc_get(collection: str, doc_id: str) -> Optional[dict]:
    ref = db.collection(_path(collection)).document(doc_id)
    snap = await asyncio.to_thread(ref.get)
    return snap.to_dict() if snap.exists else None


async def doc_update(collection: str, doc_id: str, data: dict) -> None:
    ref = db.collection(_path(collection)).document(doc_id)
    await asyncio.to_thread(ref.update, data)


async def collection_list(collection: str, **filters) -> list[dict]:
    """
    List all documents in a collection.
    Optional keyword filters: field=value pairs passed as equality where-clauses.
    """
    def _query():
        ref = db.collection(_path(collection))
        for field, value in filters.items():
            ref = ref.where(filter=FieldFilter(field, "==", value))
        return [doc.to_dict() for doc in ref.stream()]

    return await asyncio.to_thread(_query)


async def collection_where(
    collection: str,
    conditions: list[tuple[str, str, Any]],
    order_by: Optional[list[tuple[str, str]]] = None,
    limit_to: Optional[int] = None,
    start_after: Optional[dict] = None,
) -> list[dict]:
    """
    Query a collection with arbitrary where-clauses.
    conditions: list of (field, op, value) — e.g. [("ended_at", ">=", cutoff_dt)]
    Supports any Firestore operator, including "array_contains" — it's just
    forwarded straight to FieldFilter, so a condition like
    ("incident_ids", "array_contains", incident_id) already worked before this
    function grew explicit order_by/limit/cursor params below.

    order_by: list of (field, direction) — direction is "ASCENDING" or
      "DESCENDING" (Firestore's own constants; passed straight through as
      strings so this module doesn't need a google.cloud.firestore_v1.Query
      import). Applied in list order, so multi-field sorts work.
    limit_to: cap the number of documents returned.
    start_after: cursor — a dict of the same field values as the *last*
      document from a previous page's order_by fields (Firestore's
      `Query.start_after()` takes a field-value mapping, not a document
      snapshot, when you're not holding one).

    Added for org_id-scoped queries that also need to be ordered/paginated —
    unscoped equality-only lookups can keep using collection_list().
    """
    def _query():
        ref = db.collection(_path(collection))
        for field, op, value in conditions:
            ref = ref.where(filter=FieldFilter(field, op, value))
        for field, direction in (order_by or []):
            ref = ref.order_by(field, direction=direction)
        if start_after is not None:
            ref = ref.start_after(start_after)
        if limit_to is not None:
            ref = ref.limit(limit_to)
        return [doc.to_dict() for doc in ref.stream()]

    return await asyncio.to_thread(_query)


async def doc_delete(collection: str, doc_id: str) -> None:
    ref = db.collection(_path(collection)).document(doc_id)
    await asyncio.to_thread(ref.delete)


async def doc_get_cached(collection: str, doc_id: str, ttl: float = 300.0) -> Optional[dict]:
    """
    Like doc_get but backed by a short-lived in-memory TTL cache.
    Use for documents that change rarely (systems config, node assignments).
    Default TTL is 5 minutes — a write will be visible within that window.
    """
    key = f"{_path(collection)}/{doc_id}"
    now = _time.monotonic()
    entry = _doc_cache.get(key)
    if entry and now < entry[0]:
        return entry[1]
    data = await doc_get(collection, doc_id)
    _doc_cache[key] = (now + ttl, data)
    return data