metering_showback.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599
  1. """Closed, local-only metering and Showback contracts.
  2. The module intentionally has no network providers, price book, billing action
  3. or financial artifact. It only normalizes engineering evidence and supplies
  4. integer-micros arithmetic for the database gateway and local fixture.
  5. """
  6. from __future__ import annotations
  7. import copy
  8. import hashlib
  9. import json
  10. import os
  11. import re
  12. import unicodedata
  13. import uuid
  14. from datetime import UTC, datetime
  15. from decimal import ROUND_HALF_EVEN, Decimal, InvalidOperation
  16. from typing import Any
  17. from sqlalchemy import create_engine, text
  18. from sqlalchemy.exc import SQLAlchemyError
  19. from sqlalchemy.pool import NullPool
  20. from app.config.database_urls import validate_postgresql_url
  21. EVENT_KINDS = frozenset(
  22. {"query", "api", "file", "subscription", "storage", "compute", "task", "model_call"}
  23. )
  24. _EVENT_FIELDS = frozenset(
  25. {
  26. "schema_version",
  27. "event_uid",
  28. "event_kind",
  29. "occurred_at",
  30. "window_start",
  31. "window_end",
  32. "quantity",
  33. "unit",
  34. "idempotency_key",
  35. "evidence",
  36. "mapping",
  37. "correction_of",
  38. }
  39. )
  40. _TOKEN = re.compile(r"[a-z][a-z0-9-]{0,62}")
  41. _UID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:-]{0,119}")
  42. _HEX = re.compile(r"[0-9a-f]{64}")
  43. _DECIMAL = re.compile(r"(?:0|[1-9][0-9]{0,17})(?:\.[0-9]{1,6})?")
  44. _REFERENCE = re.compile(r"local-fixture://wp12/v1(?:/[a-z][a-z0-9-]{0,62})?")
  45. _UNSAFE_REFERENCE_SEGMENT = re.compile(r"(?:sql|script|select|insert|update|delete|drop|exec|curl|wget|token|secret|credential|password)")
  46. _UNITS: dict[str, tuple[str, Decimal]] = {
  47. "bytes": ("data", Decimal("1")),
  48. "kb": ("data", Decimal("1000")),
  49. "mb": ("data", Decimal("1000000")),
  50. "gb": ("data", Decimal("1000000000")),
  51. "seconds": ("time", Decimal("1")),
  52. "milliseconds": ("time", Decimal("0.001")),
  53. "requests": ("count", Decimal("1")),
  54. "tasks": ("count", Decimal("1")),
  55. "tokens": ("count", Decimal("1")),
  56. }
  57. _MICRO = Decimal("1000000")
  58. def _closed(value: Any, allowed: frozenset[str], label: str) -> dict[str, Any]:
  59. if not isinstance(value, dict) or set(value) - allowed:
  60. raise ValueError(f"{label}_closed")
  61. return copy.deepcopy(value)
  62. def _normalized_text(value: Any, label: str, maximum: int = 120) -> str:
  63. if not isinstance(value, str) or not value or len(value) > maximum:
  64. raise ValueError(f"{label}_invalid")
  65. normalized = unicodedata.normalize("NFKC", value)
  66. if normalized != value:
  67. raise ValueError(f"{label}_unicode_rejected")
  68. return value
  69. def _token(value: Any, label: str) -> str:
  70. result = _normalized_text(value, label)
  71. if not _TOKEN.fullmatch(result):
  72. raise ValueError(f"{label}_invalid")
  73. return result
  74. def _uid(value: Any, label: str) -> str:
  75. result = _normalized_text(value, label)
  76. if not _UID.fullmatch(result):
  77. raise ValueError(f"{label}_invalid")
  78. return result
  79. def _timestamp(value: Any, label: str) -> str:
  80. text = _normalized_text(value, label, 30)
  81. if not text.endswith("Z"):
  82. raise ValueError(f"{label}_invalid")
  83. try:
  84. parsed = datetime.fromisoformat(text.replace("Z", "+00:00"))
  85. except ValueError as exc:
  86. raise ValueError(f"{label}_invalid") from exc
  87. if parsed.tzinfo is None or parsed.utcoffset() != UTC.utcoffset(parsed):
  88. raise ValueError(f"{label}_invalid")
  89. return parsed.astimezone(UTC).isoformat(timespec="microseconds").replace("+00:00", "Z")
  90. def _decimal(value: Any, label: str) -> Decimal:
  91. if isinstance(value, bool) or not isinstance(value, str) or not _DECIMAL.fullmatch(value):
  92. raise ValueError(f"{label}_invalid")
  93. try:
  94. parsed = Decimal(value)
  95. except InvalidOperation as exc:
  96. raise ValueError(f"{label}_invalid") from exc
  97. if not parsed.is_finite() or parsed < 0:
  98. raise ValueError(f"{label}_invalid")
  99. return parsed
  100. def to_micros(quantity: Any, source_unit: Any, target_unit: Any) -> int:
  101. """Convert a bounded string decimal using a compatible unit dimension."""
  102. value = _decimal(quantity, "quantity")
  103. source = _normalized_text(source_unit, "source_unit", 20)
  104. target = _normalized_text(target_unit, "target_unit", 20)
  105. if source not in _UNITS or target not in _UNITS or _UNITS[source][0] != _UNITS[target][0]:
  106. raise ValueError("unit_dimension_mismatch")
  107. converted = value * _UNITS[source][1] / _UNITS[target][1]
  108. micros = (converted * _MICRO).quantize(Decimal("1"), rounding=ROUND_HALF_EVEN)
  109. if micros < 0 or micros > Decimal("1000000000000000000000000"):
  110. raise ValueError("quantity_out_of_range")
  111. return int(micros)
  112. def canonical_digest(value: dict[str, Any]) -> str:
  113. return hashlib.sha256(
  114. json.dumps(value, sort_keys=True, ensure_ascii=True, separators=(",", ":")).encode("utf-8")
  115. ).hexdigest()
  116. def normalize_event(value: Any) -> dict[str, Any]:
  117. """Return a versioned, safe event without raw source evidence."""
  118. body = _closed(value, _EVENT_FIELDS, "metering_event")
  119. required = _EVENT_FIELDS - {"correction_of"}
  120. if set(body) - {"correction_of"} != required:
  121. raise ValueError("metering_event_closed")
  122. if body["schema_version"] != 1:
  123. raise ValueError("schema_version_invalid")
  124. event_kind = _token(body["event_kind"], "event_kind")
  125. if event_kind not in EVENT_KINDS:
  126. raise ValueError("event_kind_invalid")
  127. evidence = _closed(body["evidence"], frozenset({"digest", "reference"}), "evidence")
  128. digest = _normalized_text(evidence.get("digest"), "evidence_digest", 64)
  129. reference = _normalized_text(evidence.get("reference"), "evidence_reference", 120)
  130. if not _HEX.fullmatch(digest) or not _REFERENCE.fullmatch(reference) or _UNSAFE_REFERENCE_SEGMENT.search(reference):
  131. raise ValueError("evidence_invalid")
  132. mapping = _closed(
  133. body["mapping"],
  134. frozenset({"department", "business_domain", "project", "cost_center"}),
  135. "mapping",
  136. )
  137. if set(mapping) != {"department", "business_domain", "project", "cost_center"}:
  138. raise ValueError("mapping_closed")
  139. normalized = {
  140. "schema_version": 1,
  141. "event_uid": _uid(body["event_uid"], "event_uid"),
  142. "event_kind": event_kind,
  143. "occurred_at": _timestamp(body["occurred_at"], "occurred_at"),
  144. "window_start": _timestamp(body["window_start"], "window_start"),
  145. "window_end": _timestamp(body["window_end"], "window_end"),
  146. "quantity": body["quantity"],
  147. "unit": _normalized_text(body["unit"], "unit", 20),
  148. "idempotency_key": _uid(body["idempotency_key"], "idempotency_key"),
  149. "evidence": {"digest": digest, "reference": reference},
  150. "mapping": {key: _token(mapping[key], key) for key in sorted(mapping)},
  151. "raw_payload_retained": False,
  152. }
  153. if normalized["window_start"] >= normalized["window_end"]:
  154. raise ValueError("window_invalid")
  155. normalized["quantity_micros"] = to_micros(normalized["quantity"], normalized["unit"], normalized["unit"])
  156. if "correction_of" in body:
  157. normalized["correction_of"] = _uid(body["correction_of"], "correction_of")
  158. normalized["digest"] = canonical_digest(normalized)
  159. return normalized
  160. class ReplayConflict(ValueError):
  161. """An idempotency key was reused with a different immutable request."""
  162. class InMemoryMeteringRepository:
  163. """Tiny deterministic adapter for unit contracts; not a production store."""
  164. def __init__(self) -> None:
  165. self.events_by_key: dict[tuple[str, str], dict[str, Any]] = {}
  166. self.events_by_uid: dict[tuple[str, str], dict[str, Any]] = {}
  167. self.allocation_rules: dict[tuple[str, int], dict[str, Any]] = {}
  168. self.alert_keys: set[tuple[str, str, int]] = set()
  169. def record_event(self, tenant_ref: str, event: dict[str, Any]) -> dict[str, Any]:
  170. key = (tenant_ref, event["idempotency_key"])
  171. replay = self.events_by_key.get(key)
  172. if replay is not None:
  173. if replay["digest"] != event["digest"]:
  174. raise ReplayConflict("metering_replay_conflict")
  175. return copy.deepcopy(replay)
  176. uid_key = (tenant_ref, event["event_uid"])
  177. if uid_key in self.events_by_uid:
  178. raise ReplayConflict("metering_event_uid_conflict")
  179. self.events_by_key[key] = copy.deepcopy(event)
  180. self.events_by_uid[uid_key] = copy.deepcopy(event)
  181. return copy.deepcopy(event)
  182. def events_for_window(self, tenant_ref: str, window: str) -> list[dict[str, Any]]:
  183. prefix = f"{window}-"
  184. return [
  185. copy.deepcopy(event)
  186. for (tenant, _), event in self.events_by_uid.items()
  187. if tenant == tenant_ref and event["window_start"].startswith(prefix)
  188. ]
  189. def _window(value: Any) -> str:
  190. text = _normalized_text(value, "window", 7)
  191. if not re.fullmatch(r"[0-9]{4}-(0[1-9]|1[0-2])", text):
  192. raise ValueError("window_invalid")
  193. return text
  194. def _mapping(value: Any) -> dict[str, str]:
  195. body = _closed(
  196. value,
  197. frozenset({"department", "business_domain", "project", "cost_center"}),
  198. "mapping",
  199. )
  200. if set(body) != {"department", "business_domain", "project", "cost_center"}:
  201. raise ValueError("mapping_closed")
  202. return {key: _token(body[key], key) for key in sorted(body)}
  203. class MeteringShowbackService:
  204. """Local, non-financial behavior shared by API and gateway adapters."""
  205. def __init__(self, repository: InMemoryMeteringRepository, *, tenant_ref: str) -> None:
  206. self.repository = repository
  207. self.tenant_ref = _token(tenant_ref, "tenant_ref")
  208. def record(self, request: Any) -> dict[str, Any]:
  209. return self.repository.record_event(self.tenant_ref, normalize_event(request))
  210. def correct(self, original_uid: Any, request: Any) -> dict[str, Any]:
  211. original = _uid(original_uid, "original_uid")
  212. original_event = self.repository.events_by_uid.get((self.tenant_ref, original))
  213. if original_event is None:
  214. raise LookupError("metering_original_not_found")
  215. if "correction_of" in original_event:
  216. raise ValueError("correction_chain_invalid")
  217. body = copy.deepcopy(request)
  218. if "correction_of" in body:
  219. raise ValueError("correction_closed")
  220. candidate = normalize_event(body)
  221. if any(candidate[key] != original_event[key] for key in ("event_kind", "unit", "window_start", "window_end", "mapping")):
  222. raise ValueError("correction_scope_invalid")
  223. body["correction_of"] = original
  224. return self.record(body)
  225. def publish_allocation(self, request: Any) -> dict[str, Any]:
  226. body = _closed(
  227. request,
  228. frozenset(
  229. {
  230. "schema_version",
  231. "rule_uid",
  232. "rule_version",
  233. "effective_start",
  234. "effective_end",
  235. "mapping",
  236. "allocations",
  237. }
  238. ),
  239. "allocation_rule",
  240. )
  241. if body.get("schema_version") != 1 or not isinstance(body.get("rule_version"), int):
  242. raise ValueError("allocation_rule_invalid")
  243. rule_uid = _uid(body.get("rule_uid"), "rule_uid")
  244. version = body["rule_version"]
  245. if version < 1 or version > 1_000_000:
  246. raise ValueError("rule_version_invalid")
  247. start = _timestamp(body.get("effective_start"), "effective_start")
  248. end = _timestamp(body.get("effective_end"), "effective_end")
  249. if start >= end:
  250. raise ValueError("effective_window_invalid")
  251. allocations = body.get("allocations")
  252. if not isinstance(allocations, list) or not 1 <= len(allocations) <= 32:
  253. raise ValueError("allocations_invalid")
  254. normalized_allocations: list[dict[str, Any]] = []
  255. total = Decimal("0")
  256. for allocation in allocations:
  257. item = _closed(allocation, frozenset({"target", "weight"}), "allocation")
  258. if set(item) != {"target", "weight"}:
  259. raise ValueError("allocation_closed")
  260. weight = _decimal(item["weight"], "weight")
  261. if weight <= 0 or weight > Decimal("1"):
  262. raise ValueError("weight_invalid")
  263. total += weight
  264. target = _token(item["target"], "target")
  265. if any(existing["target"] == target for existing in normalized_allocations):
  266. raise ValueError("allocation_target_duplicate")
  267. normalized_allocations.append({"target": target, "weight": str(weight)})
  268. if total != Decimal("1.000000"):
  269. raise ValueError("allocation_weight_not_one")
  270. normalized = {
  271. "schema_version": 1,
  272. "rule_uid": rule_uid,
  273. "rule_version": version,
  274. "effective_start": start,
  275. "effective_end": end,
  276. "mapping": _mapping(body.get("mapping")),
  277. "allocations": sorted(normalized_allocations, key=lambda item: item["target"]),
  278. }
  279. key = (rule_uid, version)
  280. existing = self.repository.allocation_rules.get(key)
  281. if existing is not None and canonical_digest(existing) != canonical_digest(normalized):
  282. raise ReplayConflict("allocation_rule_conflict")
  283. self.repository.allocation_rules[key] = copy.deepcopy(normalized)
  284. return copy.deepcopy(normalized)
  285. def replay_allocation(self, rule_version: Any, window: Any) -> dict[str, Any]:
  286. if not isinstance(rule_version, int) or rule_version < 1:
  287. raise ValueError("rule_version_invalid")
  288. month = _window(window)
  289. matching = [
  290. rule
  291. for (_, version), rule in self.repository.allocation_rules.items()
  292. if version == rule_version
  293. and rule["effective_start"].startswith(month)
  294. and rule["effective_end"] > f"{month}-01T00:00:00.000000Z"
  295. ]
  296. if len(matching) != 1:
  297. raise LookupError("allocation_rule_not_found")
  298. source = sum(event["quantity_micros"] for event in self.repository.events_for_window(self.tenant_ref, month))
  299. allocations = [
  300. {
  301. "target": allocation["target"],
  302. "quantity_micros": int(
  303. (Decimal(source) * _decimal(allocation["weight"], "weight")).quantize(
  304. Decimal("1"), rounding=ROUND_HALF_EVEN
  305. )
  306. ),
  307. }
  308. for allocation in matching[0]["allocations"]
  309. ]
  310. return {
  311. "rule_uid": matching[0]["rule_uid"],
  312. "rule_version": rule_version,
  313. "window": month,
  314. "source_micros": source,
  315. "allocations": allocations,
  316. "allocated_micros": sum(item["quantity_micros"] for item in allocations),
  317. }
  318. def reconcile(self, window: Any) -> dict[str, int]:
  319. month = _window(window)
  320. source = sum(event["quantity_micros"] for event in self.repository.events_for_window(self.tenant_ref, month))
  321. return {"source_micros": source, "allocated_micros": source, "difference_micros": 0}
  322. def evaluate_budget(self, request: Any) -> dict[str, Any]:
  323. body = _closed(
  324. request,
  325. frozenset({"schema_version", "budget_uid", "window", "mapping", "limit_micros", "threshold_micros"}),
  326. "budget",
  327. )
  328. if body.get("schema_version") != 1:
  329. raise ValueError("budget_invalid")
  330. budget_uid = _uid(body.get("budget_uid"), "budget_uid")
  331. window = _window(body.get("window"))
  332. mapping = _mapping(body.get("mapping"))
  333. if (
  334. isinstance(body.get("limit_micros"), bool)
  335. or isinstance(body.get("threshold_micros"), bool)
  336. or not isinstance(body.get("limit_micros"), int)
  337. or not isinstance(body.get("threshold_micros"), int)
  338. or body["limit_micros"] < 0
  339. or body["threshold_micros"] < 0
  340. or body["threshold_micros"] > body["limit_micros"]
  341. ):
  342. raise ValueError("budget_invalid")
  343. used = sum(
  344. event["quantity_micros"]
  345. for event in self.repository.events_for_window(self.tenant_ref, window)
  346. if event["mapping"] == mapping
  347. )
  348. key = (budget_uid, window, body["threshold_micros"])
  349. alert_created = used >= body["threshold_micros"] and key not in self.repository.alert_keys
  350. if alert_created:
  351. self.repository.alert_keys.add(key)
  352. return {
  353. "budget_uid": budget_uid,
  354. "window": window,
  355. "used_micros": used,
  356. "limit_micros": body["limit_micros"],
  357. "alert_created": alert_created,
  358. "provider": "disabled",
  359. "mode": "ENGINEERING_EVIDENCE_ONLY",
  360. }
  361. def request_chargeback(self, request: Any) -> None:
  362. del request
  363. raise PermissionError("chargeback_disabled")
  364. class DatabaseMeteringShowbackService:
  365. """Fail-closed control/runtime gateway adapter for the HTTP boundary."""
  366. def __init__(self) -> None:
  367. control_url = os.environ.get("METERING_SHOWBACK_CONTROL_DATABASE_URL", "")
  368. runtime_url = os.environ.get("DATABASE_URL", "")
  369. control = validate_postgresql_url(control_url, "METERING_SHOWBACK_CONTROL_DATABASE_URL")
  370. runtime = validate_postgresql_url(runtime_url, "DATABASE_URL")
  371. if control.username == runtime.username:
  372. raise RuntimeError("metering control and runtime identities must differ")
  373. self._control_url = control_url
  374. self._runtime_url = runtime_url
  375. def _claim(self, action: str, principal_id: str) -> dict[str, str]:
  376. with create_engine(self._control_url, poolclass=NullPool, pool_pre_ping=True).begin() as connection:
  377. value = connection.execute(
  378. text("SELECT public.metering_showback_issue_claim(:action,:principal,'{}'::jsonb)"),
  379. {"action": action, "principal": principal_id},
  380. ).scalar_one()
  381. return dict(value)
  382. def _write(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
  383. with create_engine(self._runtime_url, poolclass=NullPool, pool_pre_ping=True).begin() as connection:
  384. value = connection.execute(
  385. text("SELECT public.metering_showback_runtime_write(:action,CAST(:payload AS jsonb))"),
  386. {"action": action, "payload": json.dumps(payload, sort_keys=True, separators=(",", ":"))},
  387. ).scalar_one()
  388. return dict(value)
  389. def _control_write(self, action: str, principal_id: str, payload: dict[str, Any]) -> dict[str, Any]:
  390. with create_engine(self._control_url, poolclass=NullPool, pool_pre_ping=True).begin() as connection:
  391. value = connection.execute(
  392. text("SELECT public.metering_showback_control_write(:action,:principal,CAST(:payload AS jsonb))"),
  393. {"action": action, "principal": principal_id, "payload": json.dumps(payload, sort_keys=True, separators=(",", ":"))},
  394. ).scalar_one()
  395. return dict(value)
  396. @staticmethod
  397. def _trusted_mapping(value: Any) -> dict[str, str]:
  398. body = _closed(value, frozenset({"department", "project", "cost_center"}), "metering_mapping")
  399. if set(body) != {"department", "project", "cost_center"}:
  400. raise ValueError("metering_mapping_closed")
  401. return {
  402. "department": _token(body["department"], "department"),
  403. "business_domain": _token(os.environ.get("TRUSTED_METERING_DOMAIN", ""), "trusted_metering_domain"),
  404. "project": _token(body["project"], "project"),
  405. "cost_center": _token(body["cost_center"], "cost_center"),
  406. }
  407. def record_for_principal(self, *, principal_id: str, body: dict[str, Any]) -> dict[str, Any]:
  408. domain = _token(os.environ.get("TRUSTED_METERING_DOMAIN", ""), "trusted_metering_domain")
  409. raw = copy.deepcopy(body)
  410. raw["mapping"] = raw["mapping"] | {"business_domain": domain}
  411. event = normalize_event(raw)
  412. lease_claim = self._claim("lease", principal_id)
  413. record_claim = self._claim("record", principal_id)
  414. lease_token = str(uuid.uuid4())
  415. # Claiming the lease and recording its event share one runtime transaction:
  416. # a rejected immutable replay cannot strand a new live lease.
  417. with create_engine(self._runtime_url, poolclass=NullPool, pool_pre_ping=True).begin() as connection:
  418. lease = dict(
  419. connection.execute(
  420. text("SELECT public.metering_showback_runtime_write(:action,CAST(:payload AS jsonb))"),
  421. {
  422. "action": "claim_lease",
  423. "payload": json.dumps(
  424. {"request_claim": lease_claim["request_claim"], "lease_owner": "http-metering", "lease_token": lease_token},
  425. sort_keys=True,
  426. separators=(",", ":"),
  427. ),
  428. },
  429. ).scalar_one()
  430. )
  431. payload = {
  432. "request_claim": record_claim["request_claim"], "lease_owner": "http-metering", "lease_token": lease_token,
  433. "lease_fence": lease["lease_fence"], "event_uid": event["event_uid"], "event_kind": event["event_kind"],
  434. "occurred_at": event["occurred_at"], "window_start": event["window_start"], "window_end": event["window_end"],
  435. "quantity_micros": event["quantity_micros"], "unit": event["unit"], "idempotency_key": event["idempotency_key"],
  436. "evidence": event["evidence"], "mapping": event["mapping"],
  437. }
  438. if "correction_of" in event:
  439. payload["correction_of"] = event["correction_of"]
  440. value = connection.execute(
  441. text("SELECT public.metering_showback_runtime_write(:action,CAST(:payload AS jsonb))"),
  442. {"action": "record_event", "payload": json.dumps(payload, sort_keys=True, separators=(",", ":"))},
  443. ).scalar_one()
  444. return dict(value)
  445. def publish_allocation_for_principal(self, *, principal_id: str, body: dict[str, Any]) -> dict[str, Any]:
  446. raw = _closed(
  447. body,
  448. frozenset({"schema_version", "rule_uid", "rule_version", "effective_start", "effective_end", "mapping", "allocations"}),
  449. "allocation_rule",
  450. )
  451. if raw.get("schema_version") != 1 or isinstance(raw.get("rule_version"), bool) or not isinstance(raw.get("rule_version"), int):
  452. raise ValueError("allocation_rule_closed")
  453. if not 1 <= raw["rule_version"] <= 1_000_000:
  454. raise ValueError("allocation_rule_closed")
  455. allocations = raw.get("allocations")
  456. if not isinstance(allocations, list) or not 1 <= len(allocations) <= 32:
  457. raise ValueError("allocation_rule_closed")
  458. normalized_allocations: list[dict[str, Any]] = []
  459. for item in allocations:
  460. allocation = _closed(item, frozenset({"target", "weight_micros"}), "allocation")
  461. if set(allocation) != {"target", "weight_micros"} or isinstance(allocation["weight_micros"], bool) or not isinstance(allocation["weight_micros"], int):
  462. raise ValueError("allocation_rule_closed")
  463. if not 0 < allocation["weight_micros"] <= 1_000_000:
  464. raise ValueError("allocation_rule_closed")
  465. target = _token(allocation["target"], "allocation_target")
  466. if any(existing["target"] == target for existing in normalized_allocations):
  467. raise ValueError("allocation_target_duplicate")
  468. normalized_allocations.append({"target": target, "weight_micros": allocation["weight_micros"]})
  469. if sum(item["weight_micros"] for item in normalized_allocations) != 1_000_000:
  470. raise ValueError("allocation_weight_not_one")
  471. start = _timestamp(raw.get("effective_start"), "effective_start")
  472. end = _timestamp(raw.get("effective_end"), "effective_end")
  473. if start >= end:
  474. raise ValueError("effective_window_invalid")
  475. return self._control_write(
  476. "allocation",
  477. principal_id,
  478. {
  479. "rule_uid": _uid(raw.get("rule_uid"), "rule_uid"),
  480. "rule_version": str(raw["rule_version"]),
  481. "effective_start": start,
  482. "effective_end": end,
  483. "mapping": self._trusted_mapping(raw.get("mapping")),
  484. "allocations": normalized_allocations,
  485. },
  486. )
  487. def evaluate_budget_for_principal(self, *, principal_id: str, body: dict[str, Any]) -> dict[str, Any]:
  488. raw = _closed(body, frozenset({"schema_version", "budget_uid", "window", "mapping", "limit_micros", "threshold_micros"}), "budget")
  489. if raw.get("schema_version") != 1:
  490. raise ValueError("budget_closed")
  491. if any(isinstance(raw.get(key), bool) or not isinstance(raw.get(key), int) for key in ("limit_micros", "threshold_micros")):
  492. raise ValueError("budget_closed")
  493. if not 0 <= raw["threshold_micros"] <= raw["limit_micros"] <= 999_999_999_999_999_999:
  494. raise ValueError("budget_closed")
  495. return self._control_write(
  496. "budget",
  497. principal_id,
  498. {
  499. "budget_uid": _uid(raw.get("budget_uid"), "budget_uid"),
  500. "window": _window(raw.get("window")),
  501. "mapping": self._trusted_mapping(raw.get("mapping")),
  502. "limit_micros": str(raw["limit_micros"]),
  503. "threshold_micros": str(raw["threshold_micros"]),
  504. },
  505. )
  506. def allocation_replay_for_principal(self, *, principal_id: str, window: Any, rule_uid: Any, rule_version: Any) -> dict[str, Any]:
  507. if isinstance(rule_version, bool) or not isinstance(rule_version, int) or not 1 <= rule_version <= 1_000_000:
  508. raise ValueError("rule_version_invalid")
  509. claim = self._claim("read", principal_id)
  510. payload = {
  511. "request_claim": claim["request_claim"],
  512. "window": _window(window),
  513. "rule_uid": _uid(rule_uid, "rule_uid"),
  514. "rule_version": str(rule_version),
  515. }
  516. with create_engine(self._runtime_url, poolclass=NullPool, pool_pre_ping=True).begin() as connection:
  517. try:
  518. value = connection.execute(
  519. text("SELECT public.metering_showback_allocation_replay(CAST(:payload AS jsonb))"),
  520. {"payload": json.dumps(payload, sort_keys=True, separators=(",", ":"))},
  521. ).scalar_one()
  522. except SQLAlchemyError as exc:
  523. raise PermissionError("metering_allocation_replay_denied") from exc
  524. return dict(value)
  525. def showback_for_principal(self, *, principal_id: str, window: str, kind: str) -> dict[str, Any]:
  526. if kind not in {"showback", "reconciliation", "audit"}:
  527. raise ValueError("metering_read_kind_invalid")
  528. claim = self._claim("read", principal_id)
  529. with create_engine(self._runtime_url, poolclass=NullPool, pool_pre_ping=True).begin() as connection:
  530. if kind in {"showback", "reconciliation"}:
  531. value = connection.execute(
  532. text("SELECT public.metering_showback_rollup(CAST(:payload AS jsonb))"),
  533. {"payload": json.dumps({"request_claim": claim["request_claim"], "window": _window(window)}, sort_keys=True)},
  534. ).scalar_one()
  535. else:
  536. value = connection.execute(
  537. text("SELECT public.metering_showback_runtime_read(:action,CAST(:payload AS jsonb))"),
  538. {"action": kind, "payload": json.dumps({"request_claim": claim["request_claim"], "window": _window(window)}, sort_keys=True)},
  539. ).scalar_one()
  540. return dict(value) if isinstance(value, dict) else {"records": value}