feat(usage): record per-tenant usage events into the app DB
Deck #67 data-plane slice: tenant Pods record billable operations (embedding queries, pages/chunks embedded) into an app-DB usage_events table that the control plane later pulls read-only into the billing ledger and syncs to Stripe Meter Events. - migration 007: usage_events table (Postgres TIMESTAMPTZ/JSONB/UUID with portable SQLite fallbacks), indexed (occurred_at, metric) for the CP rollup's per-day range scan + GROUP BY metric. - UsageEventStore: best-effort, flag-gated writer reusing the shared RefreshTokenStorage engine; ON CONFLICT (event_id) DO NOTHING for idempotent retries; dialect-branched occurred_at bind. All work (incl. metadata JSON encode) is swallowed so a metering failure never surfaces to the user op. - USAGE_METERING_ENABLED flag (default off) wired through Settings + env map; off-path touches no storage, so OSS self-hosters get an empty table and zero write overhead. - two recording hooks: embeddings_queries (per nc_semantic_search, which nc_semantic_search_answer reuses) and pages_chunks (after dense embedding succeeds, covering both in-process and procrastinate paths). - storage.acquire()/.dialect public seams so the sibling store doesn't reach into the underscored internal. - tests parametrized over SQLite + Postgres: flag-off no-op, roundtrip, ON CONFLICT dedup, JSON/NULL metadata, and the best-effort swallow of both DB errors and unserializable metadata. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
42505f8f87
commit
1c6b1a84ea
@@ -157,6 +157,18 @@ class TestGetSettings:
|
||||
assert settings.vector_sync_processor_workers == 5
|
||||
assert settings.vector_sync_queue_max_size == 5000
|
||||
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
def test_usage_metering_disabled_by_default(self):
|
||||
"""USAGE_METERING_ENABLED defaults to False (OSS doesn't self-monitor)."""
|
||||
_reload_config()
|
||||
assert get_settings().usage_metering_enabled is False
|
||||
|
||||
@patch.dict(os.environ, {"USAGE_METERING_ENABLED": "true"}, clear=True)
|
||||
def test_usage_metering_enabled_via_env(self):
|
||||
"""USAGE_METERING_ENABLED=true maps to settings.usage_metering_enabled."""
|
||||
_reload_config()
|
||||
assert get_settings().usage_metering_enabled is True
|
||||
|
||||
|
||||
class TestChunkConfigValidation:
|
||||
"""Test document chunking configuration validation."""
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
"""Unit tests for ``UsageEventStore`` (Deck #67 usage metering, data plane).
|
||||
|
||||
Parametrized over both supported backends via the shared ``storage_backend``
|
||||
fixture: SQLite (default, always runs) and Postgres (opt-in, gated on
|
||||
``TEST_DATABASE_URL`` — bring up ``docker compose --profile postgres up -d
|
||||
postgres-test`` and export
|
||||
``TEST_DATABASE_URL=postgresql+asyncpg://mcp:mcp@localhost:5433/mcp``).
|
||||
|
||||
Covers the recording contract: flag-gated no-op, insert roundtrip, ON CONFLICT
|
||||
dedup, JSON metadata roundtrip, NULL metadata, and the best-effort guarantee
|
||||
that a DB failure is swallowed instead of surfacing to the caller.
|
||||
"""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
import nextcloud_mcp_server.usage.store as store_module
|
||||
from nextcloud_mcp_server.auth.storage import RefreshTokenStorage
|
||||
from nextcloud_mcp_server.usage.store import UsageEventStore
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def storage(storage_backend):
|
||||
"""Initialized RefreshTokenStorage backed by SQLite or Postgres."""
|
||||
key = Fernet.generate_key()
|
||||
if storage_backend["kind"] == "sqlite":
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
db_path = Path(tmpdir) / "usage.db"
|
||||
s = RefreshTokenStorage(db_path=str(db_path), encryption_key=key)
|
||||
await s.initialize()
|
||||
yield s
|
||||
else:
|
||||
s = RefreshTokenStorage(database_url=storage_backend["url"], encryption_key=key)
|
||||
await s.initialize()
|
||||
try:
|
||||
yield s
|
||||
finally:
|
||||
await storage_backend["reset"]()
|
||||
|
||||
|
||||
def _set_metering(monkeypatch, enabled: bool) -> None:
|
||||
"""Force the metering flag without mutating global dynaconf state.
|
||||
|
||||
``record_usage_event`` calls ``get_settings()`` (imported into the store
|
||||
module's namespace), and ``get_settings()`` builds a fresh Settings per
|
||||
call — so patching the symbol the store sees is the clean seam.
|
||||
"""
|
||||
|
||||
class _Settings:
|
||||
usage_metering_enabled = enabled
|
||||
|
||||
monkeypatch.setattr(store_module, "get_settings", lambda: _Settings())
|
||||
|
||||
|
||||
async def _count(storage: RefreshTokenStorage) -> int:
|
||||
async with storage.acquire() as db:
|
||||
cursor = await db.execute("SELECT COUNT(*) FROM usage_events")
|
||||
row = await cursor.fetchone()
|
||||
return row[0]
|
||||
|
||||
|
||||
async def _fetch(storage: RefreshTokenStorage, event_id: str):
|
||||
async with storage.acquire() as db:
|
||||
cursor = await db.execute(
|
||||
"SELECT event_id, occurred_at, metric, value, metadata "
|
||||
"FROM usage_events WHERE event_id = ?",
|
||||
(event_id,),
|
||||
)
|
||||
return await cursor.fetchone()
|
||||
|
||||
|
||||
async def test_flag_off_is_noop(storage, monkeypatch):
|
||||
"""With metering disabled, nothing is written (zero DB work)."""
|
||||
_set_metering(monkeypatch, False)
|
||||
store = UsageEventStore(storage)
|
||||
await store.record_usage_event(metric="pages_chunks", value=5)
|
||||
assert await _count(storage) == 0
|
||||
|
||||
|
||||
async def test_insert_roundtrip(storage, monkeypatch):
|
||||
"""A recorded event lands and reads back with the right fields."""
|
||||
_set_metering(monkeypatch, True)
|
||||
store = UsageEventStore(storage)
|
||||
eid = str(uuid.uuid4())
|
||||
await store.record_usage_event(
|
||||
metric="pages_chunks",
|
||||
value=7,
|
||||
event_id=eid,
|
||||
metadata={"provider": "gateway"},
|
||||
)
|
||||
row = await _fetch(storage, eid)
|
||||
assert row is not None
|
||||
# Postgres returns event_id as a uuid.UUID; normalize to str for compare.
|
||||
assert str(row[0]) == eid
|
||||
assert row[2] == "pages_chunks"
|
||||
assert row[3] == 7
|
||||
|
||||
|
||||
async def test_on_conflict_dedup(storage, monkeypatch):
|
||||
"""A duplicate event_id is a no-op; the first write is retained."""
|
||||
_set_metering(monkeypatch, True)
|
||||
store = UsageEventStore(storage)
|
||||
eid = str(uuid.uuid4())
|
||||
await store.record_usage_event(metric="pages_chunks", value=1, event_id=eid)
|
||||
await store.record_usage_event(metric="embeddings_queries", value=99, event_id=eid)
|
||||
assert await _count(storage) == 1
|
||||
row = await _fetch(storage, eid)
|
||||
assert row[2] == "pages_chunks" # DO NOTHING, not DO UPDATE
|
||||
assert row[3] == 1
|
||||
|
||||
|
||||
async def test_metadata_json_roundtrip(storage, monkeypatch):
|
||||
"""Nested metadata round-trips as JSON on both backends."""
|
||||
_set_metering(monkeypatch, True)
|
||||
store = UsageEventStore(storage)
|
||||
eid = str(uuid.uuid4())
|
||||
meta = {"provider": "gateway", "model": "titan", "nested": {"chunks": 3}}
|
||||
await store.record_usage_event(
|
||||
metric="pages_chunks", value=3, event_id=eid, metadata=meta
|
||||
)
|
||||
row = await _fetch(storage, eid)
|
||||
raw = row[4]
|
||||
# asyncpg returns JSONB as a JSON string (no codec); SQLite stores TEXT.
|
||||
loaded = raw if isinstance(raw, dict) else json.loads(raw)
|
||||
assert loaded == meta
|
||||
|
||||
|
||||
async def test_metadata_none_is_null(storage, monkeypatch):
|
||||
"""Omitting metadata stores SQL NULL, not the string 'null'."""
|
||||
_set_metering(monkeypatch, True)
|
||||
store = UsageEventStore(storage)
|
||||
eid = str(uuid.uuid4())
|
||||
await store.record_usage_event(
|
||||
metric="embeddings_queries", value=1, event_id=eid, metadata=None
|
||||
)
|
||||
row = await _fetch(storage, eid)
|
||||
assert row[4] is None
|
||||
|
||||
|
||||
async def test_best_effort_swallows_db_errors(storage, monkeypatch):
|
||||
"""A DB failure is logged + dropped, never raised into the caller."""
|
||||
_set_metering(monkeypatch, True)
|
||||
store = UsageEventStore(storage)
|
||||
|
||||
recorded: list[tuple] = []
|
||||
monkeypatch.setattr(
|
||||
store_module,
|
||||
"record_db_operation",
|
||||
lambda *args, **kwargs: recorded.append(args),
|
||||
)
|
||||
|
||||
def _boom():
|
||||
raise RuntimeError("db down")
|
||||
|
||||
# ``acquire`` raises on call — the store must catch and continue.
|
||||
monkeypatch.setattr(storage, "acquire", _boom)
|
||||
|
||||
# Must not raise.
|
||||
await store.record_usage_event(metric="pages_chunks", value=1)
|
||||
|
||||
assert recorded, "record_db_operation should be called on the error path"
|
||||
assert recorded[-1][3] == "error"
|
||||
|
||||
|
||||
async def test_best_effort_swallows_unserializable_metadata(storage, monkeypatch):
|
||||
"""Non-serializable metadata is swallowed, not raised into the caller.
|
||||
|
||||
json.dumps runs inside the best-effort try, so a metadata value the JSON
|
||||
encoder can't handle must drop the event like any other write failure
|
||||
rather than surfacing to the user op.
|
||||
"""
|
||||
_set_metering(monkeypatch, True)
|
||||
store = UsageEventStore(storage)
|
||||
|
||||
# An arbitrary object is not JSON-serializable; json.dumps raises TypeError.
|
||||
bad_metadata = {"obj": object()}
|
||||
|
||||
# Must not raise.
|
||||
await store.record_usage_event(
|
||||
metric="pages_chunks", value=1, metadata=bad_metadata
|
||||
)
|
||||
|
||||
# Nothing was written — the encode failed before the insert.
|
||||
assert await _count(storage) == 0
|
||||
Reference in New Issue
Block a user