Parth
e7d75c3b05
* Rewrite comments in Google developer documentation style Rewrite the comments and docstrings across the backend core modules so they read plainly. The previous prose was accurate but dense and figurative, which made it slow to skim. Applies the Google developer documentation style guide: short sentences, active voice, present tense, American spelling, and no metaphors, idioms, or rhetorical asides. Replaces em-dash chains with separate sentences.
345 lines
12 KiB
Python
345 lines
12 KiB
Python
"""How many bytes the database was actually asked for.
|
|
|
|
Query *counts* are easy to see and have never been the problem here. Both
|
|
egress blowouts this project has had were one query fetching a column nobody
|
|
read: `context_snapshot` on every action, then every memory's embedding on
|
|
every turn. Counting statements would have shown nothing wrong in either case,
|
|
and at development scale, meaning ten rows and no embeddings, so would a
|
|
stopwatch.
|
|
|
|
So this meter measures bytes, and it measures them at the only place the truth
|
|
is available: the DBAPI cursor, after the driver has decoded a row and before
|
|
SQLAlchemy has turned it into anything. Every value that crosses that line
|
|
crossed the wire.
|
|
|
|
Usage:
|
|
|
|
meter = Meter()
|
|
meter.attach(engine)
|
|
with meter.scope("turn"):
|
|
... # anything that touches the database
|
|
print(meter.render())
|
|
|
|
`attach()` wraps the pool's connection factory, so it reuses the engine the app
|
|
already configured rather than building a second one. Nothing about connect
|
|
args, pre-ping, or the SQLite foreign-key pragma is repeated here, and the
|
|
metered engine cannot diverge from the real one.
|
|
Existing pooled connections are dropped first, so a connection opened before
|
|
attaching cannot quietly stay unmetered.
|
|
|
|
Sizes are of the decoded values, not of the protocol framing: the count is what
|
|
the payload weighs, and ignores per-row and per-packet overhead. Text and bytes
|
|
are exact. Numbers are charged their binary width, which is what the Postgres
|
|
binary protocol sends and close enough elsewhere. JSON is the case that matters
|
|
most and the one to be careful about: SQLite hands back the raw string (exact),
|
|
while psycopg parses `json`/`jsonb` into Python before this sees it, so the
|
|
value is re-serialised to size it. Re-serialising is within a byte or two of
|
|
the original for machine-written JSON, which is all this project stores.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from collections import defaultdict
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass, field
|
|
|
|
# The table a statement is charged to. Only the first match is used: a join
|
|
# reads mostly from its driving table, and splitting a row across tables would
|
|
# need the result metadata, which is more machinery than the answer is worth.
|
|
_TABLE_RE = re.compile(
|
|
r"\b(?:FROM|JOIN|INTO|UPDATE)\s+\"?([A-Za-z_][A-Za-z_0-9]*)\"?", re.IGNORECASE
|
|
)
|
|
_WHITESPACE_RE = re.compile(r"\s+")
|
|
|
|
|
|
def value_bytes(value) -> int:
|
|
"""The payload weight of one decoded column value."""
|
|
if value is None:
|
|
return 0
|
|
if isinstance(value, (bytes, bytearray)):
|
|
return len(value)
|
|
if isinstance(value, memoryview):
|
|
return value.nbytes
|
|
if isinstance(value, str):
|
|
return len(value.encode("utf-8"))
|
|
if isinstance(value, bool):
|
|
return 1
|
|
if isinstance(value, (int, float)):
|
|
return 8
|
|
if isinstance(value, (dict, list)):
|
|
# A JSON column the driver already parsed. psycopg parses it and SQLite
|
|
# does not. The separators match what a database emits, with no spaces.
|
|
return len(json.dumps(value, separators=(",", ":"), default=str).encode("utf-8"))
|
|
return len(str(value).encode("utf-8"))
|
|
|
|
|
|
def row_bytes(row) -> int:
|
|
return sum(value_bytes(v) for v in row)
|
|
|
|
|
|
def _table_of(statement: str) -> str:
|
|
match = _TABLE_RE.search(statement)
|
|
if match:
|
|
return match.group(1).lower()
|
|
# The statement names no table, such as BEGIN, a PRAGMA, or a savepoint.
|
|
# Group those under the keyword, so they stay countable rather than collapse
|
|
# into one "?" bucket.
|
|
head = statement.strip().split(None, 1)
|
|
return head[0].lower()[:20] if head else "(empty)"
|
|
|
|
|
|
def _one_line(statement: str, width: int = 132) -> str:
|
|
collapsed = _WHITESPACE_RE.sub(" ", statement).strip()
|
|
return collapsed if len(collapsed) <= width else collapsed[: width - 1] + "…"
|
|
|
|
|
|
@dataclass
|
|
class Tally:
|
|
statements: int = 0
|
|
rows: int = 0
|
|
fetched: int = 0 # bytes
|
|
|
|
|
|
@dataclass
|
|
class Scope:
|
|
"""What one measured stretch of work asked the database for."""
|
|
|
|
name: str
|
|
total: Tally = field(default_factory=Tally)
|
|
by_table: dict[str, Tally] = field(default_factory=lambda: defaultdict(Tally))
|
|
by_statement: dict[str, Tally] = field(default_factory=lambda: defaultdict(Tally))
|
|
|
|
def add_statement(self, statement: str) -> None:
|
|
self.total.statements += 1
|
|
self.by_table[_table_of(statement)].statements += 1
|
|
self.by_statement[statement].statements += 1
|
|
|
|
def add_rows(self, statement: str, rows: int, nbytes: int) -> None:
|
|
for tally in (
|
|
self.total,
|
|
self.by_table[_table_of(statement)],
|
|
self.by_statement[statement],
|
|
):
|
|
tally.rows += rows
|
|
tally.fetched += nbytes
|
|
|
|
|
|
class Meter:
|
|
"""Collects what the metered cursors report, grouped by scope.
|
|
|
|
Scopes nest: a statement is charged to every scope currently open, so an
|
|
inner "retrieve memories" and an outer "turn" both see it.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._open: list[Scope] = []
|
|
self.scopes: list[Scope] = []
|
|
self._attached_pools: list = []
|
|
|
|
# -------------------------------------------------------------- recording
|
|
|
|
@contextmanager
|
|
def scope(self, name: str):
|
|
scope = Scope(name)
|
|
self._open.append(scope)
|
|
try:
|
|
yield scope
|
|
finally:
|
|
self._open.pop()
|
|
self.scopes.append(scope)
|
|
|
|
def note_statement(self, statement: str) -> None:
|
|
for scope in self._open:
|
|
scope.add_statement(statement)
|
|
|
|
def note_rows(self, statement: str, rows: int, nbytes: int) -> None:
|
|
for scope in self._open:
|
|
scope.add_rows(statement, rows, nbytes)
|
|
|
|
# ------------------------------------------------------------- attaching
|
|
|
|
def attach(self, engine) -> None:
|
|
"""Meter every connection this engine opens from now on."""
|
|
# Drop pooled connections created before now; they were built by the
|
|
# unwrapped creator and would go on reporting nothing.
|
|
engine.dispose()
|
|
pool = engine.pool
|
|
creator = pool._creator # the factory create_engine() built from the URL
|
|
if getattr(creator, "_dbmeter", None) is self:
|
|
return
|
|
meter = self
|
|
|
|
def metered_creator():
|
|
return _MeteredConnection(creator(), meter)
|
|
|
|
metered_creator._dbmeter = self
|
|
pool._creator = metered_creator
|
|
self._attached_pools.append((engine, pool, creator))
|
|
|
|
def detach(self) -> None:
|
|
"""Put every metered engine back as it was.
|
|
|
|
A script exits and takes the wrapping with it; a test does not, and one
|
|
test leaving the shared engine metered would go on charging bytes to a
|
|
scope nobody opened. Pooled connections are dropped again on the way
|
|
out for the same reason attach drops them on the way in.
|
|
"""
|
|
while self._attached_pools:
|
|
engine, pool, creator = self._attached_pools.pop()
|
|
pool._creator = creator
|
|
engine.dispose()
|
|
|
|
def __enter__(self) -> "Meter":
|
|
return self
|
|
|
|
def __exit__(self, *exc) -> None:
|
|
self.detach()
|
|
|
|
# ------------------------------------------------------------- reporting
|
|
|
|
def render(self, *, statements: int = 5) -> str:
|
|
return "\n".join(render_scope(s, statements=statements) for s in self.scopes)
|
|
|
|
|
|
# ------------------------------------------------------------ DBAPI wrappers
|
|
|
|
|
|
class _MeteredConnection:
|
|
"""A DBAPI connection that hands out metered cursors.
|
|
|
|
Everything else is delegated: the dialects reach for driver-specific
|
|
attributes (`isolation_level` on SQLite, `info` and `autocommit` on
|
|
psycopg) and this must stay transparent to all of them.
|
|
"""
|
|
|
|
def __init__(self, connection, meter: Meter) -> None:
|
|
object.__setattr__(self, "_connection", connection)
|
|
object.__setattr__(self, "_meter", meter)
|
|
|
|
def cursor(self, *args, **kwargs):
|
|
return _MeteredCursor(self._connection.cursor(*args, **kwargs), self._meter)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(object.__getattribute__(self, "_connection"), name)
|
|
|
|
def __setattr__(self, name, value):
|
|
setattr(self._connection, name, value)
|
|
|
|
def __enter__(self):
|
|
self._connection.__enter__()
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
return self._connection.__exit__(*exc)
|
|
|
|
|
|
class _MeteredCursor:
|
|
"""Counts the bytes of every row handed back.
|
|
|
|
The wrapper covers the fetch methods rather than `execute`, because what a
|
|
statement costs is not known when it is sent. `SELECT * FROM memories` and
|
|
`SELECT count(*) FROM memories` look alike going out and differ by three
|
|
megabytes coming back.
|
|
"""
|
|
|
|
def __init__(self, cursor, meter: Meter) -> None:
|
|
self.__dict__["_cursor"] = cursor
|
|
self.__dict__["_meter"] = meter
|
|
self.__dict__["_statement"] = ""
|
|
|
|
# -- execution
|
|
|
|
def execute(self, statement, *args, **kwargs):
|
|
self.__dict__["_statement"] = statement
|
|
self._meter.note_statement(statement)
|
|
return self._cursor.execute(statement, *args, **kwargs)
|
|
|
|
def executemany(self, statement, *args, **kwargs):
|
|
self.__dict__["_statement"] = statement
|
|
self._meter.note_statement(statement)
|
|
return self._cursor.executemany(statement, *args, **kwargs)
|
|
|
|
# -- fetching
|
|
|
|
def _charge(self, rows) -> None:
|
|
self._meter.note_rows(
|
|
self._statement, len(rows), sum(row_bytes(r) for r in rows)
|
|
)
|
|
|
|
def fetchone(self):
|
|
row = self._cursor.fetchone()
|
|
if row is not None:
|
|
self._charge([row])
|
|
return row
|
|
|
|
def fetchmany(self, *args, **kwargs):
|
|
rows = self._cursor.fetchmany(*args, **kwargs)
|
|
self._charge(rows)
|
|
return rows
|
|
|
|
def fetchall(self):
|
|
rows = self._cursor.fetchall()
|
|
self._charge(rows)
|
|
return rows
|
|
|
|
def __iter__(self):
|
|
# Special methods are looked up on the type, so this cannot be left to
|
|
# __getattr__ the way the rest of the driver surface is.
|
|
for row in self._cursor:
|
|
self._charge([row])
|
|
yield row
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.__dict__["_cursor"], name)
|
|
|
|
def __setattr__(self, name, value):
|
|
setattr(self.__dict__["_cursor"], name, value)
|
|
|
|
def __enter__(self):
|
|
self._cursor.__enter__()
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
return self._cursor.__exit__(*exc)
|
|
|
|
|
|
# --------------------------------------------------------------- formatting
|
|
|
|
|
|
def kb(nbytes: int) -> str:
|
|
return f"{nbytes / 1024:,.1f} kB"
|
|
|
|
|
|
def render_scope(scope: Scope, *, statements: int = 5) -> str:
|
|
lines = [
|
|
"",
|
|
f"── {scope.name} " + "─" * max(0, 62 - len(scope.name)),
|
|
f" {scope.total.statements} statements · {scope.total.rows} rows · "
|
|
f"{kb(scope.total.fetched)} fetched",
|
|
]
|
|
|
|
if scope.by_table:
|
|
lines += ["", f" {'table':<20}{'stmts':>7}{'rows':>8}{'fetched':>16}"]
|
|
ranked = sorted(
|
|
scope.by_table.items(), key=lambda kv: kv[1].fetched, reverse=True
|
|
)
|
|
for table, tally in ranked:
|
|
share = tally.fetched / scope.total.fetched if scope.total.fetched else 0
|
|
lines.append(
|
|
f" {table:<20}{tally.statements:>7}{tally.rows:>8}"
|
|
f"{kb(tally.fetched):>16}{share:>7.0%}"
|
|
)
|
|
|
|
heavy = sorted(
|
|
scope.by_statement.items(), key=lambda kv: kv[1].fetched, reverse=True
|
|
)[:statements]
|
|
heavy = [(s, t) for s, t in heavy if t.fetched]
|
|
if heavy:
|
|
lines += ["", " heaviest statements"]
|
|
for statement, tally in heavy:
|
|
lines.append(f" {kb(tally.fetched):>14} {tally.statements}x {_one_line(statement)}")
|
|
|
|
return "\n".join(lines)
|