refactor(usage): address round-1 review on PR #871

- store: add optional `enabled` param to record_usage_event so hot-path
  callers (nc_semantic_search) pass the already-resolved flag instead of
  forcing a second uncached Settings build (ADR-024); falls back to
  get_settings() when None so the store stays self-gating for standalone
  use.
- hooks: thread enabled= through both call sites; bump the outer
  shared()/construction failure log from debug → warning so "metering
  enabled but no billing data" is visible at the default INFO level.
- migration: instantiate postgresql.JSONB() to match the sibling
  TIMESTAMP(timezone=True) column.
- tests: fix the misleading "asyncpg returns JSONB as a JSON string"
  comment; add occurred_at dialect round-trip test and an enabled-param
  short-circuit test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2026-06-07 15:20:41 +02:00
co-authored by Claude Opus 4.8
parent 1c6b1a84ea
commit 702f66e6b1
6 changed files with 84 additions and 5 deletions
+1
View File
@@ -0,0 +1 @@
{"sessionId":"fdfa73f5-8734-459a-8f23-67e1eeaa804c","pid":13329,"procStart":"23713","acquiredAt":1780838157352}
@@ -60,7 +60,7 @@ def upgrade() -> None:
# (json.dumps) on SQLite. # (json.dumps) on SQLite.
sa.Column( sa.Column(
"metadata", "metadata",
postgresql.JSONB if is_pg else sa.Text, postgresql.JSONB() if is_pg else sa.Text,
nullable=True, nullable=True,
), ),
) )
+8 -1
View File
@@ -534,9 +534,16 @@ def configure_semantic_tools(mcp: FastMCP):
"fusion": fusion, "fusion": fusion,
"doc_types": doc_types, "doc_types": doc_types,
}, },
# Pass the already-resolved flag so the store doesn't
# rebuild Settings on this hot query path (ADR-024).
enabled=settings.usage_metering_enabled,
) )
except Exception: except Exception:
logger.debug( # Reached only when shared()/store construction itself
# raises (record_usage_event swallows its own write
# failures). Metering is on, so warn — a silent DEBUG line
# would hide "operator enabled metering but gets no data".
logger.warning(
"usage metering hook (embeddings_queries) skipped", "usage metering hook (embeddings_queries) skipped",
exc_info=True, exc_info=True,
) )
+10 -1
View File
@@ -72,6 +72,7 @@ class UsageEventStore:
occurred_at: datetime | None = None, occurred_at: datetime | None = None,
metadata: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None,
event_id: str | None = None, event_id: str | None = None,
enabled: bool | None = None,
) -> None: ) -> None:
"""Record one billable usage event (best-effort, flag-gated). """Record one billable usage event (best-effort, flag-gated).
@@ -86,8 +87,16 @@ class UsageEventStore:
metadata: Optional rawest-unit context (provider, model, tokens, metadata: Optional rawest-unit context (provider, model, tokens,
doc_type, ...). Stored as JSONB (Postgres) / JSON text (SQLite). doc_type, ...). Stored as JSONB (Postgres) / JSON text (SQLite).
event_id: Optional idempotency key; defaults to a fresh UUID4. event_id: Optional idempotency key; defaults to a fresh UUID4.
enabled: The resolved ``USAGE_METERING_ENABLED`` value. ``None``
(default) re-reads it via ``get_settings()`` so the store stays
self-gating for standalone/test use. Hot-path callers that
already hold the flag should pass it to avoid a second uncached
``Settings`` build (``get_settings()`` is non-cached per
ADR-024 and ``nc_semantic_search`` is on the query path).
""" """
if not get_settings().usage_metering_enabled: if enabled is None:
enabled = get_settings().usage_metering_enabled
if not enabled:
return return
start = time.time() start = time.time()
+8 -1
View File
@@ -839,9 +839,16 @@ async def _index_document(
"user_id": doc_task.user_id, "user_id": doc_task.user_id,
"total_chars": total_chars, "total_chars": total_chars,
}, },
# Pass the already-resolved flag so the store doesn't
# rebuild Settings here (ADR-024).
enabled=settings.usage_metering_enabled,
) )
except Exception: except Exception:
logger.debug( # Reached only when shared()/store construction itself
# raises (record_usage_event swallows its own write
# failures). Metering is on, so warn rather than hide the
# "enabled but no billing data" case in DEBUG logs.
logger.warning(
"usage metering hook (pages_chunks) skipped", exc_info=True "usage metering hook (pages_chunks) skipped", exc_info=True
) )
+56 -1
View File
@@ -14,6 +14,7 @@ that a DB failure is swallowed instead of surfacing to the caller.
import json import json
import tempfile import tempfile
import uuid import uuid
from datetime import datetime, timezone
from pathlib import Path from pathlib import Path
import pytest import pytest
@@ -84,6 +85,32 @@ async def test_flag_off_is_noop(storage, monkeypatch):
assert await _count(storage) == 0 assert await _count(storage) == 0
async def test_enabled_param_short_circuits_without_reading_settings(
storage, monkeypatch
):
"""An explicit ``enabled`` flag is honored without touching get_settings().
Hot-path callers pass the already-resolved flag; the store must not rebuild
Settings when given one. ``enabled=False`` is a no-op; ``enabled=True``
writes even though the (boobytrapped) settings lookup would raise.
"""
def _boom():
raise AssertionError("get_settings() must not be called when enabled is passed")
monkeypatch.setattr(store_module, "get_settings", _boom)
store = UsageEventStore(storage)
await store.record_usage_event(metric="pages_chunks", value=1, enabled=False)
assert await _count(storage) == 0
eid = str(uuid.uuid4())
await store.record_usage_event(
metric="pages_chunks", value=1, event_id=eid, enabled=True
)
assert await _count(storage) == 1
async def test_insert_roundtrip(storage, monkeypatch): async def test_insert_roundtrip(storage, monkeypatch):
"""A recorded event lands and reads back with the right fields.""" """A recorded event lands and reads back with the right fields."""
_set_metering(monkeypatch, True) _set_metering(monkeypatch, True)
@@ -127,11 +154,39 @@ async def test_metadata_json_roundtrip(storage, monkeypatch):
) )
row = await _fetch(storage, eid) row = await _fetch(storage, eid)
raw = row[4] raw = row[4]
# asyncpg returns JSONB as a JSON string (no codec); SQLite stores TEXT. # Depending on the asyncpg/SQLAlchemy JSONB codec in play, Postgres may
# return JSONB as a Python dict or as a JSON str; SQLite stores TEXT.
# Handle both so the test is robust across driver/codec versions.
loaded = raw if isinstance(raw, dict) else json.loads(raw) loaded = raw if isinstance(raw, dict) else json.loads(raw)
assert loaded == meta assert loaded == meta
async def test_occurred_at_roundtrip(storage, monkeypatch):
"""occurred_at round-trips to the same instant on both backends.
The store binds a datetime on Postgres and an ISO string on SQLite (the
only dialect-specific branch in the store); this pins that both read back
to the same instant regardless of the stored representation.
"""
_set_metering(monkeypatch, True)
store = UsageEventStore(storage)
eid = str(uuid.uuid4())
when = datetime(2026, 6, 10, 12, 0, 0, tzinfo=timezone.utc)
await store.record_usage_event(
metric="pages_chunks", value=1, event_id=eid, occurred_at=when
)
row = await _fetch(storage, eid)
stored = row[1]
# SQLite returns the ISO string we bound; Postgres returns a datetime.
parsed = stored if isinstance(stored, datetime) else datetime.fromisoformat(stored)
# Aware-datetime equality compares the instant, so a UTC value coming back
# in another session tz still matches; a naive value (none expected) is
# treated as UTC.
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=timezone.utc)
assert parsed == when
async def test_metadata_none_is_null(storage, monkeypatch): async def test_metadata_none_is_null(storage, monkeypatch):
"""Omitting metadata stores SQL NULL, not the string 'null'.""" """Omitting metadata stores SQL NULL, not the string 'null'."""
_set_metering(monkeypatch, True) _set_metering(monkeypatch, True)