lorebase-store 0.2.3__py3-none-any.whl

lorebase_store/profiling/sql_databases.py

@@ -0,0 +1,298 @@
+"""MySQL, MariaDB and SQL Server profiler: catalog reflection (SQLAlchemy) plus a bounded
+SELECT per table, profiled as records (``records.RecordProfiler``).
+
+The sample is ``sample_rows`` rows (``LIMIT`` or ``TOP``, from the dialect); with ``count_rows``
+each table is also counted, so inference knows when the sample is partial. Declared primary
+keys and foreign keys come from the catalog. MySQL sessions are read-only and bounded by
+``max_execution_time`` (``max_statement_time`` on MariaDB); SQL Server has no read-only
+transaction, so the login must only be able to read (its waits are bounded by the driver
+timeout and ``LOCK_TIMEOUT``). MySQL runs through asyncmy; pymssql is blocking and runs in a
+worker thread. The drivers are optional extras of the installing package (``SqlDriverError``
+names them when missing). Binary values (a ``rowversion``) are profiled as hex text.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+import math
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, replace
+from typing import Any, Literal
+
+from sqlalchemy import (
+    URL,
+    ColumnElement,
+    Connection,
+    column,
+    create_engine,
+    func,
+    inspect,
+    make_url,
+    select,
+    text,
+)
+from sqlalchemy import table as sql_table
+from sqlalchemy.exc import ArgumentError
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.pool import NullPool
+from sqlalchemy.types import BINARY, VARBINARY, LargeBinary, NullType, String, TypeEngine
+
+from lorebase_core.infer import DeclaredForeignKey, TableProfile
+from lorebase_store.profiling.records import RecordProfileError, RecordProfiler
+from lorebase_store.settings import RecordProfilingSettings, SqlProfilingSettings
+
+SqlDialect = Literal["mysql", "mssql"]
+
+APPLICATION_NAME = "lorebase-keeper"
+
+
+class SqlProfileError(ValueError):
+    """The requested tables cannot be profiled."""
+
+
+class SqlDriverError(SqlProfileError):
+    """The driver of the database is not installed."""
+
+
+@dataclass(frozen=True)
+class _Dialect:
+    schemes: frozenset[str]
+    driver_name: str
+    module: str
+    extra: str
+
+
+_DIALECTS: dict[SqlDialect, _Dialect] = {
+    "mysql": _Dialect(
+        frozenset({"mysql", "mariadb", "mysql+asyncmy", "mariadb+asyncmy"}),
+        "mysql+asyncmy",
+        "asyncmy",
+        "mysql",
+    ),
+    "mssql": _Dialect(
+        frozenset({"mssql", "sqlserver", "mssql+pymssql"}), "mssql+pymssql", "pymssql", "mssql"
+    ),
+}
+
+
+def _url(dialect: SqlDialect, location: str) -> URL:
+    """The URL on the driver of `dialect`; refuses other schemes without echoing the value."""
+    found = _DIALECTS[dialect]
+    try:
+        importlib.import_module(found.module)
+    except ImportError:
+        raise SqlDriverError(
+            f"{dialect} sources need {found.module}: install lorebase-keeper[{found.extra}]"
+        ) from None
+    try:
+        parsed = make_url(location)
+    except ArgumentError:
+        raise SqlProfileError("the connection is not a valid database URL") from None
+    if parsed.drivername not in found.schemes:
+        raise SqlProfileError(f"the connection must be a {dialect}:// URL")
+    return parsed.set(drivername=found.driver_name)
+
+
+def _plain(value: Any) -> Any:
+    """Binary values (a rowversion) as hex text, the rest unchanged."""
+    if isinstance(value, bytes | bytearray | memoryview):
+        return bytes(value).hex()
+    return value
+
+
+def _check_schema(
+    inspector: Any, schema: str | None, allowed: Sequence[str] | None, max_chars: int = 128
+) -> None:
+    """Refuse a schema that is not a plain name of this database's catalog, or that is not
+    among the `allowed` ones (the connection's own schema when none is named)."""
+    if schema is not None:
+        if not schema or len(schema) > max_chars or "." in schema or not schema.isprintable():
+            raise SqlProfileError("the schema is not a plain schema name")
+        if schema not in inspector.get_schema_names():
+            raise SqlProfileError(f"unknown schema {schema!r}")
+    if allowed is not None:
+        effective = schema if schema is not None else inspector.default_schema_name
+        if effective not in allowed:
+            raise SqlProfileError("the schema is not allowed for this source")
+
+
+def _readable(column_name: str, sql_type: TypeEngine[Any], source: Any, limit: int) -> Any:
+    """The select item for a column, bounded: long or unbounded text is cut to `limit`
+    characters, binary columns without a small declared length and columns of unknown type
+    are not read (`None`)."""
+    found = source.c[column_name]
+    if isinstance(sql_type, NullType):
+        return None
+    if isinstance(sql_type, LargeBinary | BINARY | VARBINARY):
+        length = getattr(sql_type, "length", None)
+        return found if length is not None and length <= limit else None
+    if isinstance(sql_type, String):
+        length = sql_type.length
+        if length is None or length > limit:
+            return func.substring(found, 1, limit).label(column_name)
+    return found
+
+
+def _profile_tables(
+    connection: Connection,
+    schema: str | None,
+    names: Sequence[str] | None,
+    settings: SqlProfilingSettings,
+    records: RecordProfilingSettings,
+    allowed_schemas: Sequence[str] | None = None,
+) -> list[TableProfile]:
+    inspector = inspect(connection)
+    _check_schema(inspector, schema, allowed_schemas)
+    wanted = set(names or ())
+    if len(wanted) > settings.max_tables:
+        raise SqlProfileError(
+            f"{len(wanted)} tables requested, more than the {settings.max_tables} allowed"
+        )
+    available = sorted(inspector.get_table_names(schema=schema))
+    if wanted:
+        missing = sorted(wanted - set(available))
+        if missing:
+            where = f"schema {schema!r}" if schema else "the database"
+            raise SqlProfileError(f"unknown tables in {where}: {', '.join(missing)}")
+        chosen = [name for name in available if name in wanted]
+    else:
+        chosen = available
+    if len(chosen) > settings.max_tables:
+        raise SqlProfileError(
+            f"{len(chosen)} tables requested, more than the {settings.max_tables} allowed"
+        )
+    profiles = []
+    deadline = time.monotonic() + settings.total_timeout_s
+    for name in chosen:
+        if time.monotonic() > deadline:
+            raise SqlProfileError(
+                f"reading the tables took longer than the {settings.total_timeout_s:g} s allowed"
+            )
+        info = inspector.get_columns(name, schema=schema)
+        source = sql_table(name, *(column(str(item["name"])) for item in info), schema=schema)
+        limit = settings.max_value_chars
+        items = [
+            (str(item["name"]), _readable(str(item["name"]), item["type"], source, limit))
+            for item in info
+        ]
+        read: list[tuple[str, ColumnElement[Any]]] = [(n, e) for n, e in items if e is not None]
+        if not read:
+            continue  # nothing readable: no columns, or only unreadable ones
+        columns = [n for n, _ in read]
+        profiler = RecordProfiler(records)
+        query = select(*(e for _, e in read)).select_from(source).limit(settings.sample_rows)
+        for row in connection.execute(query):
+            profiler.add({c: _plain(value) for c, value in zip(columns, row, strict=True)})
+        try:
+            profile = profiler.profile(name, columns)
+        except RecordProfileError as error:
+            raise SqlProfileError(f"{name}: {error}") from None
+        sampled = profile.row_count or 0
+        count = (
+            connection.execute(select(func.count()).select_from(source)).scalar_one()
+            if settings.count_rows
+            else None
+        )
+        primary = tuple(
+            inspector.get_pk_constraint(name, schema=schema).get("constrained_columns") or ()
+        )
+        foreign_keys = tuple(
+            DeclaredForeignKey(
+                tuple(fk["constrained_columns"]),
+                fk["referred_table"],
+                tuple(fk["referred_columns"]),
+            )
+            for fk in inspector.get_foreign_keys(name, schema=schema)
+            if fk.get("referred_schema") in (None, schema)
+        )
+        profiles.append(
+            replace(
+                profile,
+                row_count=profile.row_count if count is None else max(int(count), sampled),
+                primary_key=primary or None,
+                foreign_keys=foreign_keys,
+            )
+        )
+    return profiles
+
+
+async def _profile_mysql(
+    url: URL,
+    settings: SqlProfilingSettings,
+    run: Callable[[Connection], list[TableProfile]],
+    engine_factory: Callable[..., Any],
+) -> list[TableProfile]:
+    connect_args: dict[str, Any] = {
+        "connect_timeout": settings.connect_timeout_s,
+        "program_name": APPLICATION_NAME,
+    }
+    if "charset" not in url.query:
+        connect_args["charset"] = "utf8mb4"
+    engine = engine_factory(url, poolclass=NullPool, connect_args=connect_args)
+    try:
+        async with engine.connect() as connection:
+            # Session-wide, so that it holds for the transaction the first query opens.
+            await connection.execute(text("SET SESSION TRANSACTION READ ONLY"))
+            timeout = int(settings.statement_timeout_ms)
+            if getattr(connection.dialect, "is_mariadb", False):
+                limit = f"SET SESSION max_statement_time = {math.ceil(timeout / 1000)}"
+            else:
+                limit = f"SET SESSION max_execution_time = {timeout}"
+            await connection.execute(text(limit))
+            await connection.rollback()
+            return list(await connection.run_sync(run))
+    finally:
+        await engine.dispose()
+
+
+def _mssql_work(
+    url: URL,
+    settings: SqlProfilingSettings,
+    run: Callable[[Connection], list[TableProfile]],
+    engine_factory: Callable[..., Any],
+) -> list[TableProfile]:
+    # pymssql cannot send ApplicationIntent=ReadOnly: the login must be able to read only.
+    timeout_ms = int(settings.statement_timeout_ms)
+    engine = engine_factory(
+        url,
+        poolclass=NullPool,
+        connect_args={
+            "login_timeout": math.ceil(settings.connect_timeout_s),
+            "timeout": math.ceil(timeout_ms / 1000),
+            "appname": APPLICATION_NAME,
+        },
+    )
+    try:
+        with engine.connect() as connection:
+            connection.execute(text(f"SET LOCK_TIMEOUT {timeout_ms}"))
+            return run(connection)
+    finally:
+        engine.dispose()
+
+
+async def profile_sql_database(
+    dialect: SqlDialect,
+    location: str,
+    settings: SqlProfilingSettings,
+    records: RecordProfilingSettings,
+    *,
+    schema: str | None = None,
+    tables: Sequence[str] | None = None,
+    allowed_schemas: Sequence[str] | None = None,
+    engine_factory: Callable[..., Any] | None = None,
+) -> list[TableProfile]:
+    """Profiles of ``tables`` (default: every table) of a MySQL, MariaDB or SQL Server database
+    named by the connection URL ``location``; ``schema`` defaults to the connection's own and
+    must be a plain name found in the catalog (and in ``allowed_schemas`` when given; the
+    connection's own schema is then checked too). ``engine_factory`` replaces the engine
+    constructor (tests)."""
+    url = _url(dialect, location)
+
+    def run(connection: Connection) -> list[TableProfile]:
+        return _profile_tables(connection, schema, tables, settings, records, allowed_schemas)
+
+    if dialect == "mysql":
+        return await _profile_mysql(url, settings, run, engine_factory or create_async_engine)
+    return await asyncio.to_thread(_mssql_work, url, settings, run, engine_factory or create_engine)

lorebase_store/ratelimit.py

@@ -0,0 +1,129 @@
+"""Token buckets shared by every Keeper (`keeper.limits.rate.backend: postgres`).
+
+A bucket is one row of the UNLOGGED `rate_limit_buckets`, keyed by the SHA-256 of its name.
+`lease` refills a bucket to the database clock and takes up to `lease` whole tokens in one
+statement (an upsert, so the row lock makes it atomic across Keepers); a Keeper then serves
+requests from the tokens it leased without another round trip. Tokens leave the shared bucket
+only through leases, so all Keepers together never exceed the limit; leased tokens a Keeper
+does not use in time are lost, which only makes the limit stricter. Time is the database's
+(`clock_timestamp()`), so the clocks of the Keepers do not matter.
+
+Rows of buckets that are full again are redundant (a missing bucket is a full one): `prune`
+deletes them.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+_NOW = "extract(epoch FROM clock_timestamp())"
+_REFILLED = (
+    "LEAST(CAST(:capacity AS float8), b.tokens + GREATEST(0, EXCLUDED.updated - b.updated) "
+    "* CAST(:refill AS float8))"
+)
+_GRANT = (
+    f"CASE WHEN {_REFILLED} >= CAST(:cost AS float8) "
+    f"THEN GREATEST(CAST(:cost AS float8), LEAST(CAST(:lease AS float8), floor({_REFILLED}))) "
+    "ELSE 0 END"
+)
+_FULL_AT = "{updated} + (CAST(:capacity AS float8) - {tokens}) / CAST(:refill AS float8)"
+
+_LEASE = text(
+    f"""
+    INSERT INTO rate_limit_buckets AS b (key, tokens, updated, full_at, granted)
+    SELECT :key, CAST(:capacity AS float8) - g.taken, g.now,
+           g.now + g.taken / CAST(:refill AS float8), g.taken
+    FROM (SELECT {_NOW} AS now,
+                 LEAST(CAST(:lease AS float8), floor(CAST(:capacity AS float8))) AS taken) AS g
+    ON CONFLICT (key) DO UPDATE SET
+        tokens = {_REFILLED} - {_GRANT},
+        updated = GREATEST(b.updated, EXCLUDED.updated),
+        full_at = {
+        _FULL_AT.format(
+            updated="GREATEST(b.updated, EXCLUDED.updated)", tokens=f"({_REFILLED} - {_GRANT})"
+        )
+    },
+        granted = {_GRANT}
+    RETURNING tokens, granted, updated
+    """
+)
+
+_PEEK = text(f"SELECT tokens, updated, {_NOW} FROM rate_limit_buckets WHERE key = :key")
+_CLOCK = text(f"SELECT {_NOW}")
+_PRUNE = text(
+    f"""
+    DELETE FROM rate_limit_buckets WHERE key IN (
+        SELECT key FROM rate_limit_buckets WHERE full_at < {_NOW} LIMIT :limit
+    )
+    """
+)
+
+
+def bucket_key(name: str) -> bytes:
+    """Stored key of a bucket name."""
+    return hashlib.sha256(name.encode("utf-8")).digest()
+
+
+@dataclass(frozen=True, slots=True)
+class Lease:
+    """Tokens taken from a shared bucket: `granted` (0 when it held fewer than the cost), the
+    tokens left in it after the lease and the database time, in epoch seconds."""
+
+    granted: float
+    remaining: float
+    now: float
+
+
+@dataclass(frozen=True, slots=True)
+class BucketReading:
+    """Tokens of a bucket at `updated`, and the database time `now` (epoch seconds). A bucket
+    that has no row is full: `tokens` is None."""
+
+    tokens: float | None
+    updated: float
+    now: float
+
+
+async def lease(
+    session: AsyncSession,
+    name: str,
+    *,
+    capacity: int,
+    refill_per_s: float,
+    lease: int,
+    cost: float = 1.0,
+) -> Lease:
+    """Refill the bucket and take up to `lease` whole tokens, at least `cost`, or none when it
+    holds fewer than `cost`. A bucket without a row starts full."""
+    row = (
+        await session.execute(
+            _LEASE,
+            {
+                "key": bucket_key(name),
+                "capacity": capacity,
+                "refill": refill_per_s,
+                "lease": max(lease, cost),
+                "cost": cost,
+            },
+        )
+    ).one()
+    return Lease(granted=float(row[1]), remaining=float(row[0]), now=float(row[2]))
+
+
+async def read(session: AsyncSession, name: str) -> BucketReading:
+    """The stored state of a bucket, without changing it."""
+    row = (await session.execute(_PEEK, {"key": bucket_key(name)})).one_or_none()
+    if row is None:
+        now = float((await session.execute(_CLOCK)).scalar_one())
+        return BucketReading(None, now, now)
+    return BucketReading(float(row[0]), float(row[1]), float(row[2]))
+
+
+async def prune(session: AsyncSession, *, limit: int) -> int:
+    """Delete up to `limit` rows of buckets that are full again; the number deleted."""
+    result = await session.execute(_PRUNE, {"limit": limit})
+    return int(getattr(result, "rowcount", 0) or 0)