diff --git a/tests/test_usage_metering.py b/tests/test_usage_metering.py new file mode 100644 index 0000000..543d3ba --- /dev/null +++ b/tests/test_usage_metering.py @@ -0,0 +1,222 @@ +"""Tests for usage_metering: durable per-key usage counters and monthly quotas. + +All tests use a temporary SQLite file and inject `now` explicitly, so they are +deterministic and touch no network. +""" + +import tempfile +import threading +import unittest +from datetime import datetime +from pathlib import Path + +from usage_metering import QuotaExceededError, UsageStore + +JAN = datetime(2026, 1, 15, 12, 0, 0) +JAN_LATER = datetime(2026, 1, 28, 23, 59, 59) +FEB = datetime(2026, 2, 1, 0, 0, 0) + + +class UsageStoreTestCase(unittest.TestCase): + """Base fixture: a fresh UsageStore on a tmp sqlite file per test.""" + + def setUp(self): + self._tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(self._tmpdir.cleanup) + self.db_path = Path(self._tmpdir.name) / "usage.db" + self.store = UsageStore(self.db_path) + + +class TestRecordAndUsage(UsageStoreTestCase): + """record() accumulates within a period; periods are independent.""" + + def test_record_accumulates_within_period(self): + self.store.record("key-a", input_bytes=100, output_bytes=40, now=JAN) + self.store.record("key-a", input_bytes=200, output_bytes=60, now=JAN_LATER) + got = self.store.usage("key-a", JAN) + self.assertEqual( + got, + { + "period": "2026-01", + "conversions": 2, + "input_bytes": 300, + "output_bytes": 100, + }, + ) + + def test_new_period_starts_fresh(self): + self.store.record("key-a", input_bytes=100, output_bytes=40, now=JAN) + got = self.store.usage("key-a", FEB) + self.assertEqual( + got, + { + "period": "2026-02", + "conversions": 0, + "input_bytes": 0, + "output_bytes": 0, + }, + ) + # January's totals are untouched. + self.assertEqual(self.store.usage("key-a", JAN)["conversions"], 1) + + def test_unknown_key_reports_zero_usage(self): + got = self.store.usage("never-seen", JAN) + self.assertEqual( + got, + { + "period": "2026-01", + "conversions": 0, + "input_bytes": 0, + "output_bytes": 0, + }, + ) + + def test_keys_are_isolated(self): + self.store.record("key-a", input_bytes=10, output_bytes=5, now=JAN) + self.store.record("key-b", input_bytes=99, output_bytes=1, now=JAN) + self.assertEqual(self.store.usage("key-a", JAN)["input_bytes"], 10) + self.assertEqual(self.store.usage("key-b", JAN)["input_bytes"], 99) + + def test_negative_bytes_rejected(self): + with self.assertRaises(ValueError): + self.store.record("key-a", input_bytes=-1, output_bytes=0, now=JAN) + with self.assertRaises(ValueError): + self.store.record("key-a", input_bytes=0, output_bytes=-1, now=JAN) + + def test_usage_survives_reopen(self): + self.store.record("key-a", input_bytes=7, output_bytes=3, now=JAN) + reopened = UsageStore(self.db_path) + self.assertEqual(reopened.usage("key-a", JAN)["conversions"], 1) + + +class TestQuota(UsageStoreTestCase): + """check_quota() allows under the limit and denies at the boundary.""" + + def test_allowed_under_limit(self): + self.store.record("key-a", input_bytes=10, output_bytes=10, now=JAN) + self.assertTrue( + self.store.check_quota( + "key-a", JAN, max_conversions=2, max_bytes=100 + ) + ) + + def test_unlimited_when_no_limits_given(self): + for _ in range(50): + self.store.record("key-a", input_bytes=1, output_bytes=1, now=JAN) + self.assertTrue(self.store.check_quota("key-a", JAN)) + + def test_denied_at_conversion_boundary(self): + self.store.record("key-a", input_bytes=1, output_bytes=1, now=JAN) + self.store.record("key-a", input_bytes=1, output_bytes=1, now=JAN) + with self.assertRaises(QuotaExceededError) as ctx: + self.store.check_quota("key-a", JAN, max_conversions=2) + self.assertEqual(ctx.exception.limit_name, "max_conversions") + self.assertEqual(ctx.exception.limit, 2) + self.assertEqual(ctx.exception.used, 2) + self.assertEqual(ctx.exception.api_key, "key-a") + + def test_allowed_one_below_conversion_boundary(self): + self.store.record("key-a", input_bytes=1, output_bytes=1, now=JAN) + self.assertTrue(self.store.check_quota("key-a", JAN, max_conversions=2)) + + def test_denied_at_byte_boundary(self): + # 60 + 40 = 100 total bytes, exactly at the limit -> denied. + self.store.record("key-a", input_bytes=60, output_bytes=40, now=JAN) + with self.assertRaises(QuotaExceededError) as ctx: + self.store.check_quota("key-a", JAN, max_bytes=100) + self.assertEqual(ctx.exception.limit_name, "max_bytes") + self.assertEqual(ctx.exception.used, 100) + + def test_allowed_just_below_byte_boundary(self): + self.store.record("key-a", input_bytes=60, output_bytes=39, now=JAN) + self.assertTrue(self.store.check_quota("key-a", JAN, max_bytes=100)) + + def test_error_message_names_the_limit(self): + self.store.record("key-a", input_bytes=1, output_bytes=1, now=JAN) + with self.assertRaises(QuotaExceededError) as ctx: + self.store.check_quota("key-a", JAN, max_conversions=1) + self.assertIn("max_conversions", str(ctx.exception)) + self.assertIn("key-a", str(ctx.exception)) + + def test_quota_resets_in_new_period(self): + self.store.record("key-a", input_bytes=1, output_bytes=1, now=JAN) + with self.assertRaises(QuotaExceededError): + self.store.check_quota("key-a", JAN, max_conversions=1) + self.assertTrue(self.store.check_quota("key-a", FEB, max_conversions=1)) + + def test_unknown_key_always_allowed(self): + self.assertTrue( + self.store.check_quota("never-seen", JAN, max_conversions=1, max_bytes=1) + ) + + +class TestAdminOperations(UsageStoreTestCase): + """reset() and all_usage() admin helpers.""" + + def test_reset_clears_all_periods_for_key(self): + self.store.record("key-a", input_bytes=1, output_bytes=1, now=JAN) + self.store.record("key-a", input_bytes=1, output_bytes=1, now=FEB) + self.store.record("key-b", input_bytes=5, output_bytes=5, now=JAN) + self.store.reset("key-a") + self.assertEqual(self.store.usage("key-a", JAN)["conversions"], 0) + self.assertEqual(self.store.usage("key-a", FEB)["conversions"], 0) + # Other keys are unaffected. + self.assertEqual(self.store.usage("key-b", JAN)["conversions"], 1) + + def test_all_usage_lists_only_that_period(self): + self.store.record("key-a", input_bytes=10, output_bytes=2, now=JAN) + self.store.record("key-b", input_bytes=20, output_bytes=4, now=JAN) + self.store.record("key-c", input_bytes=30, output_bytes=6, now=FEB) + got = self.store.all_usage("2026-01") + self.assertEqual( + got, + { + "key-a": {"conversions": 1, "input_bytes": 10, "output_bytes": 2}, + "key-b": {"conversions": 1, "input_bytes": 20, "output_bytes": 4}, + }, + ) + + def test_all_usage_empty_period(self): + self.assertEqual(self.store.all_usage("1999-12"), {}) + + +class TestConcurrency(UsageStoreTestCase): + """Concurrent record() calls from many threads must sum correctly.""" + + def test_concurrent_records_sum_correctly(self): + threads_n, per_thread = 8, 25 + + def worker(): + for _ in range(per_thread): + self.store.record("key-a", input_bytes=3, output_bytes=2, now=JAN) + + threads = [threading.Thread(target=worker) for _ in range(threads_n)] + for t in threads: + t.start() + for t in threads: + t.join() + + got = self.store.usage("key-a", JAN) + total = threads_n * per_thread + self.assertEqual(got["conversions"], total) + self.assertEqual(got["input_bytes"], 3 * total) + self.assertEqual(got["output_bytes"], 2 * total) + + +class TestMalformedDbPath(unittest.TestCase): + """Constructor errors are clear for bad paths.""" + + def test_missing_parent_directory(self): + with self.assertRaises(ValueError) as ctx: + UsageStore("/nonexistent-dir-xyz/deeper/usage.db") + self.assertIn("parent directory does not exist", str(ctx.exception)) + + def test_path_is_a_directory(self): + with tempfile.TemporaryDirectory() as tmpdir: + with self.assertRaises(ValueError) as ctx: + UsageStore(tmpdir) + self.assertIn("directory", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/usage_metering.py b/usage_metering.py new file mode 100644 index 0000000..d0860be --- /dev/null +++ b/usage_metering.py @@ -0,0 +1,272 @@ +"""Per-API-key usage metering and monthly quota enforcement. + +This module is the missing piece between API-key authentication (*who* is +calling) and billing (*how much* they owe): it durably records per-key usage +counters (conversion count, input bytes, output bytes) in a SQLite database +and enforces monthly quotas against those counters. + +Design notes: + +* Storage is stdlib :mod:`sqlite3` in WAL mode so readers never block the + writer and counters survive process restarts. +* Usage is bucketed by *period*, a ``"YYYY-MM"`` string derived from a + caller-supplied :class:`datetime.datetime`. Callers pass ``now`` explicitly + (the module never calls ``datetime.now()`` itself) so tests are fully + deterministic. +* Every operation opens a short-lived connection, which combined with a + process-level lock makes :class:`UsageStore` safe to share across threads. + +Example:: + + store = UsageStore("/var/lib/carver/usage.db") + now = datetime.now(timezone.utc) + store.check_quota(key, now, max_conversions=100) # raises if over + store.record(key, input_bytes=len(payload), output_bytes=len(result), now=now) +""" + +from __future__ import annotations + +import sqlite3 +import threading +from datetime import datetime +from pathlib import Path + +__all__ = ["QuotaExceededError", "UsageStore"] + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS usage ( + api_key TEXT NOT NULL, + period TEXT NOT NULL, + conversions INTEGER NOT NULL DEFAULT 0, + input_bytes INTEGER NOT NULL DEFAULT 0, + output_bytes INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (api_key, period) +) +""" + + +class QuotaExceededError(Exception): + """Raised when an API key has exhausted its quota for the current period. + + Attributes: + api_key: The offending API key. + limit_name: Which limit was exceeded (``"max_conversions"`` or + ``"max_bytes"``). + limit: The configured limit value. + used: The usage value that met or exceeded the limit. + """ + + def __init__(self, api_key: str, limit_name: str, limit: int, used: int) -> None: + """Store the quota-violation details and build a readable message. + + Args: + api_key: The API key whose quota was exceeded. + limit_name: Name of the exceeded limit. + limit: Configured maximum for that limit. + used: Current usage that hit the limit. + """ + self.api_key = api_key + self.limit_name = limit_name + self.limit = limit + self.used = used + super().__init__( + f"quota exceeded for API key {api_key!r}: " + f"{limit_name}={limit} reached (used={used})" + ) + + +def _period(now: datetime) -> str: + """Return the monthly billing period for ``now`` as a ``"YYYY-MM"`` string. + + Args: + now: The instant to bucket. Naive or aware datetimes both work; the + period is derived from the datetime's own year and month. + + Returns: + The zero-padded period string, e.g. ``"2026-07"``. + """ + return f"{now.year:04d}-{now.month:02d}" + + +class UsageStore: + """Durable per-API-key usage counters with monthly quota checks. + + Each instance owns one SQLite database file. Connections are short-lived + (opened per operation) and writes are serialized behind an internal lock, + so a single instance may be shared freely across threads. + """ + + def __init__(self, db_path: str | Path) -> None: + """Open (creating if necessary) the usage database at ``db_path``. + + Args: + db_path: Filesystem path for the SQLite database. The parent + directory must already exist. + + Raises: + ValueError: If the parent directory of ``db_path`` does not exist + or ``db_path`` points at a directory, so misconfiguration + surfaces as a clear error instead of an opaque sqlite failure. + """ + path = Path(db_path) + if path.is_dir(): + raise ValueError(f"db_path points at a directory, not a file: {path}") + if not path.parent.is_dir(): + raise ValueError( + f"cannot create usage database: parent directory does not exist: " + f"{path.parent}" + ) + self._db_path = path + self._lock = threading.Lock() + with self._connect() as conn: + conn.execute(_SCHEMA) + + def _connect(self) -> sqlite3.Connection: + """Open a new short-lived connection with WAL mode enabled. + + Returns: + A fresh :class:`sqlite3.Connection` to the store's database. + """ + conn = sqlite3.connect(self._db_path, timeout=30.0) + conn.execute("PRAGMA journal_mode=WAL") + return conn + + def record( + self, api_key: str, *, input_bytes: int, output_bytes: int, now: datetime + ) -> None: + """Record one completed conversion for ``api_key`` in ``now``'s period. + + Increments the key's conversion count by one and adds the given byte + counts to its running totals for the current monthly period. The write + is durable once this method returns. + + Args: + api_key: The API key that performed the conversion. + input_bytes: Size of the input media in bytes (must be >= 0). + output_bytes: Size of the converted output in bytes (must be >= 0). + now: The current time, supplied by the caller for determinism. + + Raises: + ValueError: If ``input_bytes`` or ``output_bytes`` is negative. + """ + if input_bytes < 0 or output_bytes < 0: + raise ValueError("input_bytes and output_bytes must be non-negative") + period = _period(now) + with self._lock, self._connect() as conn: + conn.execute( + """ + INSERT INTO usage (api_key, period, conversions, input_bytes, output_bytes) + VALUES (?, ?, 1, ?, ?) + ON CONFLICT (api_key, period) DO UPDATE SET + conversions = conversions + 1, + input_bytes = input_bytes + excluded.input_bytes, + output_bytes = output_bytes + excluded.output_bytes + """, + (api_key, period, input_bytes, output_bytes), + ) + + def usage(self, api_key: str, now: datetime) -> dict: + """Return ``api_key``'s usage totals for the period containing ``now``. + + Args: + api_key: The API key to look up. + now: The current time; selects the monthly period to report. + + Returns: + A dict with keys ``period``, ``conversions``, ``input_bytes``, and + ``output_bytes``. Unknown keys (or keys with no usage this period) + report zero for all counters. + """ + period = _period(now) + with self._connect() as conn: + row = conn.execute( + "SELECT conversions, input_bytes, output_bytes FROM usage " + "WHERE api_key = ? AND period = ?", + (api_key, period), + ).fetchone() + conversions, input_bytes, output_bytes = row if row else (0, 0, 0) + return { + "period": period, + "conversions": conversions, + "input_bytes": input_bytes, + "output_bytes": output_bytes, + } + + def check_quota( + self, + api_key: str, + now: datetime, + *, + max_conversions: int | None = None, + max_bytes: int | None = None, + ) -> bool: + """Check whether ``api_key`` may perform another conversion right now. + + Call this *before* doing work; it compares current-period usage + against the given limits. A key exactly at a limit is denied (the + limit is the total number allowed, so the next request would exceed + it). + + Args: + api_key: The API key to check. + now: The current time; selects the monthly period to check. + max_conversions: Maximum conversions allowed per period, or + ``None`` for unlimited. + max_bytes: Maximum combined input+output bytes allowed per period, + or ``None`` for unlimited. + + Returns: + ``True`` if the key is within all supplied limits. + + Raises: + QuotaExceededError: If any supplied limit has been reached, with + :attr:`~QuotaExceededError.limit_name` naming which one. + """ + current = self.usage(api_key, now) + if max_conversions is not None and current["conversions"] >= max_conversions: + raise QuotaExceededError( + api_key, "max_conversions", max_conversions, current["conversions"] + ) + total_bytes = current["input_bytes"] + current["output_bytes"] + if max_bytes is not None and total_bytes >= max_bytes: + raise QuotaExceededError(api_key, "max_bytes", max_bytes, total_bytes) + return True + + def reset(self, api_key: str) -> None: + """Delete all recorded usage for ``api_key`` across every period. + + Admin operation, e.g. after a plan change or a billing dispute. + + Args: + api_key: The API key whose usage rows should be removed. + """ + with self._lock, self._connect() as conn: + conn.execute("DELETE FROM usage WHERE api_key = ?", (api_key,)) + + def all_usage(self, period: str) -> dict: + """Return usage for every API key active in ``period``. + + Admin/billing-export operation. + + Args: + period: The monthly period to report, as a ``"YYYY-MM"`` string. + + Returns: + A dict mapping each API key with recorded usage in ``period`` to + a dict of its ``conversions``, ``input_bytes``, and + ``output_bytes``. Keys with no usage that period are absent. + """ + with self._connect() as conn: + rows = conn.execute( + "SELECT api_key, conversions, input_bytes, output_bytes " + "FROM usage WHERE period = ? ORDER BY api_key", + (period,), + ).fetchall() + return { + api_key: { + "conversions": conversions, + "input_bytes": input_bytes, + "output_bytes": output_bytes, + } + for api_key, conversions, input_bytes, output_bytes in rows + }