From 1c01ccc5ace173803a2b1f9d64cece6c557ecf9f Mon Sep 17 00:00:00 2001 From: AivanF Date: Thu, 6 Aug 2026 14:32:26 +0300 Subject: [PATCH 1/5] BigQuery user-scoped OAuth creds --- docs/configuration/datasources.md | 18 +- slayer/core/models.py | 14 ++ slayer/engine/cache.py | 4 +- slayer/engine/query_engine.py | 41 +++-- slayer/engine/schema_drift.py | 5 +- slayer/sql/client.py | 97 +++++++++++ slayer/sql/dialects/base.py | 37 ++++ slayer/sql/dialects/bigquery.py | 190 +++++++++++++++++--- slayer/sql/engine_factory.py | 144 ++++++++++++++-- tests/dialects/test_bigquery.py | 192 ++++++++++++++++++++- tests/dialects/test_tsql.py | 3 +- tests/integration/test_in_memory_sqlite.py | 6 +- tests/test_engine_factory.py | 156 +++++++++++++++++ tests/test_query_cache.py | 2 +- tests/test_sql_client.py | 95 ++++++++++ tests/test_sql_generator.py | 10 +- 16 files changed, 950 insertions(+), 64 deletions(-) diff --git a/docs/configuration/datasources.md b/docs/configuration/datasources.md index 33f424da..a716ab87 100644 --- a/docs/configuration/datasources.md +++ b/docs/configuration/datasources.md @@ -160,12 +160,28 @@ Statement-level timeout is enforced via | `username` | string | No | Database username | | `password` | string | No | Database password | | `connection_string` | string | No | Full connection string (alternative to individual fields) | -| `credentials_json` | string | No | Credentials JSON (used for BigQuery service accounts) | +| `credentials_json` | string | No | BigQuery service-account key JSON | +| `oauth_credentials_json` | string | No | BigQuery OAuth authorized-user grant JSON | | `schema_name` | string | No | Default schema name | !!! note Both `username` and `user` field names are accepted. The `user` alias is automatically mapped to `username` for compatibility with common database tooling conventions. +### BigQuery credentials + +Three ways to authenticate, in the order SLayer prefers them: + +1. **`oauth_credentials_json`** — an OAuth *authorized user* grant, the shape `google.oauth2.credentials.Credentials.from_authorized_user_info` accepts (`token`, `refresh_token`, `client_id`, `client_secret`, `token_uri`). Queries run as that end user, against that user's own BigQuery permissions. The grant carries no project, so the connection string must name one: `bigquery:///`. +2. **`credentials_json`** — a service-account key file's contents. One shared identity for every caller. +3. **Neither** — Application Default Credentials (`GOOGLE_APPLICATION_CREDENTIALS`, or the attached compute identity). + +Setting both (1) and (2) is an error rather than a silent precedence win: they name different identities, and guessing is how a per-user query quietly runs as the shared service account. + +An authorized-user grant will **not** work in `credentials_json` — `sqlalchemy-bigquery` routes that field to `service_account.Credentials.from_service_account_info`, which only understands key files. SLayer detects the mix-up and says so. + +!!! note "Engine caching and credentials" + Cached engines are keyed by credentials as well as by URL, so two identities never share a connection pool. The cache is bounded (default 64 engines, `SLAYER_MAX_CACHED_ENGINES`) and evicts least-recently-used, which matters once per-user credentials make its size track *users* rather than datasources. For OAuth grants the key uses the durable part of the grant, so a refreshed access token reuses its engine instead of leaking a new one. An engine whose credentials get rejected — a revoked grant, a rotated key — is dropped and rebuilt on the next call. + ## Ingesting at Startup To run idempotent auto-ingestion across every configured datasource each time `slayer serve` or `slayer mcp` boots, pass `--ingest-on-startup` (or set `SLAYER_INGEST_ON_STARTUP=1`). See [Ingesting at Startup](../concepts/ingestion.md#ingesting-at-startup) for the full contract. diff --git a/slayer/core/models.py b/slayer/core/models.py index 293638d6..fba03c76 100644 --- a/slayer/core/models.py +++ b/slayer/core/models.py @@ -792,6 +792,20 @@ class DatasourceConfig(BaseModel): # When unset, BigQuery falls back to Application Default Credentials # (``GOOGLE_APPLICATION_CREDENTIALS`` env var or attached compute identity). credentials_json: str | None = Field(default=None, repr=False) + # BigQuery-specific. A Google OAuth *authorized user* grant as a JSON + # string — the shape ``google.oauth2.credentials.Credentials + # .from_authorized_user_info`` consumes (``token``, ``refresh_token``, + # ``client_id``, ``client_secret``, ``token_uri``, ``scopes``). This is + # the per-end-user auth path: the caller obtains the grant from its own + # OAuth flow and hands SLayer a datasource carrying it, so queries run + # as that user with that user's BigQuery permissions. + # + # Mutually exclusive with ``credentials_json`` (service account, one + # shared identity for everyone). ``credentials_json`` cannot carry an + # OAuth grant: ``sqlalchemy-bigquery`` feeds it to + # ``service_account.Credentials.from_service_account_info``, which only + # understands service-account keys. + oauth_credentials_json: str | None = Field(default=None, repr=False) @model_validator(mode="before") @classmethod diff --git a/slayer/engine/cache.py b/slayer/engine/cache.py index 8ca2631f..0090542b 100644 --- a/slayer/engine/cache.py +++ b/slayer/engine/cache.py @@ -27,6 +27,8 @@ from sqlglot.optimizer.normalize_identifiers import normalize_identifiers from sqlglot.optimizer.scope import Scope, traverse_scope +from slayer.sql.engine_factory import EngineCacheKey + # Marker alias prefix for refresh-key scan projections. Also lets tests # distinguish a refresh-key scan query from a data query. _RK_ALIAS_PREFIX = "slayer_rk_" @@ -123,7 +125,7 @@ class _CacheEntry(BaseModel): sql: str ds_fingerprint: str dialect: str - ds_key: tuple[str, str] + ds_key: EngineCacheKey resolved_data_source: str | None = None original_input: Any = None variables: dict[str, Any] | None = None diff --git a/slayer/engine/query_engine.py b/slayer/engine/query_engine.py index 8f2c0fc9..8e48cb49 100644 --- a/slayer/engine/query_engine.py +++ b/slayer/engine/query_engine.py @@ -73,7 +73,8 @@ from slayer.sql.client import SlayerSQLClient from slayer.sql.dialects import SqlDialect, dialect_for_ds_type, get_dialect from slayer.sql import engine_factory -from slayer.sql.engine_factory import _runtime_fingerprint +from slayer.sql.engine_factory import EngineCacheKey +from slayer.sql.engine_factory import _cache_key as _engine_cache_key from slayer.sql.generator import SQLGenerator from slayer.sql.reserved_keywords import SLAYER_RESERVED_KEYWORDS from slayer.sql.session_policy import ScopedTable, apply_session_policy @@ -82,15 +83,17 @@ logger = logging.getLogger(__name__) -def _sql_client_cache_key(datasource: DatasourceConfig) -> tuple[str, str]: +def _sql_client_cache_key(datasource: DatasourceConfig) -> EngineCacheKey: """Cache key for ``SlayerQueryEngine._sql_clients``. - Mirrors ``engine_factory``'s cache key so two datasources differing - in (e.g.) Snowflake ``warehouse`` get distinct ``SlayerSQLClient`` - instances (and therefore distinct factory-cached engines with the - correct per-connection ``USE`` listener). + Delegates to ``engine_factory``'s key builder rather than re-deriving it: + each client memoizes the engine it got from the factory, so if the two keys + disagreed a caller could be handed a client whose engine was built for + different credentials. Sharing one implementation makes that impossible. """ - return (datasource.get_connection_string(), _runtime_fingerprint(datasource)) + return _engine_cache_key( + datasource, datasource.get_connection_string(), + ) class _ResolvedItem(BaseModel): @@ -495,7 +498,7 @@ class _Prepared(BaseModel): sql: str attributes: "ResponseAttributes" expected_columns: list[str] - ds_key: tuple[str, str] + ds_key: EngineCacheKey ds_fingerprint: str @@ -549,7 +552,7 @@ def __init__( # ``engine_factory``'s cache so Snowflake datasources sharing a # connection_name but differing in warehouse/role/database/schema # get distinct clients (DEV-1551). - self._sql_clients: dict[tuple[str, str], SlayerSQLClient] = {} + self._sql_clients: dict[EngineCacheKey, SlayerSQLClient] = {} # DEV-1578: immutable, engine-global forced-filter policy. When set, # every generated SQL is rewritten to scope each physical table to the # configured tenant before execution / dry-run / explain / profiling. @@ -563,7 +566,7 @@ def __init__( # datasource, for the correlated-subquery join-rule gate. ``None`` (or # a missing entry) fails closed. Populated by # ``_preflight_clickhouse_correlated`` before the policy rewrite. - self._ch_version_cache: dict[tuple[str, str], tuple[int, int] | None] = {} + self._ch_version_cache: dict[EngineCacheKey, tuple[int, int] | None] = {} # DEV-1587: per-engine, in-memory, opt-in query result cache. The # cache is local to this engine instance so two engines with # different RLS / connection settings keep separate caches. @@ -1357,7 +1360,7 @@ async def _execute_pipeline( return await self._run_and_build(prepared=prepared, client=client) def _get_client( - self, datasource: DatasourceConfig, ds_key: tuple[str, str] + self, datasource: DatasourceConfig, ds_key: EngineCacheKey ) -> SlayerSQLClient: """Reuse (or open) the SQL client + connection pool for a datasource. @@ -1604,21 +1607,21 @@ async def _run_and_cleanup() -> RefreshResult: async def _refresh_scan_all( self, snapshot: "dict[str, _CacheEntry]", result: RefreshResult, - ) -> "tuple[dict[tuple[tuple[str, str], str, str], Any], set]": + ) -> "tuple[dict[tuple[EngineCacheKey, str, str], Any], set]": """Collate applicable ``(ds_key, table)`` scan targets across all entries and run ONE batched scan per pair. Scan failures are recorded as continue-on-error ``RefreshError(phase="refresh_key_scan")`` and the table is marked failed so its dependent entries are left unchanged.""" - targets: dict[tuple[tuple[str, str], str], set] = {} - dialects: dict[tuple[str, str], str] = {} - clients: dict[tuple[str, str], SlayerSQLClient] = {} + targets: dict[tuple[EngineCacheKey, str], set] = {} + dialects: dict[EngineCacheKey, str] = {} + clients: dict[EngineCacheKey, SlayerSQLClient] = {} for entry in snapshot.values(): ds_key = tuple(entry.ds_key) dialects[ds_key] = entry.dialect for table, expr in entry.applicable: targets.setdefault((ds_key, table), set()).add(expr) - fresh: dict[tuple[tuple[str, str], str, str], Any] = {} + fresh: dict[tuple[EngineCacheKey, str, str], Any] = {} failed: set = set() for (ds_key, table), exprs in targets.items(): expr_list = sorted(exprs) @@ -1639,9 +1642,9 @@ async def _refresh_scan_all( async def _client_for_refresh( self, - ds_key: tuple[str, str], + ds_key: EngineCacheKey, snapshot: "dict[str, _CacheEntry]", - clients: dict[tuple[str, str], SlayerSQLClient], + clients: dict[EngineCacheKey, SlayerSQLClient], ) -> SlayerSQLClient: """Get the cached client for ``ds_key``, reopening it from an entry's recorded resolved datasource name if it isn't cached.""" @@ -1671,7 +1674,7 @@ async def _refresh_one( *, key: str, entry: _CacheEntry, - fresh_values: "dict[tuple[tuple[str, str], str, str], Any]", + fresh_values: "dict[tuple[EngineCacheKey, str, str], Any]", failed_tables: set, result: RefreshResult, ) -> None: diff --git a/slayer/engine/schema_drift.py b/slayer/engine/schema_drift.py index 5997e846..faf45d62 100644 --- a/slayer/engine/schema_drift.py +++ b/slayer/engine/schema_drift.py @@ -49,6 +49,7 @@ _sa_type_to_data_type, ) from slayer.sql.client import SlayerSQLClient +from slayer.sql.engine_factory import EngineCacheKey logger = logging.getLogger(__name__) @@ -2051,7 +2052,7 @@ async def _collect_sql_diffs( *, datasource: DatasourceConfig, sql_models: list[SlayerModel], - sql_clients: dict[tuple[str, str], SlayerSQLClient] | None, + sql_clients: dict[EngineCacheKey, SlayerSQLClient] | None, ) -> dict[str, tuple[ToDeleteEntry | None, set[str]]]: """Trial-execute each sql-mode model concurrently and produce its diff.""" out: dict[str, tuple[ToDeleteEntry | None, set[str]]] = {} @@ -2087,7 +2088,7 @@ async def validate_datasource( *, datasource: DatasourceConfig, models: list[SlayerModel], - sql_clients: dict[tuple[str, str], SlayerSQLClient] | None = None, + sql_clients: dict[EngineCacheKey, SlayerSQLClient] | None = None, ) -> list[ToDeleteEntry]: """Validate every persisted model in ``models`` (all in the same DS) against the live schema of ``datasource``. Read-only. diff --git a/slayer/sql/client.py b/slayer/sql/client.py index 4d66a302..eef82b07 100644 --- a/slayer/sql/client.py +++ b/slayer/sql/client.py @@ -528,12 +528,49 @@ def _get_sync_engine_for_client(self) -> sa.Engine | None: self._sync_engine = engine_factory.get_engine(self.datasource) return self._sync_engine + def _discard_engine_on_auth_failure(self, exc: BaseException) -> None: + """Drop the cached engine when the failure was a credential rejection. + + The credentials an engine authenticates with are fixed at construction + (BigQuery bakes a client object into the engine), so a revoked OAuth + grant or rotated key poisons that engine permanently — every later call + through the cache reproduces the same failure. Evicting means the next + call rebuilds from whatever credentials the datasource now carries, + which is exactly what a caller that just refreshed its grant expects. + + Best-effort and non-fatal: the original error is what the caller needs + to see, so a failure to clean up must not displace it. + """ + if not _is_auth_failure(exc): + return + from slayer.sql import engine_factory # noqa: PLC0415 + self._sync_engine = None + try: + engine_factory.invalidate_engine(self.datasource) + except Exception: + logger.warning( + "Failed to invalidate engine for datasource %r after an " + "authentication failure.", self.datasource.name, exc_info=True, + ) + async def execute( self, sql: str, timeout_seconds: int = 120, ) -> list[dict[str, Any]]: """Execute SQL asynchronously.""" + try: + return await self._execute(sql=sql, timeout_seconds=timeout_seconds) + except Exception as exc: + self._discard_engine_on_auth_failure(exc) + raise + + async def _execute( + self, + *, + sql: str, + timeout_seconds: int, + ) -> list[dict[str, Any]]: async_engine = self._get_async_engine() db_type = self.datasource.type if async_engine is not None: @@ -647,6 +684,66 @@ def _is_transient_db_error(exc: BaseException) -> bool: return any(sig in msg for sig in _TRANSIENT_DB_ERROR_SIGNALS) +# Credential rejection, as opposed to a transient blip. Distinct from +# _TRANSIENT_DB_ERROR_SIGNALS on purpose: retrying these is pointless (the +# credentials are baked into the engine, so every attempt fails identically), +# and the right response is to throw the engine away rather than sleep. +# +# Deliberately narrow. Postgres' table-level "permission denied for table" +# is NOT here — the credentials worked fine, the grant didn't, and evicting +# a healthy engine over it is pure pool churn. +_AUTH_ERROR_SIGNALS = ( + "invalid_grant", # OAuth refresh token revoked / expired + "invalid_client", + "unauthorized_client", + "token has been expired or revoked", + "could not refresh access token", + "reauthentication is needed", + "authentication failed", + "password authentication failed", # libpq + "invalid credentials", + "invalid username or password", +) + +# Matched by class name so this stays dependency-free — google-auth and +# google-api-core ship only with the optional 'bigquery' extra. +_AUTH_ERROR_TYPE_NAMES = frozenset({ + "RefreshError", # google.auth.exceptions + "DefaultCredentialsError", # google.auth.exceptions + "Unauthorized", # google.api_core.exceptions — HTTP 401 +}) + + +def _is_auth_failure(exc: BaseException) -> bool: + """Return True when the server rejected the *credentials* themselves. + + Walks the cause/context chain (plus SQLAlchemy's ``orig``) because the + signal is almost always a driver exception wrapped one or two layers deep + by the time it surfaces. + """ + seen: set[int] = set() + pending: list[BaseException] = [exc] + while pending: + current = pending.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + if type(current).__name__ in _AUTH_ERROR_TYPE_NAMES: + return True + text = str(current).lower() + if any(signal in text for signal in _AUTH_ERROR_SIGNALS): + return True + pending.extend( + nested for nested in ( + getattr(current, "orig", None), + current.__cause__, + current.__context__, + ) + if isinstance(nested, BaseException) + ) + return False + + async def _retry_with_backoff( *, sql: str, diff --git a/slayer/sql/dialects/base.py b/slayer/sql/dialects/base.py index 7af32f60..e5cfdadb 100644 --- a/slayer/sql/dialects/base.py +++ b/slayer/sql/dialects/base.py @@ -12,6 +12,7 @@ from __future__ import annotations +import hashlib from functools import lru_cache from typing import TYPE_CHECKING, Any from collections.abc import Callable @@ -167,6 +168,17 @@ def _sqlglot_backslash_escapes(sqlglot_name: str) -> bool: return "\\" in escapes +def _digest(secret: str | None) -> str: + """Short, stable, non-reversible id for secret material used in cache keys. + + Truncated to 16 hex chars: collision risk is negligible for the number of + live credentials in a process, and short keys stay readable in logs. + """ + if not secret: + return "" + return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16] + + class SqlDialect(BaseModel): """Strategy class encapsulating one database's SQL-generation quirks. @@ -593,6 +605,31 @@ def build_engine( """ return None + # ------------------------------------------------------------------ + # Credential identity (engine-cache safety) + # ------------------------------------------------------------------ + + def credential_fingerprint(self, datasource: "DatasourceConfig") -> str: + """Opaque identity of the credentials this datasource authenticates with. + + Engines are cached per ``(connection_string, runtime_fingerprint, + credential_fingerprint)``. Any dialect whose secret does **not** appear + in the connection string MUST override this, or two callers holding + different credentials for the same URL will silently share one engine — + the first caller's identity then serves everyone for the life of the + process. + + The default covers the common case safely: username/password dialects + embed their credentials in the URL, so the connection string already + distinguishes them and ``""`` adds nothing. ``credentials_json`` is + hashed here rather than left out, so a dialect that carries it + out-of-band is keyed correctly even before it overrides this. + + Return a digest or a stable subject id — **never** raw secret material, + since cache keys reach logs and error messages. + """ + return _digest(datasource.credentials_json) + def apply_session_overrides( self, dbapi_connection: Any, diff --git a/slayer/sql/dialects/bigquery.py b/slayer/sql/dialects/bigquery.py index 0eace4ac..d8fa8460 100644 --- a/slayer/sql/dialects/bigquery.py +++ b/slayer/sql/dialects/bigquery.py @@ -35,7 +35,7 @@ from slayer.core.enums import TimeGranularity from slayer.sql.dialects._alias_mangle import decode_alias, encode_alias -from slayer.sql.dialects.base import SqlDialect +from slayer.sql.dialects.base import SqlDialect, _digest if TYPE_CHECKING: from slayer.core.models import DatasourceConfig @@ -67,6 +67,58 @@ _DOTTED_ALIAS_RE = re.compile(r"`(\w+(?:\.\w+)+)`", re.ASCII) +# --------------------------------------------------------------------------- +# Credential parsing +# --------------------------------------------------------------------------- + + +# ``type`` marker Google writes into an OAuth authorized-user JSON, as +# opposed to ``"service_account"`` in a key file. +_AUTHORIZED_USER_TYPE = "authorized_user" + +# Fields of an authorized-user grant that change on every token refresh +# without changing *whose* grant it is. +_ROTATING_OAUTH_FIELDS = frozenset({"token", "access_token", "expiry", "id_token"}) + + +def _parse_credentials_object( + *, + raw: str | None, + field: str, + datasource_name: str, +) -> dict[str, Any]: + """Decode a credentials JSON string, or raise naming the offending field.""" + try: + parsed = json.loads(raw or "") + except json.JSONDecodeError as exc: + raise ValueError( + f"Datasource '{datasource_name}': {field} is not valid JSON: {exc}" + ) from exc + if not isinstance(parsed, dict): + raise ValueError( + f"Datasource '{datasource_name}': {field} must be a JSON object" + ) + return parsed + + +def _durable_oauth_material(raw: str) -> str: + """Canonical string identifying *whose* OAuth grant ``raw`` is. + + Strips the rotating token fields when a refresh token is present, so a + refreshed grant keeps its cache identity. Unparseable input falls back + to the raw string: a bad blob still gets a distinct identity, and + ``build_engine`` is where it earns its error message. + """ + try: + info = json.loads(raw) + except json.JSONDecodeError: + return raw + if not isinstance(info, dict) or not info.get("refresh_token"): + return raw + durable = {k: v for k, v in info.items() if k not in _ROTATING_OAUTH_FIELDS} + return json.dumps(durable, sort_keys=True, default=str) + + # --------------------------------------------------------------------------- # BigqueryDialect — Tier 1 (has logic, not just scalar config) # --------------------------------------------------------------------------- @@ -158,32 +210,128 @@ def build_engine( *, connection_string: str, ) -> "sa.Engine | None": - """Construct the SQLAlchemy engine with inline service-account JSON - when ``DatasourceConfig.credentials_json`` is set. - - ``sqlalchemy-bigquery`` accepts a ``credentials_info`` kwarg on - ``create_engine`` — a dict matching the service-account key file's - shape. Parse the JSON string into a dict and pass it through; the - BigQuery client builds credentials from it directly, no temp file - needed. When ``credentials_json`` is unset, return ``None`` so - ``engine_factory`` falls back to the default ``create_engine`` and - the BigQuery client picks up Application Default Credentials. + """Construct the SQLAlchemy engine for whichever auth path the + datasource configures. + + Three paths, in precedence order: + + 1. ``oauth_credentials_json`` — a per-end-user OAuth grant. Queries + run as that user against their own BigQuery permissions. See + :meth:`_build_oauth_engine`. + 2. ``credentials_json`` — inline service-account key. ``sqlalchemy + -bigquery`` accepts a ``credentials_info`` kwarg on + ``create_engine``, a dict matching the key file's shape, so the + BigQuery client builds credentials from it directly with no temp + file. One shared identity for every caller. + 3. Neither — return ``None`` so ``engine_factory`` falls back to the + default ``create_engine`` and the BigQuery client picks up + Application Default Credentials. + + Setting both (1) and (2) is a configuration error rather than a + silent precedence win: the two mean different identities, and + guessing which one the caller meant is how a per-user query + quietly runs as the shared service account. """ - credentials_json = datasource.credentials_json - if not credentials_json: - return None - try: - credentials_info = json.loads(credentials_json) - except json.JSONDecodeError as exc: + if datasource.oauth_credentials_json and datasource.credentials_json: raise ValueError( - f"Datasource '{datasource.name}': credentials_json is not valid JSON: {exc}" - ) from exc - if not isinstance(credentials_info, dict): + f"Datasource '{datasource.name}': credentials_json and " + f"oauth_credentials_json are mutually exclusive — they select " + f"different identities (shared service account vs. end user). " + f"Set exactly one." + ) + if datasource.oauth_credentials_json: + return self._build_oauth_engine( + datasource=datasource, connection_string=connection_string, + ) + if not datasource.credentials_json: + return None + credentials_info = _parse_credentials_object( + raw=datasource.credentials_json, + field="credentials_json", + datasource_name=datasource.name, + ) + if credentials_info.get("type") == _AUTHORIZED_USER_TYPE: raise ValueError( - f"Datasource '{datasource.name}': credentials_json must be a JSON object" + f"Datasource '{datasource.name}': credentials_json holds an " + f"'{_AUTHORIZED_USER_TYPE}' OAuth grant, but it only accepts a " + f"service-account key. Put OAuth grants in " + f"oauth_credentials_json instead." ) return sa.create_engine( connection_string, credentials_info=credentials_info, pool_pre_ping=True, ) + + def _build_oauth_engine( + self, + *, + datasource: "DatasourceConfig", + connection_string: str, + ) -> "sa.Engine": + """Build an engine bound to a caller-supplied OAuth user grant. + + ``sqlalchemy-bigquery`` has no kwarg for OAuth credentials — its + ``credentials_info``/``credentials_path``/``credentials_base64`` + kwargs all route to ``service_account.Credentials``. The supported + escape hatch is its ``user_supplied_client`` URL flag: with it set, + ``create_connect_args`` skips building a client of its own and takes + ours from ``connect_args={"client": ...}``. Without the flag the + driver would first construct an Application-Default-Credentials + client (failing outright where no ADC exists) before ours displaced + it, so the flag is load-bearing, not decorative. + + The BigQuery project is not part of an OAuth grant the way it is + part of a service-account key, so it has to come from the + connection string's host (``bigquery:///``). + """ + from google.cloud import bigquery # noqa: PLC0415 (optional 'bigquery' extra) + from google.oauth2.credentials import Credentials # noqa: PLC0415 + + info = _parse_credentials_object( + raw=datasource.oauth_credentials_json, + field="oauth_credentials_json", + datasource_name=datasource.name, + ) + url = sa.engine.make_url(connection_string) + project = url.host or info.get("quota_project_id") + if not project: + raise ValueError( + f"Datasource '{datasource.name}': OAuth credentials carry no " + f"project, so the BigQuery project must be given in the " + f"connection string as 'bigquery:///'." + ) + try: + credentials = Credentials.from_authorized_user_info(info) + except ValueError as exc: + raise ValueError( + f"Datasource '{datasource.name}': oauth_credentials_json is not " + f"a usable authorized-user grant: {exc}" + ) from exc + return sa.create_engine( + url.update_query_dict({"user_supplied_client": "true"}), + connect_args={"client": bigquery.Client( + project=project, credentials=credentials, + )}, + pool_pre_ping=True, + ) + + def credential_fingerprint(self, datasource: "DatasourceConfig") -> str: + """Identity of both BigQuery auth paths, so cached engines never + cross between a service account and an end user, or between two + end users. + + The OAuth half deliberately digests the *durable* grant rather than + the raw JSON: an access token rotates, and keying on it would mint + (and leak) a fresh engine on every refresh. Dropping the rotating + fields is only safe while a refresh token pins the identity — + without one the access token IS the whole identity, and removing it + would let two different users share one engine. + """ + material = [datasource.credentials_json or ""] + raw_oauth = datasource.oauth_credentials_json + if raw_oauth: + material.append(_durable_oauth_material(raw_oauth)) + if not any(material): + return "" + return _digest("\x00".join(material)) diff --git a/slayer/sql/engine_factory.py b/slayer/sql/engine_factory.py index 5cd46904..63391ce0 100644 --- a/slayer/sql/engine_factory.py +++ b/slayer/sql/engine_factory.py @@ -22,6 +22,8 @@ from __future__ import annotations import logging +import os +from collections import OrderedDict import sqlalchemy as sa import sqlalchemy.event as sa_event @@ -33,8 +35,93 @@ logger = logging.getLogger(__name__) -# Engine cache. Key = (connection_string, runtime_fingerprint). -_engine_cache: dict[tuple[str, str], sa.Engine] = {} +#: Identity of a cached engine: URL + runtime fields + credentials. Exported so +#: callers that key their own caches by engine identity stay in lockstep with +#: this module instead of hard-coding the arity. +EngineCacheKey = tuple[str, str, str] + +# Engine cache. Key = (connection_string, runtime_fingerprint, credential_fingerprint). +# The credential leg is load-bearing for security: a dialect whose secret is not +# in the URL (BigQuery's service-account / OAuth credentials) would otherwise +# have two differently-authenticated callers share one engine, silently running +# one identity's queries under the other's credentials. +# +# Ordered least- to most-recently-used, and bounded: once credentials are part +# of the key, cardinality follows the number of distinct *identities*, not +# datasources. A deployment handing SLayer per-end-user credentials would +# otherwise accumulate one engine — and one connection pool — per user who ever +# ran a query, for the life of the process. +_engine_cache: "OrderedDict[EngineCacheKey, sa.Engine]" = OrderedDict() + +#: Cap on simultaneously cached engines. Sized for "every datasource, times a +#: working set of active users" rather than for a whole user base; a miss costs +#: one engine construction, so over-eviction degrades latency, not correctness. +DEFAULT_MAX_CACHED_ENGINES = 64 + +#: Env override for :data:`DEFAULT_MAX_CACHED_ENGINES`. ``0`` disables caching. +MAX_CACHED_ENGINES_ENV = "SLAYER_MAX_CACHED_ENGINES" + + +def _max_cached_engines() -> int: + """Resolve the cache cap, falling back to the default on junk input. + + Read per call rather than at import so tests and hosts can retune it + without reloading the module. + """ + raw = os.environ.get(MAX_CACHED_ENGINES_ENV) + if raw is None: + return DEFAULT_MAX_CACHED_ENGINES + try: + value = int(raw) + except ValueError: + logger.warning( + "%s=%r is not an integer; using default %d.", + MAX_CACHED_ENGINES_ENV, raw, DEFAULT_MAX_CACHED_ENGINES, + ) + return DEFAULT_MAX_CACHED_ENGINES + if value < 0: + logger.warning( + "%s=%d is negative; using default %d.", + MAX_CACHED_ENGINES_ENV, value, DEFAULT_MAX_CACHED_ENGINES, + ) + return DEFAULT_MAX_CACHED_ENGINES + return value + + +def _dispose_quietly(*, engine: sa.Engine, reason: str) -> None: + """Release an engine's pooled connections, logging rather than raising. + + ``Engine.dispose()`` closes checked-in connections and swaps in a fresh + pool; connections checked out by an in-flight query are detached and + closed when returned. So this is safe to call on an engine another caller + still holds a reference to — it reclaims sockets without breaking them. + """ + try: + engine.dispose() + except Exception: + logger.warning("Failed to dispose engine (%s).", reason, exc_info=True) + + +def _evict_to_limit() -> None: + """Drop least-recently-used engines until the cache fits its cap.""" + limit = _max_cached_engines() + while len(_engine_cache) > limit: + _, evicted = _engine_cache.popitem(last=False) + _dispose_quietly(engine=evicted, reason="evicted from engine cache") + + +def _cache_key(datasource: DatasourceConfig, connection_string: str) -> EngineCacheKey: + """Cache identity for ``datasource``: URL + runtime fields + credentials. + + Kept in one place because ``query_engine._sql_client_cache_key`` must agree + with it; a divergence between the two caches means a caller can get a client + whose engine was built for different credentials. + """ + return ( + connection_string, + _runtime_fingerprint(datasource), + dialect_for_ds_type(datasource.type).credential_fingerprint(datasource), + ) def _runtime_fingerprint(datasource: DatasourceConfig) -> str: @@ -147,15 +234,50 @@ def get_engine(datasource: DatasourceConfig) -> sa.Engine: apply the wrong USE statements. """ connection_string = datasource.get_connection_string() - cache_key = (connection_string, _runtime_fingerprint(datasource)) - if cache_key not in _engine_cache: - _engine_cache[cache_key] = _build_engine( - datasource=datasource, connection_string=connection_string, - ) - return _engine_cache[cache_key] + cache_key = _cache_key(datasource, connection_string) + cached = _engine_cache.get(cache_key) + if cached is not None: + _engine_cache.move_to_end(cache_key) + return cached + engine = _build_engine( + datasource=datasource, connection_string=connection_string, + ) + _engine_cache[cache_key] = engine + _evict_to_limit() + return engine + +def invalidate_engine(datasource: DatasourceConfig) -> bool: + """Drop and dispose the cached engine for ``datasource``. Returns whether + one was cached. -def reset_cache() -> None: - """Discard every cached engine. Used by tests that need fresh pools; - not called by production code.""" + Meant for the case where the *credentials* an engine was built with have + stopped working — a revoked OAuth grant, a rotated service-account key. + Those engines are poisoned for good: the credential object is baked in at + construction, so every retry through the cache fails identically until + someone throws the engine away. Retrying a transient network blip, by + contrast, wants the pool kept. + + ``get_engine`` rebuilds on the next call, picking up whatever credentials + the datasource now carries. + """ + connection_string = datasource.get_connection_string() + evicted = _engine_cache.pop(_cache_key(datasource, connection_string), None) + if evicted is None: + return False + _dispose_quietly(engine=evicted, reason=f"credentials rejected for '{datasource.name}'") + return True + + +def reset_cache(*, dispose: bool = False) -> None: + """Discard every cached engine. + + ``dispose`` defaults to False to preserve the long-standing test-fixture + behaviour of dropping references without touching pools. Hosts tearing a + process down want ``dispose=True`` so server-side connections close + promptly instead of waiting on garbage collection. + """ + if dispose: + for key, engine in list(_engine_cache.items()): + _dispose_quietly(engine=engine, reason=f"cache reset ({key[0]})") _engine_cache.clear() diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py index fade31de..aba69c9d 100644 --- a/tests/dialects/test_bigquery.py +++ b/tests/dialects/test_bigquery.py @@ -12,6 +12,7 @@ from __future__ import annotations +import json import re import tempfile from unittest.mock import patch @@ -24,6 +25,7 @@ from slayer.engine.enriched import EnrichedQuery from slayer.engine.enrichment import enrich_query from slayer.engine.query_engine import SlayerQueryEngine +from slayer.engine.query_engine import _sql_client_cache_key from slayer.sql.dialects import ( BigqueryDialect, PostgresDialect, @@ -473,7 +475,7 @@ async def test_engine_dispatches_through_decode_result_keys_hook() -> None: ) await storage.save_model(model) engine = SlayerQueryEngine(storage=storage) - engine._sql_clients[(ds.get_connection_string(), "")] = _FakeBigQueryClient( + engine._sql_clients[_sql_client_cache_key(ds)] = _FakeBigQueryClient( rows=[{"orders.status": "paid"}] ) with patch.object( @@ -539,7 +541,7 @@ async def _build_bigquery_engine(rows: list[dict]) -> tuple[SlayerQueryEngine, t ) await storage.save_model(model) engine = SlayerQueryEngine(storage=storage) - engine._sql_clients[(ds.get_connection_string(), "")] = _FakeBigQueryClient(rows) + engine._sql_clients[_sql_client_cache_key(ds)] = _FakeBigQueryClient(rows) return engine, tmp, ds @@ -664,3 +666,189 @@ def test_build_engine_with_non_object_credentials_json_raises(payload: str) -> N dialect = BigqueryDialect() with pytest.raises(ValueError, match="credentials_json must be a JSON object"): dialect.build_engine(ds, connection_string="bigquery://my-project") + + +# --------------------------------------------------------------------------- +# build_engine — per-end-user OAuth grant +# +# ``sqlalchemy-bigquery`` routes every credentials kwarg it has to +# ``service_account.Credentials``, so OAuth user grants have to go through +# its ``user_supplied_client`` escape hatch instead. These tests pin that +# wiring and the mutual exclusion with the service-account path. +# --------------------------------------------------------------------------- + + +def _oauth_info(**overrides) -> dict: + info = { + "type": "authorized_user", + "client_id": "cid.apps.googleusercontent.com", + "client_secret": "csecret", + "refresh_token": "rtok-alice", + "token": "access-token-1", + "token_uri": "https://oauth2.googleapis.com/token", + } + info.update(overrides) + return info + + +def _oauth_ds(name: str = "bq", **overrides) -> DatasourceConfig: + return DatasourceConfig( + name=name, + type="bigquery", + oauth_credentials_json=json.dumps(_oauth_info(**overrides)), + ) + + +def test_build_engine_oauth_uses_user_supplied_client() -> None: + """OAuth path must set the ``user_supplied_client`` URL flag AND pass the + client via ``connect_args``. Without the flag, ``sqlalchemy-bigquery`` + builds an ADC client first (and raises where no ADC exists).""" + dialect = BigqueryDialect() + captured: dict = {} + + def fake_create_engine(url, **kwargs): + captured["url"] = url + captured["kwargs"] = kwargs + return object() + + fake_client = object() + with ( + patch("slayer.sql.dialects.bigquery.sa.create_engine", side_effect=fake_create_engine), + patch("google.cloud.bigquery.Client", return_value=fake_client) as mk_client, + patch("google.oauth2.credentials.Credentials.from_authorized_user_info") as mk_creds, + ): + engine = dialect.build_engine( + _oauth_ds(), connection_string="bigquery://my-project/my_dataset", + ) + + assert engine is not None + assert captured["url"].query["user_supplied_client"] == "true" + assert captured["kwargs"]["connect_args"] == {"client": fake_client} + assert captured["kwargs"]["pool_pre_ping"] is True + # No ``credentials_info`` — that kwarg would send us back through + # ``service_account.Credentials`` and defeat the whole path. + assert "credentials_info" not in captured["kwargs"] + assert mk_client.call_args.kwargs["project"] == "my-project" + assert mk_client.call_args.kwargs["credentials"] is mk_creds.return_value + assert mk_creds.call_args.args[0] == _oauth_info() + + +def test_build_engine_oauth_without_project_raises() -> None: + """An OAuth grant carries no project, so omitting it from the connection + string is a config error rather than a confusing downstream 404.""" + dialect = BigqueryDialect() + with pytest.raises(ValueError, match="must be given in the connection string"): + dialect.build_engine(_oauth_ds(), connection_string="bigquery://") + + +def test_build_engine_oauth_falls_back_to_quota_project() -> None: + """A grant carrying ``quota_project_id`` supplies the project when the + connection string doesn't.""" + dialect = BigqueryDialect() + with ( + patch("slayer.sql.dialects.bigquery.sa.create_engine", return_value=object()), + patch("google.cloud.bigquery.Client", return_value=object()) as mk_client, + patch("google.oauth2.credentials.Credentials.from_authorized_user_info"), + ): + dialect.build_engine( + _oauth_ds(quota_project_id="quota-proj"), connection_string="bigquery://", + ) + assert mk_client.call_args.kwargs["project"] == "quota-proj" + + +def test_build_engine_rejects_both_credential_kinds() -> None: + """Setting both is a config error, not a silent precedence win: guessing + is how a per-user query quietly runs as the shared service account.""" + ds = DatasourceConfig( + name="bq", + type="bigquery", + credentials_json=json.dumps({"type": "service_account"}), + oauth_credentials_json=json.dumps(_oauth_info()), + ) + with pytest.raises(ValueError, match="mutually exclusive"): + BigqueryDialect().build_engine(ds, connection_string="bigquery://p/d") + + +def test_build_engine_rejects_oauth_grant_in_credentials_json() -> None: + """An authorized-user grant in ``credentials_json`` cannot work — the + driver hands it to ``from_service_account_info``. Say so up front.""" + ds = DatasourceConfig( + name="bq", type="bigquery", credentials_json=json.dumps(_oauth_info()), + ) + with pytest.raises(ValueError, match="Put OAuth grants in oauth_credentials_json"): + BigqueryDialect().build_engine(ds, connection_string="bigquery://p/d") + + +@pytest.mark.parametrize( + argnames="payload,message", + argvalues=[ + ("not json at all", "oauth_credentials_json is not valid JSON"), + ("[]", "oauth_credentials_json must be a JSON object"), + ], +) +def test_build_engine_oauth_malformed_raises(payload: str, message: str) -> None: + ds = DatasourceConfig(name="bq", type="bigquery", oauth_credentials_json=payload) + with pytest.raises(ValueError, match=message): + BigqueryDialect().build_engine(ds, connection_string="bigquery://p/d") + + +# --------------------------------------------------------------------------- +# credential_fingerprint — cached engines must not cross identities +# --------------------------------------------------------------------------- + + +def test_credential_fingerprint_empty_without_credentials() -> None: + """ADC datasources keep the empty fingerprint, so their cache key shape + is unchanged.""" + ds = DatasourceConfig(name="bq", type="bigquery", database="p") + assert BigqueryDialect().credential_fingerprint(ds) == "" + + +def test_credential_fingerprint_differs_between_oauth_users() -> None: + """Two end users on the same project must never share a cached engine.""" + dialect = BigqueryDialect() + alice = dialect.credential_fingerprint(_oauth_ds(refresh_token="rtok-alice")) + bob = dialect.credential_fingerprint(_oauth_ds(refresh_token="rtok-bob")) + assert alice != bob + assert alice and bob + + +def test_credential_fingerprint_differs_between_oauth_and_service_account() -> None: + dialect = BigqueryDialect() + oauth = dialect.credential_fingerprint(_oauth_ds()) + svc = dialect.credential_fingerprint(DatasourceConfig( + name="bq", type="bigquery", + credentials_json=json.dumps({"type": "service_account", "project_id": "p"}), + )) + assert oauth != svc + + +def test_credential_fingerprint_stable_across_token_refresh() -> None: + """A refreshed access token is the same user. Keying on it would mint — + and leak — a fresh engine on every refresh.""" + dialect = BigqueryDialect() + before = dialect.credential_fingerprint(_oauth_ds(token="access-1", expiry="2026-01-01")) + after = dialect.credential_fingerprint(_oauth_ds(token="access-2", expiry="2026-01-02")) + assert before == after + + +def test_credential_fingerprint_keeps_token_when_no_refresh_token() -> None: + """Without a refresh token the access token IS the whole identity, so it + must stay in the digest or two users collide on one engine.""" + dialect = BigqueryDialect() + info = _oauth_info() + info.pop("refresh_token") + def ds_for(token: str) -> DatasourceConfig: + return DatasourceConfig( + name="bq", type="bigquery", + oauth_credentials_json=json.dumps({**info, "token": token}), + ) + assert dialect.credential_fingerprint(ds_for("tok-alice")) != dialect.credential_fingerprint(ds_for("tok-bob")) + + +def test_credential_fingerprint_leaks_no_secret_material() -> None: + """The fingerprint lands in an in-memory cache key and log lines; it must + not be reversible to the grant.""" + fp = BigqueryDialect().credential_fingerprint(_oauth_ds()) + for secret in ("rtok-alice", "csecret", "access-token-1"): + assert secret not in fp diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py index 5622f2e0..ba0bf293 100644 --- a/tests/dialects/test_tsql.py +++ b/tests/dialects/test_tsql.py @@ -29,6 +29,7 @@ from slayer.core.query import ColumnRef, OrderItem, SlayerQuery, TimeDimension from slayer.engine.enrichment import enrich_query from slayer.engine.query_engine import SlayerQueryEngine +from slayer.engine.query_engine import _sql_client_cache_key from slayer.sql.generator import SQLGenerator from slayer.sql.dialects.tsql import TsqlDialect from slayer.storage.yaml_storage import YAMLStorage @@ -702,7 +703,7 @@ async def _build_tsql_engine(rows: list[dict]) -> tuple[SlayerQueryEngine, tempf ) await storage.save_model(model) engine = SlayerQueryEngine(storage=storage) - engine._sql_clients[(ds.get_connection_string(), "")] = _FakeTsqlClient(rows) + engine._sql_clients[_sql_client_cache_key(ds)] = _FakeTsqlClient(rows) return engine, tmp, ds diff --git a/tests/integration/test_in_memory_sqlite.py b/tests/integration/test_in_memory_sqlite.py index 551bcdb4..53bb1098 100644 --- a/tests/integration/test_in_memory_sqlite.py +++ b/tests/integration/test_in_memory_sqlite.py @@ -85,9 +85,9 @@ async def test_file_backed_sqlite_engine_cached_in_module(tmp_path: Path) -> Non engine_factory.reset_cache() client = SlayerSQLClient(datasource=ds) await client.execute("SELECT 1") - # Engine is cached in the factory under the (connection_string, '') - # key (empty runtime fingerprint for non-snowflake datasources). - assert (conn_str, "") in engine_factory._engine_cache + # Engine is cached in the factory under the datasource's engine identity + # (connection_string plus empty runtime/credential fingerprints here). + assert engine_factory._cache_key(ds, conn_str) in engine_factory._engine_cache # Per-client engine now holds the factory-cached instance directly, # so a second client on the same datasource reuses the same engine. client2 = SlayerSQLClient(datasource=ds) diff --git a/tests/test_engine_factory.py b/tests/test_engine_factory.py index 27b7528e..38fecc94 100644 --- a/tests/test_engine_factory.py +++ b/tests/test_engine_factory.py @@ -228,3 +228,159 @@ def test_sql_client_uses_engine_factory_for_engine_creation(self) -> None: from slayer.sql import client as sql_client source = open(sql_client.__file__).read() assert "engine_factory" in source + + +class TestCredentialKeying: + """The cache key's credential leg. Without it, two callers whose only + difference is *who they authenticate as* share one engine — and one + silently runs the other's queries under the wrong identity.""" + + @staticmethod + def _bq(name: str, credentials_json: str | None) -> DatasourceConfig: + return DatasourceConfig( + name=name, type="bigquery", + connection_string="bigquery://proj/dset", + credentials_json=credentials_json, + ) + + def test_same_url_different_credentials_get_different_keys(self) -> None: + alice = self._bq("bq", '{"type": "service_account", "client_email": "alice@x"}') + bob = self._bq("bq", '{"type": "service_account", "client_email": "bob@x"}') + conn = "bigquery://proj/dset" + assert engine_factory._cache_key(alice, conn) != engine_factory._cache_key(bob, conn) + + def test_key_carries_no_secret_material(self) -> None: + secret = "super-secret-private-key" + ds = self._bq("bq", '{"type": "service_account", "private_key": "%s"}' % secret) + key = engine_factory._cache_key(ds, "bigquery://proj/dset") + assert secret not in "".join(key) + + def test_credential_free_datasource_keeps_empty_leg(self) -> None: + ds = DatasourceConfig(name="pg", type="postgres", host="h", database="db") + key = engine_factory._cache_key(ds, ds.get_connection_string()) + assert key[2] == "" + + def test_query_engine_agrees_with_factory(self) -> None: + """The two caches must key identically, or a caller can be handed a + client whose engine was built for someone else's credentials.""" + from slayer.engine.query_engine import _sql_client_cache_key + ds = self._bq("bq", '{"type": "service_account", "client_email": "alice@x"}') + assert _sql_client_cache_key(ds) == engine_factory._cache_key( + ds, ds.get_connection_string(), + ) + + +class TestCacheBounding: + """Per-identity keys make cache cardinality track *users*, not + datasources, so the cache has to be bounded and evictions must actually + release the pooled connections.""" + + @staticmethod + def _lite(n: int) -> DatasourceConfig: + return DatasourceConfig(name=f"lite{n}", type="sqlite", database=f"/tmp/slayer-cache-{n}.db") + + def test_cache_evicts_least_recently_used_over_limit(self, monkeypatch) -> None: + engine_factory.reset_cache() + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "2") + first, second, third = (self._lite(i) for i in range(3)) + engine_factory.get_engine(first) + engine_factory.get_engine(second) + engine_factory.get_engine(third) + assert len(engine_factory._engine_cache) == 2 + cached = list(engine_factory._engine_cache) + assert engine_factory._cache_key(first, first.get_connection_string()) not in cached + engine_factory.reset_cache() + + def test_reuse_refreshes_recency(self, monkeypatch) -> None: + """A hit must move the entry to the MRU end, otherwise the cap + degenerates into FIFO and evicts the hottest engine.""" + engine_factory.reset_cache() + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "2") + first, second, third = (self._lite(i) for i in range(3)) + engine_factory.get_engine(first) + engine_factory.get_engine(second) + engine_factory.get_engine(first) # first is now most-recently used + engine_factory.get_engine(third) + remaining = list(engine_factory._engine_cache) + assert engine_factory._cache_key(first, first.get_connection_string()) in remaining + assert engine_factory._cache_key(second, second.get_connection_string()) not in remaining + engine_factory.reset_cache() + + def test_eviction_disposes_the_engine(self, monkeypatch) -> None: + engine_factory.reset_cache() + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "1") + first, second = self._lite(0), self._lite(1) + evicted = engine_factory.get_engine(first) + with patch.object(evicted, "dispose") as disposed: + engine_factory.get_engine(second) + disposed.assert_called_once() + engine_factory.reset_cache() + + def test_dispose_failure_does_not_break_caching(self, monkeypatch) -> None: + """A pool that refuses to close must not take the whole factory with + it — the new engine still has to reach the caller.""" + engine_factory.reset_cache() + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "1") + first, second = self._lite(0), self._lite(1) + stuck = engine_factory.get_engine(first) + with patch.object(stuck, "dispose", side_effect=RuntimeError("pool stuck")): + assert isinstance(engine_factory.get_engine(second), sa.Engine) + engine_factory.reset_cache() + + @pytest.mark.parametrize(argnames="raw", argvalues=["nonsense", "-1", ""]) + def test_bad_limit_env_falls_back_to_default(self, monkeypatch, raw: str) -> None: + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, raw) + assert engine_factory._max_cached_engines() == engine_factory.DEFAULT_MAX_CACHED_ENGINES + + def test_zero_limit_disables_caching(self, monkeypatch) -> None: + engine_factory.reset_cache() + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "0") + engine_factory.get_engine(self._lite(0)) + assert len(engine_factory._engine_cache) == 0 + + +class TestInvalidateEngine: + """Credentials baked into an engine can be revoked out from under it. + Those engines are poisoned permanently, so retrying through the cache + reproduces the failure forever unless something evicts them.""" + + @staticmethod + def _lite() -> DatasourceConfig: + return DatasourceConfig(name="lite", type="sqlite", database="/tmp/slayer-invalidate.db") + + def test_invalidate_removes_and_disposes(self) -> None: + engine_factory.reset_cache() + ds = self._lite() + engine = engine_factory.get_engine(ds) + with patch.object(engine, "dispose") as disposed: + assert engine_factory.invalidate_engine(ds) is True + disposed.assert_called_once() + assert engine_factory._cache_key(ds, ds.get_connection_string()) not in engine_factory._engine_cache + + def test_invalidate_is_a_noop_when_uncached(self) -> None: + engine_factory.reset_cache() + assert engine_factory.invalidate_engine(self._lite()) is False + + def test_next_get_engine_rebuilds(self) -> None: + engine_factory.reset_cache() + ds = self._lite() + before = engine_factory.get_engine(ds) + engine_factory.invalidate_engine(ds) + assert engine_factory.get_engine(ds) is not before + engine_factory.reset_cache() + + +class TestResetCacheDisposal: + + def test_reset_disposes_only_when_asked(self) -> None: + engine_factory.reset_cache() + ds = DatasourceConfig(name="lite", type="sqlite", database="/tmp/slayer-reset.db") + engine = engine_factory.get_engine(ds) + with patch.object(engine, "dispose") as disposed: + engine_factory.reset_cache() + disposed.assert_not_called() + + engine = engine_factory.get_engine(ds) + with patch.object(engine, "dispose") as disposed: + engine_factory.reset_cache(dispose=True) + disposed.assert_called_once() diff --git a/tests/test_query_cache.py b/tests/test_query_cache.py index a07a479c..2bf56a08 100644 --- a/tests/test_query_cache.py +++ b/tests/test_query_cache.py @@ -63,7 +63,7 @@ def _make_entry(*, created_at: float = 0.0, sql: str = "SELECT 1", response=None sql=sql, ds_fingerprint="fp", dialect="sqlite", - ds_key=("conn", "rt"), + ds_key=("conn", "rt", "cred"), created_at=created_at, ) diff --git a/tests/test_sql_client.py b/tests/test_sql_client.py index a903f567..ec7dd7b2 100644 --- a/tests/test_sql_client.py +++ b/tests/test_sql_client.py @@ -2,16 +2,19 @@ import logging import sqlite3 +from unittest.mock import patch import pytest import sqlalchemy.exc +from slayer.core.models import DatasourceConfig from slayer.sql import client as sql_client from slayer.sql.client import ( _build_type_probe_sql, _execute_with_retry_async, _execute_with_retry_sync, _execute_with_retry_threaded, + _is_auth_failure, _is_transient_db_error, _map_type_code, ) @@ -444,3 +447,95 @@ def test_tsql_alias_uses_top_0(self) -> None: def test_none_db_type_uses_limit(self) -> None: sql = _build_type_probe_sql(self.BASE, db_type=None) assert "LIMIT 0" in sql + + +class TestIsAuthFailure: + """Credential rejection is classified separately from transient errors: + retrying it is pointless (the credentials are baked into the engine), and + the right response is to throw the engine away.""" + + def test_oauth_invalid_grant_is_auth_failure(self) -> None: + assert _is_auth_failure(Exception("('invalid_grant: Token has been expired or revoked.')")) + + def test_libpq_password_failure_is_auth_failure(self) -> None: + assert _is_auth_failure(Exception('FATAL: password authentication failed for user "svc"')) + + def test_signal_found_through_sqlalchemy_orig(self) -> None: + """Drivers surface wrapped; the signal is a layer or two down.""" + inner = Exception("invalid_grant") + wrapped = sqlalchemy.exc.OperationalError("SELECT 1", {}, inner) + assert _is_auth_failure(wrapped) + + def test_signal_found_through_cause_chain(self) -> None: + inner = Exception("Reauthentication is needed") + outer = RuntimeError("query failed") + outer.__cause__ = inner + assert _is_auth_failure(outer) + + def test_google_refresh_error_matched_by_type_name(self) -> None: + """google-auth ships only with the optional 'bigquery' extra, so the + classifier matches on class name rather than importing it.""" + class RefreshError(Exception): + pass + assert _is_auth_failure(RefreshError("bad news")) + + def test_transient_errors_are_not_auth_failures(self) -> None: + for message in ("database is locked", "deadlock detected", "server closed the connection"): + assert not _is_auth_failure(Exception(message)), message + + def test_table_permission_denied_is_not_an_auth_failure(self) -> None: + """The credentials worked; the grant didn't. Evicting a healthy engine + over this is pure pool churn.""" + assert not _is_auth_failure(Exception("permission denied for table orders")) + + def test_cyclic_cause_chain_terminates(self) -> None: + first, second = Exception("a"), Exception("b") + first.__cause__ = second + second.__cause__ = first + assert _is_auth_failure(first) is False + + +class TestClientDiscardsEngineOnAuthFailure: + + @staticmethod + def _client(): + return sql_client.SlayerSQLClient( + datasource=DatasourceConfig(name="bq", type="bigquery", connection_string="bigquery://p/d"), + ) + + async def test_auth_failure_invalidates_cached_engine(self) -> None: + client = self._client() + client._sync_engine = object() + boom = Exception("invalid_grant: Token has been expired or revoked.") + with ( + patch.object(sql_client.SlayerSQLClient, "_execute", side_effect=boom), + patch("slayer.sql.engine_factory.invalidate_engine") as invalidate, + ): + with pytest.raises(Exception, match="invalid_grant"): + await client.execute("SELECT 1") + invalidate.assert_called_once_with(client.datasource) + assert client._sync_engine is None + + async def test_non_auth_failure_keeps_the_engine(self) -> None: + client = self._client() + engine = object() + client._sync_engine = engine + with ( + patch.object(sql_client.SlayerSQLClient, "_execute", side_effect=Exception("no such table: orders")), + patch("slayer.sql.engine_factory.invalidate_engine") as invalidate, + ): + with pytest.raises(Exception, match="no such table"): + await client.execute("SELECT 1") + invalidate.assert_not_called() + assert client._sync_engine is engine + + async def test_cleanup_failure_does_not_mask_the_original_error(self) -> None: + """The auth error is what the caller needs to see; a failed eviction + must not displace it.""" + client = self._client() + with ( + patch.object(sql_client.SlayerSQLClient, "_execute", side_effect=Exception("invalid_grant")), + patch("slayer.sql.engine_factory.invalidate_engine", side_effect=RuntimeError("cache busted")), + ): + with pytest.raises(Exception, match="invalid_grant"): + await client.execute("SELECT 1") diff --git a/tests/test_sql_generator.py b/tests/test_sql_generator.py index fa83d15a..1498ea98 100644 --- a/tests/test_sql_generator.py +++ b/tests/test_sql_generator.py @@ -5422,6 +5422,9 @@ async def test_expression_measure_sql_not_corrupted(self) -> None: mock_ds = MagicMock() mock_ds.get_connection_string.return_value = "sqlite://" mock_ds.type = "sqlite" + # Real attribute, not an auto-Mock: it feeds the cache key's + # credential digest, which hashes it. + mock_ds.credentials_json = None with patch.object(engine, "_resolve_datasource", new_callable=AsyncMock, return_value=mock_ds): captured_sql = [] @@ -5431,7 +5434,7 @@ async def capture_sql(sql): mock_client = MagicMock() mock_client.get_column_types = capture_sql - engine._sql_clients[("sqlite://", "")] = mock_client + engine._sql_clients[("sqlite://", "", "")] = mock_client await engine.get_column_types("orders") @@ -5480,6 +5483,9 @@ async def test_cross_model_measures_probed_via_engine(self) -> None: mock_ds = MagicMock() mock_ds.get_connection_string.return_value = "sqlite://" mock_ds.type = "sqlite" + # Real attribute, not an auto-Mock: it feeds the cache key's + # credential digest, which hashes it. + mock_ds.credentials_json = None with patch.object(engine, "_resolve_datasource", new_callable=AsyncMock, return_value=mock_ds), \ patch.object(engine, "_enrich", new_callable=AsyncMock, return_value=mock_enriched): @@ -5489,7 +5495,7 @@ async def capture_types(sql): mock_client = MagicMock() mock_client.get_column_types = capture_types - engine._sql_clients[("sqlite://", "")] = mock_client + engine._sql_clients[("sqlite://", "", "")] = mock_client result = await engine.get_column_types("orders") From e181daf037b8fc3ccd957c8fa569896a3f7a0e70 Mon Sep 17 00:00:00 2001 From: AivanF Date: Thu, 6 Aug 2026 14:54:03 +0300 Subject: [PATCH 2/5] Fixes from review --- docs/configuration/datasources.md | 2 +- slayer/engine/query_engine.py | 2 +- slayer/sql/client.py | 53 ++++-- slayer/sql/dialects/bigquery.py | 20 ++- slayer/sql/engine_factory.py | 101 ++++++++++-- tests/dialects/test_bigquery.py | 45 ++++- tests/test_engine_factory.py | 264 +++++++++++++++++++++++++++++- tests/test_sql_client.py | 63 ++++++- 8 files changed, 507 insertions(+), 43 deletions(-) diff --git a/docs/configuration/datasources.md b/docs/configuration/datasources.md index a716ab87..089a7dba 100644 --- a/docs/configuration/datasources.md +++ b/docs/configuration/datasources.md @@ -171,7 +171,7 @@ Statement-level timeout is enforced via Three ways to authenticate, in the order SLayer prefers them: -1. **`oauth_credentials_json`** — an OAuth *authorized user* grant, the shape `google.oauth2.credentials.Credentials.from_authorized_user_info` accepts (`token`, `refresh_token`, `client_id`, `client_secret`, `token_uri`). Queries run as that end user, against that user's own BigQuery permissions. The grant carries no project, so the connection string must name one: `bigquery:///`. +1. **`oauth_credentials_json`** — an OAuth *authorized user* grant, the shape `google.oauth2.credentials.Credentials.from_authorized_user_info` accepts (`token`, `refresh_token`, `client_id`, `client_secret`, `token_uri`). Queries run as that end user, against that user's own BigQuery permissions. An OAuth grant carries no project of its own, so one must come from the connection string — `bigquery:///` — or from a `quota_project_id` in the grant, which SLayer falls back to when the connection string omits it. 2. **`credentials_json`** — a service-account key file's contents. One shared identity for every caller. 3. **Neither** — Application Default Credentials (`GOOGLE_APPLICATION_CREDENTIALS`, or the attached compute identity). diff --git a/slayer/engine/query_engine.py b/slayer/engine/query_engine.py index 8e48cb49..e22d7c17 100644 --- a/slayer/engine/query_engine.py +++ b/slayer/engine/query_engine.py @@ -92,7 +92,7 @@ def _sql_client_cache_key(datasource: DatasourceConfig) -> EngineCacheKey: different credentials. Sharing one implementation makes that impossible. """ return _engine_cache_key( - datasource, datasource.get_connection_string(), + datasource=datasource, connection_string=datasource.get_connection_string(), ) diff --git a/slayer/sql/client.py b/slayer/sql/client.py index eef82b07..05dc091f 100644 --- a/slayer/sql/client.py +++ b/slayer/sql/client.py @@ -528,8 +528,9 @@ def _get_sync_engine_for_client(self) -> sa.Engine | None: self._sync_engine = engine_factory.get_engine(self.datasource) return self._sync_engine - def _discard_engine_on_auth_failure(self, exc: BaseException) -> None: - """Drop the cached engine when the failure was a credential rejection. + def _discard_sync_engine_on_auth_failure(self, exc: BaseException) -> bool: + """Drop the sync engine when the failure was a credential rejection. + Returns whether ``exc`` was one. The credentials an engine authenticates with are fixed at construction (BigQuery bakes a client object into the engine), so a revoked OAuth @@ -542,7 +543,7 @@ def _discard_engine_on_auth_failure(self, exc: BaseException) -> None: to see, so a failure to clean up must not displace it. """ if not _is_auth_failure(exc): - return + return False from slayer.sql import engine_factory # noqa: PLC0415 self._sync_engine = None try: @@ -552,6 +553,18 @@ def _discard_engine_on_auth_failure(self, exc: BaseException) -> None: "Failed to invalidate engine for datasource %r after an " "authentication failure.", self.datasource.name, exc_info=True, ) + return True + + async def _discard_engines_on_auth_failure(self, exc: BaseException) -> None: + """Async-path cleanup: the sync engine, plus this client's async one. + + Native-async dialects hold a second pool that ``invalidate_engine`` + knows nothing about, and disposing it needs a loop — hence the split + from the sync variant, which ``execute_sync`` uses. + """ + if not self._discard_sync_engine_on_auth_failure(exc): + return + await self.aclose() async def execute( self, @@ -562,7 +575,7 @@ async def execute( try: return await self._execute(sql=sql, timeout_seconds=timeout_seconds) except Exception as exc: - self._discard_engine_on_auth_failure(exc) + await self._discard_engines_on_auth_failure(exc) raise async def _execute( @@ -602,6 +615,13 @@ async def get_column_types(self, sql: str) -> dict[str, str]: Returns {column_name: type_category} where type_category is "number", "string", "time", or "boolean". """ + try: + return await self._get_column_types(sql=sql) + except Exception as exc: + await self._discard_engines_on_auth_failure(exc) + raise + + async def _get_column_types(self, *, sql: str) -> dict[str, str]: async_engine = self._get_async_engine() if async_engine is not None: return await _get_column_types_async( @@ -626,14 +646,23 @@ def execute_sync( sql: str, timeout_seconds: int = 120, ) -> list[dict[str, Any]]: - """Execute SQL synchronously (for CLI, notebooks, tests).""" - return _execute_with_retry_sync( - sql=sql, - connection_string=self.datasource.get_connection_string(), - db_type=self.datasource.type, - timeout_seconds=timeout_seconds, - engine=self._get_sync_engine_for_client(), - ) + """Execute SQL synchronously (for CLI, notebooks, tests). + + Only the sync engine is discarded on a credential rejection: this + method has no loop to dispose an async pool on, and the paths that do + (``execute`` / ``get_column_types``) handle it. + """ + try: + return _execute_with_retry_sync( + sql=sql, + connection_string=self.datasource.get_connection_string(), + db_type=self.datasource.type, + timeout_seconds=timeout_seconds, + engine=self._get_sync_engine_for_client(), + ) + except Exception as exc: + self._discard_sync_engine_on_auth_failure(exc) + raise # --------------------------------------------------------------------------- diff --git a/slayer/sql/dialects/bigquery.py b/slayer/sql/dialects/bigquery.py index d8fa8460..47d84ed5 100644 --- a/slayer/sql/dialects/bigquery.py +++ b/slayer/sql/dialects/bigquery.py @@ -283,11 +283,13 @@ def _build_oauth_engine( The BigQuery project is not part of an OAuth grant the way it is part of a service-account key, so it has to come from the - connection string's host (``bigquery:///``). - """ - from google.cloud import bigquery # noqa: PLC0415 (optional 'bigquery' extra) - from google.oauth2.credentials import Credentials # noqa: PLC0415 + connection string's host (``bigquery:///``), or + from the grant's ``quota_project_id`` when it carries one. + Config is validated before the ``google.*`` imports so a + misconfigured datasource reports *that* rather than a missing + optional dependency — those ship only with the 'bigquery' extra. + """ info = _parse_credentials_object( raw=datasource.oauth_credentials_json, field="oauth_credentials_json", @@ -297,10 +299,14 @@ def _build_oauth_engine( project = url.host or info.get("quota_project_id") if not project: raise ValueError( - f"Datasource '{datasource.name}': OAuth credentials carry no " - f"project, so the BigQuery project must be given in the " - f"connection string as 'bigquery:///'." + f"Datasource '{datasource.name}': no BigQuery project resolved. " + f"An OAuth grant carries none of its own, so it must be given in " + f"the connection string as 'bigquery:///', or " + f"as 'quota_project_id' inside oauth_credentials_json." ) + from google.cloud import bigquery # noqa: PLC0415 (optional 'bigquery' extra) + from google.oauth2.credentials import Credentials # noqa: PLC0415 + try: credentials = Credentials.from_authorized_user_info(info) except ValueError as exc: diff --git a/slayer/sql/engine_factory.py b/slayer/sql/engine_factory.py index 63391ce0..815eb23a 100644 --- a/slayer/sql/engine_factory.py +++ b/slayer/sql/engine_factory.py @@ -23,6 +23,7 @@ import logging import os +import threading from collections import OrderedDict import sqlalchemy as sa @@ -53,6 +54,18 @@ # ran a query, for the life of the process. _engine_cache: "OrderedDict[EngineCacheKey, sa.Engine]" = OrderedDict() +# Guards every read, recency bump, insert, eviction and reset of the cache +# above. Engines are reached from worker threads (``_run_sync_in_thread`` in +# slayer.sql.client) as well as from the event loop, so the lookup/move_to_end +# pair is not safe unsynchronised: an ``invalidate_engine`` landing between the +# two raises KeyError, and two threads missing at once build two pools for one +# key and silently orphan the loser. +# +# Never held across engine construction or ``dispose()`` — both do real I/O, +# and serialising unrelated datasources behind them would cost more than the +# race it prevents. +_cache_lock = threading.Lock() + #: Cap on simultaneously cached engines. Sized for "every datasource, times a #: working set of active users" rather than for a whole user base; a miss costs #: one engine construction, so over-eviction degrades latency, not correctness. @@ -102,12 +115,19 @@ def _dispose_quietly(*, engine: sa.Engine, reason: str) -> None: logger.warning("Failed to dispose engine (%s).", reason, exc_info=True) -def _evict_to_limit() -> None: - """Drop least-recently-used engines until the cache fits its cap.""" +def _take_evictions_over_limit() -> list[sa.Engine]: + """Pop least-recently-used entries until the cache fits its cap. + + Returns the evicted engines rather than disposing them, because callers + run this holding ``_cache_lock`` and ``dispose()`` closes sockets. Dispose + the result after releasing the lock. + """ limit = _max_cached_engines() + evicted: list[sa.Engine] = [] while len(_engine_cache) > limit: - _, evicted = _engine_cache.popitem(last=False) - _dispose_quietly(engine=evicted, reason="evicted from engine cache") + _, engine = _engine_cache.popitem(last=False) + evicted.append(engine) + return evicted def _cache_key(datasource: DatasourceConfig, connection_string: str) -> EngineCacheKey: @@ -116,6 +136,12 @@ def _cache_key(datasource: DatasourceConfig, connection_string: str) -> EngineCa Kept in one place because ``query_engine._sql_client_cache_key`` must agree with it; a divergence between the two caches means a caller can get a client whose engine was built for different credentials. + + ``datasource`` is read, not retained. Only ``get_engine`` snapshots it + first, because only there does a slow engine build sit between the key and + the credentials it claims to describe. Callers that compute a key and use + it immediately are safe: a config mutated underneath them yields a key that + simply misses and rebuilds. """ return ( connection_string, @@ -232,18 +258,59 @@ def get_engine(datasource: DatasourceConfig) -> sa.Engine: that two datasources differing in (e.g.) warehouse get different cached engines — otherwise the connect listener would silently apply the wrong USE statements. + + The cap is re-applied on hits as well as inserts, so lowering + ``SLAYER_MAX_CACHED_ENGINES`` — or setting it to ``0`` to turn caching off + — takes effect on the next call instead of waiting for a miss. """ - connection_string = datasource.get_connection_string() - cache_key = _cache_key(datasource, connection_string) - cached = _engine_cache.get(cache_key) - if cached is not None: - _engine_cache.move_to_end(cache_key) + # Snapshot first: DatasourceConfig is mutable, and the key is computed long + # before _build_engine re-reads the credentials — engine construction sits + # between them. A rotation landing in that window would cache an engine + # under a fingerprint that doesn't describe what it authenticates as, which + # is the exact confusion the credential leg of the key exists to prevent. + # Shallow copy suffices: every field is a scalar, so this is already + # detached from later writes to the caller's object. + snapshot = datasource.model_copy() + connection_string = snapshot.get_connection_string() + cache_key = _cache_key(datasource=snapshot, connection_string=connection_string) + with _cache_lock: + cached = _engine_cache.get(cache_key) + if cached is None: + trimmed, reuse = [], False + else: + _engine_cache.move_to_end(cache_key) + # Re-apply the cap on the hit path too. A cache that only trims on + # insert stays oversized until the next miss, and a cap of 0 would + # keep serving entries admitted before it was set. + trimmed = _take_evictions_over_limit() + # A cap of 0 trims the entry we just touched, which is what "no + # caching" has to mean: fall through and build instead of handing + # back an engine we are about to dispose. + reuse = cache_key in _engine_cache + for stale in trimmed: + _dispose_quietly(engine=stale, reason="cache limit lowered") + if reuse: return cached + # Built outside the lock — construction does real work (BigQuery mints an + # API client), and every other datasource would queue behind it. engine = _build_engine( - datasource=datasource, connection_string=connection_string, + datasource=snapshot, connection_string=connection_string, ) - _engine_cache[cache_key] = engine - _evict_to_limit() + with _cache_lock: + winner = _engine_cache.get(cache_key) + if winner is not None: + # Another caller built the same engine while we were constructing. + # Converge on theirs so one key never backs two live pools. + _engine_cache.move_to_end(cache_key) + loser, engine = engine, winner + else: + loser = None + _engine_cache[cache_key] = engine + evicted = _take_evictions_over_limit() + if loser is not None: + _dispose_quietly(engine=loser, reason="lost a concurrent build race") + for stale in evicted: + _dispose_quietly(engine=stale, reason="evicted from engine cache") return engine @@ -262,7 +329,9 @@ def invalidate_engine(datasource: DatasourceConfig) -> bool: the datasource now carries. """ connection_string = datasource.get_connection_string() - evicted = _engine_cache.pop(_cache_key(datasource, connection_string), None) + cache_key = _cache_key(datasource=datasource, connection_string=connection_string) + with _cache_lock: + evicted = _engine_cache.pop(cache_key, None) if evicted is None: return False _dispose_quietly(engine=evicted, reason=f"credentials rejected for '{datasource.name}'") @@ -277,7 +346,9 @@ def reset_cache(*, dispose: bool = False) -> None: process down want ``dispose=True`` so server-side connections close promptly instead of waiting on garbage collection. """ + with _cache_lock: + dropped = list(_engine_cache.items()) + _engine_cache.clear() if dispose: - for key, engine in list(_engine_cache.items()): + for key, engine in dropped: _dispose_quietly(engine=engine, reason=f"cache reset ({key[0]})") - _engine_cache.clear() diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py index aba69c9d..8a5a157c 100644 --- a/tests/dialects/test_bigquery.py +++ b/tests/dialects/test_bigquery.py @@ -679,7 +679,7 @@ def test_build_engine_with_non_object_credentials_json_raises(payload: str) -> N def _oauth_info(**overrides) -> dict: - info = { + info = { # NOSONAR(S2068) — test fixture; placeholder grant, not real credentials "type": "authorized_user", "client_id": "cid.apps.googleusercontent.com", "client_secret": "csecret", @@ -737,8 +737,9 @@ def test_build_engine_oauth_without_project_raises() -> None: """An OAuth grant carries no project, so omitting it from the connection string is a config error rather than a confusing downstream 404.""" dialect = BigqueryDialect() + ds = _oauth_ds() with pytest.raises(ValueError, match="must be given in the connection string"): - dialect.build_engine(_oauth_ds(), connection_string="bigquery://") + dialect.build_engine(ds, connection_string="bigquery://") def test_build_engine_oauth_falls_back_to_quota_project() -> None: @@ -765,8 +766,9 @@ def test_build_engine_rejects_both_credential_kinds() -> None: credentials_json=json.dumps({"type": "service_account"}), oauth_credentials_json=json.dumps(_oauth_info()), ) + dialect = BigqueryDialect() with pytest.raises(ValueError, match="mutually exclusive"): - BigqueryDialect().build_engine(ds, connection_string="bigquery://p/d") + dialect.build_engine(ds, connection_string="bigquery://p/d") def test_build_engine_rejects_oauth_grant_in_credentials_json() -> None: @@ -775,8 +777,9 @@ def test_build_engine_rejects_oauth_grant_in_credentials_json() -> None: ds = DatasourceConfig( name="bq", type="bigquery", credentials_json=json.dumps(_oauth_info()), ) + dialect = BigqueryDialect() with pytest.raises(ValueError, match="Put OAuth grants in oauth_credentials_json"): - BigqueryDialect().build_engine(ds, connection_string="bigquery://p/d") + dialect.build_engine(ds, connection_string="bigquery://p/d") @pytest.mark.parametrize( @@ -788,8 +791,9 @@ def test_build_engine_rejects_oauth_grant_in_credentials_json() -> None: ) def test_build_engine_oauth_malformed_raises(payload: str, message: str) -> None: ds = DatasourceConfig(name="bq", type="bigquery", oauth_credentials_json=payload) + dialect = BigqueryDialect() with pytest.raises(ValueError, match=message): - BigqueryDialect().build_engine(ds, connection_string="bigquery://p/d") + dialect.build_engine(ds, connection_string="bigquery://p/d") # --------------------------------------------------------------------------- @@ -852,3 +856,34 @@ def test_credential_fingerprint_leaks_no_secret_material() -> None: fp = BigqueryDialect().credential_fingerprint(_oauth_ds()) for secret in ("rtok-alice", "csecret", "access-token-1"): assert secret not in fp + + +def test_credential_fingerprint_tolerates_malformed_oauth_json() -> None: + """The fingerprint runs on every cache-key lookup, so a stored grant that + won't parse has to yield a digest rather than raise — otherwise a bad + datasource breaks engine lookup instead of reaching ``build_engine``'s + clear error.""" + ds = DatasourceConfig( + name="bq", type="bigquery", oauth_credentials_json="not json at all", + ) + assert BigqueryDialect().credential_fingerprint(ds) + + +def test_credential_fingerprint_distinguishes_malformed_payloads() -> None: + """Two unparseable grants are still two different identities.""" + dialect = BigqueryDialect() + + def ds_for(payload: str) -> DatasourceConfig: + return DatasourceConfig(name="bq", type="bigquery", oauth_credentials_json=payload) + + assert dialect.credential_fingerprint(ds_for("garbage-alice")) != dialect.credential_fingerprint(ds_for("garbage-bob")) + + +def test_build_engine_oauth_validates_before_importing_optional_driver() -> None: + """Config errors must surface as themselves even where the optional + 'bigquery' extra is absent, so validation precedes the google.* imports.""" + ds = DatasourceConfig(name="bq", type="bigquery", oauth_credentials_json="not json") + dialect = BigqueryDialect() + with patch.dict("sys.modules", {"google.cloud": None, "google.oauth2.credentials": None}): + with pytest.raises(ValueError, match="is not valid JSON"): + dialect.build_engine(ds, connection_string="bigquery://p/d") diff --git a/tests/test_engine_factory.py b/tests/test_engine_factory.py index 38fecc94..b42a9a2a 100644 --- a/tests/test_engine_factory.py +++ b/tests/test_engine_factory.py @@ -15,6 +15,8 @@ bare ``sa.create_engine``. """ +import json +import threading from unittest.mock import MagicMock, patch import pytest @@ -139,7 +141,14 @@ def test_session_listener_invokes_dialect_apply_session_overrides(self) -> None: assert apply_mock.call_count >= 1 # Listener calls ``apply_session_overrides(dbapi_connection=..., datasource=...)`` # by name; the datasource is the kwarg, not a positional arg. - assert apply_mock.call_args.kwargs["datasource"] is ds + # + # It receives the snapshot ``get_engine`` took, not the caller's object. + # That is deliberate: warehouse / role / database / schema_name are all + # in this engine's cache-key fingerprint, so the USE statements have to + # keep matching the values it was admitted under. + applied = apply_mock.call_args.kwargs["datasource"] + assert applied is not ds + assert applied == ds class TestCacheKeying: @@ -327,6 +336,61 @@ def test_dispose_failure_does_not_break_caching(self, monkeypatch) -> None: assert isinstance(engine_factory.get_engine(second), sa.Engine) engine_factory.reset_cache() + def test_lowering_the_limit_trims_on_the_next_hit(self, monkeypatch) -> None: + """A cache that only trims on insert stays oversized until the next + miss. Re-applying the cap on hits makes the new bound take effect on + the next call.""" + engine_factory.reset_cache() + sources = [self._lite(i) for i in range(4)] + for ds in sources: + engine_factory.get_engine(ds) + assert len(engine_factory._engine_cache) == 4 + + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "2") + # A pure hit — no miss to piggyback the trim on. + hot = sources[-1] + assert engine_factory.get_engine(hot) is not None + assert len(engine_factory._engine_cache) == 2 + # The entry just touched is the most-recently-used, so it survives. + assert engine_factory._cache_key( + datasource=hot, connection_string=hot.get_connection_string(), + ) in engine_factory._engine_cache + engine_factory.reset_cache() + + def test_hit_trim_disposes_outside_the_lock(self, monkeypatch) -> None: + """Trimmed engines are disposed after ``_cache_lock`` is released — + ``dispose()`` closes sockets and must not run under the lock.""" + engine_factory.reset_cache() + cold, hot = self._lite(0), self._lite(1) + engine_factory.get_engine(cold) + engine_factory.get_engine(hot) + + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "1") + held: list[bool] = [] + + def record_lock_state(*, engine, reason): + held.append(engine_factory._cache_lock.locked()) + + with patch.object(engine_factory, "_dispose_quietly", side_effect=record_lock_state): + engine_factory.get_engine(hot) + assert held == [False], "disposal ran while holding the cache lock" + engine_factory.reset_cache() + + def test_zero_limit_bypasses_reuse_on_a_hit(self, monkeypatch) -> None: + """With caching off, a previously-cached key must not be served from + the cache — the entry is dropped and a fresh engine built, rather than + handing back one we are about to dispose.""" + engine_factory.reset_cache() + ds = self._lite(0) + first = engine_factory.get_engine(ds) + assert len(engine_factory._engine_cache) == 1 + + monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "0") + second = engine_factory.get_engine(ds) + assert second is not first, "a zero cap must not reuse the cached engine" + assert len(engine_factory._engine_cache) == 0 + engine_factory.reset_cache() + @pytest.mark.parametrize(argnames="raw", argvalues=["nonsense", "-1", ""]) def test_bad_limit_env_falls_back_to_default(self, monkeypatch, raw: str) -> None: monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, raw) @@ -384,3 +448,201 @@ def test_reset_disposes_only_when_asked(self) -> None: with patch.object(engine, "dispose") as disposed: engine_factory.reset_cache(dispose=True) disposed.assert_called_once() + + +class TestCacheConcurrency: + """Engines are reached from worker threads as well as the event loop, so + the cache's read/bump/insert sequences need one lock around them.""" + + @staticmethod + def _lite(n: int) -> DatasourceConfig: + return DatasourceConfig(name=f"lite{n}", type="sqlite", database=f"/tmp/slayer-conc-{n}.db") + + @staticmethod + def _join_all(threads: list[threading.Thread], *, timeout: float = 10.0) -> None: + """Join every worker and fail if any is still running. + + Without this, ``join(timeout=...)`` silently returns on a deadlocked + worker and the test passes — which is the one failure mode adding a + lock introduces. Workers are daemons so a genuine deadlock trips this + assertion instead of hanging the whole pytest process. + """ + for thread in threads: + thread.join(timeout=timeout) + stuck = [t.name for t in threads if t.is_alive()] + assert not stuck, f"workers still running after {timeout}s (deadlock?): {stuck}" + + def test_concurrent_misses_yield_one_shared_engine(self) -> None: + """Two threads missing on the same key must converge on one pool, and + the losing engine must be disposed rather than orphaned.""" + engine_factory.reset_cache() + ds = self._lite(0) + built: list[sa.Engine] = [] + gate = threading.Barrier(2) + real_build = engine_factory._build_engine + + def slow_build(**kwargs): + gate.wait(timeout=5) # force both threads past the first lookup + engine = real_build(**kwargs) + built.append(engine) + return engine + + results: list[sa.Engine] = [] + with ( + patch.object(engine_factory, "_build_engine", side_effect=slow_build), + patch.object(engine_factory, "_dispose_quietly") as dispose, + ): + threads = [ + threading.Thread( + target=lambda: results.append(engine_factory.get_engine(ds)), + daemon=True, + ) + for _ in range(2) + ] + for t in threads: + t.start() + self._join_all(threads) + + assert len(built) == 2, "both threads should have raced past the lookup" + assert results[0] is results[1], "callers must converge on one engine" + assert len(engine_factory._engine_cache) == 1 + # The build the cache didn't keep. Leaving it undisposed is the + # pool leak the double-check exists to prevent, so assert on the + # engine identity — not on the log-facing reason string. + loser = next(engine for engine in built if engine is not results[0]) + dispose.assert_called_once() + assert dispose.call_args.kwargs["engine"] is loser + + engine_factory.reset_cache() + + def test_invalidate_racing_a_lookup_does_not_raise(self) -> None: + """Pre-lock, an invalidate landing between ``get`` and ``move_to_end`` + raised KeyError.""" + engine_factory.reset_cache() + sources = [self._lite(i) for i in range(4)] + for ds in sources: + engine_factory.get_engine(ds) + + errors: list[BaseException] = [] + stop = threading.Event() + + def hammer(fn): + try: + while not stop.is_set(): + for ds in sources: + fn(ds) + except BaseException as exc: # noqa: BLE001 — the point is to catch anything + errors.append(exc) + + threads = [ + threading.Thread(target=hammer, args=(engine_factory.get_engine,), daemon=True), + threading.Thread(target=hammer, args=(engine_factory.invalidate_engine,), daemon=True), + ] + for t in threads: + t.start() + stop.wait(timeout=1.0) + stop.set() + self._join_all(threads) + + assert not errors, f"concurrent access raised: {errors[:3]}" + engine_factory.reset_cache() + + +class TestConfigSnapshot: + """``get_engine`` snapshots the config before deriving anything from it. + A ``DatasourceConfig`` is mutable, and engine construction sits between the + cache key and the dialect's second read of the credentials.""" + + @staticmethod + def _oauth_ds(refresh_token: str) -> DatasourceConfig: + return DatasourceConfig( + name="bq", type="bigquery", connection_string="bigquery://proj/dset", + oauth_credentials_json=json.dumps({ + "type": "authorized_user", # NOSONAR(S2068) — placeholder grant + "refresh_token": refresh_token, + "client_id": "cid", + "client_secret": "csecret", + }), + ) + + def test_rotation_mid_build_cannot_desync_key_from_engine(self) -> None: + """A grant refresh landing while the engine is under construction must + not leave that engine cached under a fingerprint describing the *other* + set of credentials — the confusion the credential leg exists to stop.""" + engine_factory.reset_cache() + ds = self._oauth_ds("before") + key_for_before = engine_factory._cache_key( + datasource=ds, connection_string="bigquery://proj/dset", + ) + seen: dict = {} + + def rotate_then_build(*, datasource, connection_string): + # Stands in for a concurrent token refresh mutating the caller's + # config while we are inside _build_engine. + ds.oauth_credentials_json = json.dumps({"refresh_token": "after"}) + seen["oauth"] = datasource.oauth_credentials_json + return MagicMock(spec=sa.Engine) + + with patch.object(engine_factory, "_build_engine", side_effect=rotate_then_build): + engine_factory.get_engine(ds) + + assert list(engine_factory._engine_cache) == [key_for_before] + # The decisive assertion: the build saw the credentials that key + # describes. Without the snapshot it would have seen "after" while the + # entry sat under the "before" key. + assert json.loads(seen["oauth"])["refresh_token"] == "before" + engine_factory.reset_cache() + + def test_caller_mutation_does_not_leak_into_the_cached_engine(self) -> None: + """Mutating the config after the call is self-correcting: the next + lookup keys off the new credentials, misses, and rebuilds.""" + engine_factory.reset_cache() + ds = self._oauth_ds("before") + with patch.object( + engine_factory, "_build_engine", + side_effect=lambda **_: MagicMock(spec=sa.Engine), + ): + first = engine_factory.get_engine(ds) + ds.oauth_credentials_json = json.dumps({"refresh_token": "after"}) + second = engine_factory.get_engine(ds) + assert first is not second, "rotated credentials must not reuse the old engine" + assert len(engine_factory._engine_cache) == 2 + engine_factory.reset_cache() + + def test_snapshot_is_detached_from_the_callers_object(self) -> None: + """Sanity-check the copy depth: every field is scalar, so a shallow + model_copy already detaches. If a mutable field is ever added, this is + where the assumption breaks.""" + ds = self._oauth_ds("before") + snapshot = ds.model_copy() + ds.oauth_credentials_json = "mutated" + ds.warehouse = "mutated" + assert snapshot.oauth_credentials_json != "mutated" + assert snapshot.warehouse != "mutated" + + def test_session_overrides_listener_is_insulated_from_later_mutation(self) -> None: + """The listener fires on every checkout for the life of the engine. It + must keep applying the fields this engine was cached under, not whatever + the caller's config says later — those fields *are* its cache key.""" + pytest.importorskip("snowflake.connector") + pytest.importorskip("snowflake.sqlalchemy") + engine_factory.reset_cache() + ds = DatasourceConfig( + name="sf", type="snowflake", connection_name="default", warehouse="WH_AT_BUILD", + ) + real_engine = sa.create_engine("sqlite:///:memory:") + with ( + patch( + "slayer.sql.dialects.snowflake.SnowflakeDialect.build_engine", + return_value=real_engine, + ), + patch( + "slayer.sql.dialects.snowflake.SnowflakeDialect.apply_session_overrides", + ) as apply_mock, + ): + engine = engine_factory.get_engine(ds) + ds.warehouse = "WH_MUTATED_AFTER" + with engine.connect() as _: + pass # NOSONAR(S108) — opening + closing fires the checkout listener + assert apply_mock.call_args.kwargs["datasource"].warehouse == "WH_AT_BUILD" + engine_factory.reset_cache() diff --git a/tests/test_sql_client.py b/tests/test_sql_client.py index ec7dd7b2..f3093db5 100644 --- a/tests/test_sql_client.py +++ b/tests/test_sql_client.py @@ -2,7 +2,7 @@ import logging import sqlite3 -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest import sqlalchemy.exc @@ -529,6 +529,67 @@ async def test_non_auth_failure_keeps_the_engine(self) -> None: invalidate.assert_not_called() assert client._sync_engine is engine + async def test_async_engine_is_disposed_too(self) -> None: + """Native-async dialects hold a second pool that ``invalidate_engine`` + knows nothing about; it has to go as well.""" + client = self._client() + async_engine = AsyncMock() + client._async_engine = async_engine + with ( + patch.object(sql_client.SlayerSQLClient, "_execute", side_effect=Exception("invalid_grant")), + patch("slayer.sql.engine_factory.invalidate_engine"), + ): + with pytest.raises(Exception, match="invalid_grant"): + await client.execute("SELECT 1") + async_engine.dispose.assert_awaited_once() + assert client._async_engine is None + + async def test_get_column_types_also_discards(self) -> None: + """It runs its own SQL against the same cached engine, so it needs the + same cleanup ``execute`` gets.""" + client = self._client() + client._sync_engine = object() + with ( + patch.object( + sql_client.SlayerSQLClient, "_get_column_types", + side_effect=Exception("invalid_grant"), + ), + patch("slayer.sql.engine_factory.invalidate_engine") as invalidate, + ): + with pytest.raises(Exception, match="invalid_grant"): + await client.get_column_types("SELECT 1") + invalidate.assert_called_once_with(client.datasource) + assert client._sync_engine is None + + def test_execute_sync_also_discards(self) -> None: + """The sync path shares the factory-cached engine. It cannot dispose an + async pool (no loop to do it on), so it only drops the sync one.""" + client = self._client() + client._sync_engine = object() + with ( + patch("slayer.sql.client._execute_with_retry_sync", side_effect=Exception("invalid_grant")), + patch.object(sql_client.SlayerSQLClient, "_get_sync_engine_for_client", return_value=None), + patch("slayer.sql.engine_factory.invalidate_engine") as invalidate, + ): + with pytest.raises(Exception, match="invalid_grant"): + client.execute_sync("SELECT 1") + invalidate.assert_called_once_with(client.datasource) + assert client._sync_engine is None + + def test_execute_sync_keeps_engine_on_non_auth_failure(self) -> None: + client = self._client() + engine = object() + client._sync_engine = engine + with ( + patch("slayer.sql.client._execute_with_retry_sync", side_effect=Exception("no such table")), + patch.object(sql_client.SlayerSQLClient, "_get_sync_engine_for_client", return_value=None), + patch("slayer.sql.engine_factory.invalidate_engine") as invalidate, + ): + with pytest.raises(Exception, match="no such table"): + client.execute_sync("SELECT 1") + invalidate.assert_not_called() + assert client._sync_engine is engine + async def test_cleanup_failure_does_not_mask_the_original_error(self) -> None: """The auth error is what the caller needs to see; a failed eviction must not displace it.""" From ed3627220e4fa95f50e9fddbb0243a6462614f3a Mon Sep 17 00:00:00 2001 From: AivanF Date: Fri, 7 Aug 2026 13:22:46 +0300 Subject: [PATCH 3/5] Fixes from review --- slayer/sql/dialects/bigquery.py | 4 +-- slayer/sql/engine_factory.py | 19 +++++++++-- tests/dialects/test_bigquery.py | 16 +++++---- tests/test_engine_factory.py | 57 +++++++++++++++++++++++++++++++++ 4 files changed, 86 insertions(+), 10 deletions(-) diff --git a/slayer/sql/dialects/bigquery.py b/slayer/sql/dialects/bigquery.py index 47d84ed5..c0a2a7a5 100644 --- a/slayer/sql/dialects/bigquery.py +++ b/slayer/sql/dialects/bigquery.py @@ -258,7 +258,7 @@ def build_engine( f"oauth_credentials_json instead." ) return sa.create_engine( - connection_string, + url=connection_string, credentials_info=credentials_info, pool_pre_ping=True, ) @@ -315,7 +315,7 @@ def _build_oauth_engine( f"a usable authorized-user grant: {exc}" ) from exc return sa.create_engine( - url.update_query_dict({"user_supplied_client": "true"}), + url=url.update_query_dict({"user_supplied_client": "true"}), connect_args={"client": bigquery.Client( project=project, credentials=credentials, )}, diff --git a/slayer/sql/engine_factory.py b/slayer/sql/engine_factory.py index 815eb23a..2e425983 100644 --- a/slayer/sql/engine_factory.py +++ b/slayer/sql/engine_factory.py @@ -31,7 +31,7 @@ from slayer.core.models import DatasourceConfig from slayer.sql.dialects import dialect_for_ds_type -from slayer.sql.dialects.base import SqlDialect +from slayer.sql.dialects.base import SqlDialect, _digest logger = logging.getLogger(__name__) @@ -101,6 +101,18 @@ def _max_cached_engines() -> int: return value +def loggable_key(key: EngineCacheKey) -> str: + """A short, stable id for a cache key that is safe to put in a log line. + + ``key[0]`` is the connection string, and for username/password dialects + ``DatasourceConfig.get_connection_string`` renders it with + ``hide_password=False`` — so the raw key carries a plaintext password and + must never be logged. Digesting the whole key keeps entries correlatable + across log lines while reducing every leg to non-reversible hex. + """ + return _digest("\x00".join(key)) + + def _dispose_quietly(*, engine: sa.Engine, reason: str) -> None: """Release an engine's pooled connections, logging rather than raising. @@ -108,6 +120,9 @@ def _dispose_quietly(*, engine: sa.Engine, reason: str) -> None: pool; connections checked out by an in-flight query are detached and closed when returned. So this is safe to call on an engine another caller still holds a reference to — it reclaims sockets without breaking them. + + ``reason`` reaches a log line, so callers must keep secrets out of it — + see :func:`loggable_key` for cache keys. """ try: engine.dispose() @@ -351,4 +366,4 @@ def reset_cache(*, dispose: bool = False) -> None: _engine_cache.clear() if dispose: for key, engine in dropped: - _dispose_quietly(engine=engine, reason=f"cache reset ({key[0]})") + _dispose_quietly(engine=engine, reason=f"cache reset ({loggable_key(key)})") diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py index 8a5a157c..58aa83b3 100644 --- a/tests/dialects/test_bigquery.py +++ b/tests/dialects/test_bigquery.py @@ -24,8 +24,7 @@ from slayer.core.query import ColumnRef, SlayerQuery from slayer.engine.enriched import EnrichedQuery from slayer.engine.enrichment import enrich_query -from slayer.engine.query_engine import SlayerQueryEngine -from slayer.engine.query_engine import _sql_client_cache_key +from slayer.engine.query_engine import SlayerQueryEngine, _sql_client_cache_key from slayer.sql.dialects import ( BigqueryDialect, PostgresDialect, @@ -814,7 +813,10 @@ def test_credential_fingerprint_differs_between_oauth_users() -> None: alice = dialect.credential_fingerprint(_oauth_ds(refresh_token="rtok-alice")) bob = dialect.credential_fingerprint(_oauth_ds(refresh_token="rtok-bob")) assert alice != bob - assert alice and bob + # Neither may collapse to the empty "no credentials" fingerprint, which + # would drop both users into the Application-Default-Credentials bucket. + assert alice != "" + assert bob != "" def test_credential_fingerprint_differs_between_oauth_and_service_account() -> None: @@ -884,6 +886,8 @@ def test_build_engine_oauth_validates_before_importing_optional_driver() -> None 'bigquery' extra is absent, so validation precedes the google.* imports.""" ds = DatasourceConfig(name="bq", type="bigquery", oauth_credentials_json="not json") dialect = BigqueryDialect() - with patch.dict("sys.modules", {"google.cloud": None, "google.oauth2.credentials": None}): - with pytest.raises(ValueError, match="is not valid JSON"): - dialect.build_engine(ds, connection_string="bigquery://p/d") + with ( + patch.dict("sys.modules", {"google.cloud": None, "google.oauth2.credentials": None}), + pytest.raises(ValueError, match="is not valid JSON"), + ): + dialect.build_engine(ds, connection_string="bigquery://p/d") diff --git a/tests/test_engine_factory.py b/tests/test_engine_factory.py index b42a9a2a..468b360d 100644 --- a/tests/test_engine_factory.py +++ b/tests/test_engine_factory.py @@ -434,6 +434,63 @@ def test_next_get_engine_rebuilds(self) -> None: engine_factory.reset_cache() +class TestLogSafety: + """Cache keys carry the connection string, which for username/password + dialects is rendered with the password in plaintext. None of it may reach + a log line.""" + + @staticmethod + def _pg_with_password(password: str) -> DatasourceConfig: + return DatasourceConfig( + name="pg", type="postgres", host="h", username="u", + password=password, database="db", + ) + + def test_loggable_key_hides_the_connection_string(self) -> None: + secret = "hunter2-plaintext" # NOSONAR(S2068) — test fixture + ds = self._pg_with_password(secret) + key = engine_factory._cache_key( + datasource=ds, connection_string=ds.get_connection_string(), + ) + assert secret in key[0], "precondition: the raw key does carry the password" + rendered = engine_factory.loggable_key(key) + for fragment in (secret, "u", "postgresql://"): + assert fragment not in rendered + + def test_loggable_key_is_stable_and_distinguishing(self) -> None: + """Useful for correlating log lines: same key -> same id, different + credentials -> different id.""" + alice = self._pg_with_password("alice-pw") # NOSONAR(S2068) — test fixture + # A separate object carrying the same values: the id must follow the + # credentials, not the identity of the config object. + alice_again = self._pg_with_password("alice-pw") # NOSONAR(S2068) — test fixture + bob = self._pg_with_password("bob-pw") # NOSONAR(S2068) — test fixture + + def log_id(ds): + return engine_factory.loggable_key(engine_factory._cache_key( + datasource=ds, connection_string=ds.get_connection_string(), + )) + + assert log_id(alice) == log_id(alice_again) + assert log_id(alice) != log_id(bob) + + def test_reset_disposal_reason_carries_no_credentials(self) -> None: + """``reset_cache(dispose=True)`` builds its reason string from the + cache key; ``_dispose_quietly`` writes that into a warning when + ``dispose()`` raises.""" + secret = "reset-time-secret" # NOSONAR(S2068) — test fixture + engine_factory.reset_cache() + engine_factory.get_engine(self._pg_with_password(secret)) + reasons: list[str] = [] + with patch.object( + engine_factory, "_dispose_quietly", + side_effect=lambda *, engine, reason: reasons.append(reason), + ): + engine_factory.reset_cache(dispose=True) + assert reasons, "precondition: disposal ran" + assert not any(secret in reason for reason in reasons), reasons + + class TestResetCacheDisposal: def test_reset_disposes_only_when_asked(self) -> None: From bc079622c3d3b2c7f6438d6cbbb5edab60e46ac4 Mon Sep 17 00:00:00 2001 From: AivanF Date: Sat, 8 Aug 2026 16:52:12 +0300 Subject: [PATCH 4/5] Neater doc-strings --- slayer/core/models.py | 18 ++--- slayer/sql/client.py | 57 ++++++-------- slayer/sql/dialects/bigquery.py | 74 +++++++----------- slayer/sql/engine_factory.py | 132 +++++++++++--------------------- tests/dialects/test_bigquery.py | 52 +++++-------- tests/test_engine_factory.py | 82 +++++++++----------- tests/test_sql_client.py | 29 ++++--- 7 files changed, 171 insertions(+), 273 deletions(-) diff --git a/slayer/core/models.py b/slayer/core/models.py index fba03c76..f4c6b19c 100644 --- a/slayer/core/models.py +++ b/slayer/core/models.py @@ -792,19 +792,11 @@ class DatasourceConfig(BaseModel): # When unset, BigQuery falls back to Application Default Credentials # (``GOOGLE_APPLICATION_CREDENTIALS`` env var or attached compute identity). credentials_json: str | None = Field(default=None, repr=False) - # BigQuery-specific. A Google OAuth *authorized user* grant as a JSON - # string — the shape ``google.oauth2.credentials.Credentials - # .from_authorized_user_info`` consumes (``token``, ``refresh_token``, - # ``client_id``, ``client_secret``, ``token_uri``, ``scopes``). This is - # the per-end-user auth path: the caller obtains the grant from its own - # OAuth flow and hands SLayer a datasource carrying it, so queries run - # as that user with that user's BigQuery permissions. - # - # Mutually exclusive with ``credentials_json`` (service account, one - # shared identity for everyone). ``credentials_json`` cannot carry an - # OAuth grant: ``sqlalchemy-bigquery`` feeds it to - # ``service_account.Credentials.from_service_account_info``, which only - # understands service-account keys. + # BigQuery-specific. A Google OAuth authorized-user grant as JSON, in the + # shape ``Credentials.from_authorized_user_info`` consumes — the per-end-user + # auth path, so queries run with that user's permissions. Mutually exclusive + # with ``credentials_json``, which cannot carry a grant: the driver feeds it + # to ``from_service_account_info``. oauth_credentials_json: str | None = Field(default=None, repr=False) @model_validator(mode="before") diff --git a/slayer/sql/client.py b/slayer/sql/client.py index 05dc091f..6dce6860 100644 --- a/slayer/sql/client.py +++ b/slayer/sql/client.py @@ -529,18 +529,14 @@ def _get_sync_engine_for_client(self) -> sa.Engine | None: return self._sync_engine def _discard_sync_engine_on_auth_failure(self, exc: BaseException) -> bool: - """Drop the sync engine when the failure was a credential rejection. - Returns whether ``exc`` was one. - - The credentials an engine authenticates with are fixed at construction - (BigQuery bakes a client object into the engine), so a revoked OAuth - grant or rotated key poisons that engine permanently — every later call - through the cache reproduces the same failure. Evicting means the next - call rebuilds from whatever credentials the datasource now carries, - which is exactly what a caller that just refreshed its grant expects. - - Best-effort and non-fatal: the original error is what the caller needs - to see, so a failure to clean up must not displace it. + """Drop the sync engine on a credential rejection; returns whether + ``exc`` was one. + + An engine's credentials are fixed at construction, so a revoked grant + poisons it permanently — every later call through the cache fails the + same way. Evicting lets the next one rebuild from whatever the + datasource now carries. Best-effort: cleanup must not displace the + original error. """ if not _is_auth_failure(exc): return False @@ -556,11 +552,11 @@ def _discard_sync_engine_on_auth_failure(self, exc: BaseException) -> bool: return True async def _discard_engines_on_auth_failure(self, exc: BaseException) -> None: - """Async-path cleanup: the sync engine, plus this client's async one. + """Async-path cleanup: the sync engine plus this client's async one. - Native-async dialects hold a second pool that ``invalidate_engine`` - knows nothing about, and disposing it needs a loop — hence the split - from the sync variant, which ``execute_sync`` uses. + Native-async dialects hold a second pool ``invalidate_engine`` knows + nothing about, and disposing it needs a loop — hence the split from the + sync variant used by ``execute_sync``. """ if not self._discard_sync_engine_on_auth_failure(exc): return @@ -648,9 +644,8 @@ def execute_sync( ) -> list[dict[str, Any]]: """Execute SQL synchronously (for CLI, notebooks, tests). - Only the sync engine is discarded on a credential rejection: this - method has no loop to dispose an async pool on, and the paths that do - (``execute`` / ``get_column_types``) handle it. + Discards only the sync engine on a credential rejection — no loop here + to dispose an async pool on; ``execute`` / ``get_column_types`` cover that. """ try: return _execute_with_retry_sync( @@ -713,14 +708,11 @@ def _is_transient_db_error(exc: BaseException) -> bool: return any(sig in msg for sig in _TRANSIENT_DB_ERROR_SIGNALS) -# Credential rejection, as opposed to a transient blip. Distinct from -# _TRANSIENT_DB_ERROR_SIGNALS on purpose: retrying these is pointless (the -# credentials are baked into the engine, so every attempt fails identically), -# and the right response is to throw the engine away rather than sleep. -# -# Deliberately narrow. Postgres' table-level "permission denied for table" -# is NOT here — the credentials worked fine, the grant didn't, and evicting -# a healthy engine over it is pure pool churn. +# Credential rejection, deliberately distinct from _TRANSIENT_DB_ERROR_SIGNALS: +# retrying is pointless (the credentials are baked into the engine), so the +# response is to throw the engine away rather than sleep. Kept narrow — +# Postgres' table-level "permission denied for table" is absent on purpose, +# since the credentials worked and evicting over it is pure pool churn. _AUTH_ERROR_SIGNALS = ( "invalid_grant", # OAuth refresh token revoked / expired "invalid_client", @@ -734,8 +726,8 @@ def _is_transient_db_error(exc: BaseException) -> bool: "invalid username or password", ) -# Matched by class name so this stays dependency-free — google-auth and -# google-api-core ship only with the optional 'bigquery' extra. +# Matched by class name to stay dependency-free: google-auth ships only with +# the optional 'bigquery' extra. _AUTH_ERROR_TYPE_NAMES = frozenset({ "RefreshError", # google.auth.exceptions "DefaultCredentialsError", # google.auth.exceptions @@ -744,11 +736,10 @@ def _is_transient_db_error(exc: BaseException) -> bool: def _is_auth_failure(exc: BaseException) -> bool: - """Return True when the server rejected the *credentials* themselves. + """True when the server rejected the *credentials* themselves. - Walks the cause/context chain (plus SQLAlchemy's ``orig``) because the - signal is almost always a driver exception wrapped one or two layers deep - by the time it surfaces. + Walks the cause/context chain plus SQLAlchemy's ``orig`` — the signal + surfaces wrapped a layer or two deep. """ seen: set[int] = set() pending: list[BaseException] = [exc] diff --git a/slayer/sql/dialects/bigquery.py b/slayer/sql/dialects/bigquery.py index c0a2a7a5..fd895c6a 100644 --- a/slayer/sql/dialects/bigquery.py +++ b/slayer/sql/dialects/bigquery.py @@ -210,27 +210,18 @@ def build_engine( *, connection_string: str, ) -> "sa.Engine | None": - """Construct the SQLAlchemy engine for whichever auth path the - datasource configures. + """Build the engine for whichever auth path is configured, in order: - Three paths, in precedence order: - - 1. ``oauth_credentials_json`` — a per-end-user OAuth grant. Queries - run as that user against their own BigQuery permissions. See + 1. ``oauth_credentials_json`` — per-end-user grant, see :meth:`_build_oauth_engine`. - 2. ``credentials_json`` — inline service-account key. ``sqlalchemy - -bigquery`` accepts a ``credentials_info`` kwarg on - ``create_engine``, a dict matching the key file's shape, so the - BigQuery client builds credentials from it directly with no temp - file. One shared identity for every caller. - 3. Neither — return ``None`` so ``engine_factory`` falls back to the - default ``create_engine`` and the BigQuery client picks up - Application Default Credentials. - - Setting both (1) and (2) is a configuration error rather than a - silent precedence win: the two mean different identities, and - guessing which one the caller meant is how a per-user query - quietly runs as the shared service account. + 2. ``credentials_json`` — service-account key, passed straight through + as ``credentials_info``. One shared identity for every caller. + 3. Neither — ``None``, so the factory falls back to a plain + ``create_engine`` and the client picks up ADC. + + Setting both is an error rather than a silent precedence win: they are + different identities, and guessing is how a per-user query quietly runs + as the service account. """ if datasource.oauth_credentials_json and datasource.credentials_json: raise ValueError( @@ -271,24 +262,16 @@ def _build_oauth_engine( ) -> "sa.Engine": """Build an engine bound to a caller-supplied OAuth user grant. - ``sqlalchemy-bigquery`` has no kwarg for OAuth credentials — its - ``credentials_info``/``credentials_path``/``credentials_base64`` - kwargs all route to ``service_account.Credentials``. The supported - escape hatch is its ``user_supplied_client`` URL flag: with it set, - ``create_connect_args`` skips building a client of its own and takes - ours from ``connect_args={"client": ...}``. Without the flag the - driver would first construct an Application-Default-Credentials - client (failing outright where no ADC exists) before ours displaced - it, so the flag is load-bearing, not decorative. - - The BigQuery project is not part of an OAuth grant the way it is - part of a service-account key, so it has to come from the - connection string's host (``bigquery:///``), or - from the grant's ``quota_project_id`` when it carries one. - - Config is validated before the ``google.*`` imports so a - misconfigured datasource reports *that* rather than a missing - optional dependency — those ship only with the 'bigquery' extra. + ``sqlalchemy-bigquery`` has no OAuth kwarg — every credentials kwarg it + has routes to ``service_account.Credentials``. Its ``user_supplied_client`` + URL flag is the supported escape hatch: with it set, the driver takes our + client from ``connect_args`` instead of first building an ADC one (which + fails outright where no ADC exists), so the flag is load-bearing. + + A grant carries no project, so it comes from the URL host or the grant's + ``quota_project_id``. Config is validated before the ``google.*`` + imports — they ship only with the 'bigquery' extra, and a misconfigured + datasource should report that, not a missing dependency. """ info = _parse_credentials_object( raw=datasource.oauth_credentials_json, @@ -323,16 +306,13 @@ def _build_oauth_engine( ) def credential_fingerprint(self, datasource: "DatasourceConfig") -> str: - """Identity of both BigQuery auth paths, so cached engines never - cross between a service account and an end user, or between two - end users. - - The OAuth half deliberately digests the *durable* grant rather than - the raw JSON: an access token rotates, and keying on it would mint - (and leak) a fresh engine on every refresh. Dropping the rotating - fields is only safe while a refresh token pins the identity — - without one the access token IS the whole identity, and removing it - would let two different users share one engine. + """Identity across both auth paths, so a cached engine never crosses + between a service account and an end user, or between two end users. + + The OAuth half digests the *durable* grant: keying on a rotating access + token would mint a fresh engine per refresh. Dropping those fields is + only safe while a refresh token pins the identity — without one the + access token is the whole identity. """ material = [datasource.credentials_json or ""] raw_oauth = datasource.oauth_credentials_json diff --git a/slayer/sql/engine_factory.py b/slayer/sql/engine_factory.py index 2e425983..87fed6b8 100644 --- a/slayer/sql/engine_factory.py +++ b/slayer/sql/engine_factory.py @@ -36,39 +36,23 @@ logger = logging.getLogger(__name__) -#: Identity of a cached engine: URL + runtime fields + credentials. Exported so -#: callers that key their own caches by engine identity stay in lockstep with -#: this module instead of hard-coding the arity. +#: (connection_string, runtime_fingerprint, credential_fingerprint). Exported so +#: callers keying their own caches by engine identity don't hard-code the arity. EngineCacheKey = tuple[str, str, str] -# Engine cache. Key = (connection_string, runtime_fingerprint, credential_fingerprint). -# The credential leg is load-bearing for security: a dialect whose secret is not -# in the URL (BigQuery's service-account / OAuth credentials) would otherwise -# have two differently-authenticated callers share one engine, silently running -# one identity's queries under the other's credentials. -# -# Ordered least- to most-recently-used, and bounded: once credentials are part -# of the key, cardinality follows the number of distinct *identities*, not -# datasources. A deployment handing SLayer per-end-user credentials would -# otherwise accumulate one engine — and one connection pool — per user who ever -# ran a query, for the life of the process. +# LRU-ordered, bounded. The credential leg is a security boundary: where the +# secret isn't in the URL (BigQuery), two identities would otherwise share one +# engine. Including it also makes cardinality track users, hence the cap. _engine_cache: "OrderedDict[EngineCacheKey, sa.Engine]" = OrderedDict() -# Guards every read, recency bump, insert, eviction and reset of the cache -# above. Engines are reached from worker threads (``_run_sync_in_thread`` in -# slayer.sql.client) as well as from the event loop, so the lookup/move_to_end -# pair is not safe unsynchronised: an ``invalidate_engine`` landing between the -# two raises KeyError, and two threads missing at once build two pools for one -# key and silently orphan the loser. -# -# Never held across engine construction or ``dispose()`` — both do real I/O, -# and serialising unrelated datasources behind them would cost more than the -# race it prevents. +# Engines are reached from worker threads as well as the event loop, so the +# lookup/move_to_end pair needs guarding: an interleaved invalidate raises +# KeyError, and simultaneous misses orphan a pool. Never held across engine +# construction or dispose() — both do I/O. _cache_lock = threading.Lock() -#: Cap on simultaneously cached engines. Sized for "every datasource, times a -#: working set of active users" rather than for a whole user base; a miss costs -#: one engine construction, so over-eviction degrades latency, not correctness. +#: Sized for "datasources x working set of active users". A miss costs one +#: engine build, so over-eviction hurts latency, not correctness. DEFAULT_MAX_CACHED_ENGINES = 64 #: Env override for :data:`DEFAULT_MAX_CACHED_ENGINES`. ``0`` disables caching. @@ -76,11 +60,7 @@ def _max_cached_engines() -> int: - """Resolve the cache cap, falling back to the default on junk input. - - Read per call rather than at import so tests and hosts can retune it - without reloading the module. - """ + """Cache cap, defaulting on junk input. Read per call so it stays retunable.""" raw = os.environ.get(MAX_CACHED_ENGINES_ENV) if raw is None: return DEFAULT_MAX_CACHED_ENGINES @@ -102,13 +82,11 @@ def _max_cached_engines() -> int: def loggable_key(key: EngineCacheKey) -> str: - """A short, stable id for a cache key that is safe to put in a log line. + """Short, stable, log-safe id for a cache key. - ``key[0]`` is the connection string, and for username/password dialects - ``DatasourceConfig.get_connection_string`` renders it with - ``hide_password=False`` — so the raw key carries a plaintext password and - must never be logged. Digesting the whole key keeps entries correlatable - across log lines while reducing every leg to non-reversible hex. + ``key[0]`` is the connection string, rendered with ``hide_password=False`` + — so the raw key must never reach a log line. The digest stays correlatable + across lines without being reversible. """ return _digest("\x00".join(key)) @@ -116,13 +94,9 @@ def loggable_key(key: EngineCacheKey) -> str: def _dispose_quietly(*, engine: sa.Engine, reason: str) -> None: """Release an engine's pooled connections, logging rather than raising. - ``Engine.dispose()`` closes checked-in connections and swaps in a fresh - pool; connections checked out by an in-flight query are detached and - closed when returned. So this is safe to call on an engine another caller - still holds a reference to — it reclaims sockets without breaking them. - - ``reason`` reaches a log line, so callers must keep secrets out of it — - see :func:`loggable_key` for cache keys. + Safe on an engine someone still holds: ``dispose()`` swaps in a fresh pool + and lets in-flight connections close on return. ``reason`` is logged, so + keep secrets out of it — see :func:`loggable_key`. """ try: engine.dispose() @@ -131,11 +105,10 @@ def _dispose_quietly(*, engine: sa.Engine, reason: str) -> None: def _take_evictions_over_limit() -> list[sa.Engine]: - """Pop least-recently-used entries until the cache fits its cap. + """Pop LRU entries until the cache fits its cap. - Returns the evicted engines rather than disposing them, because callers - run this holding ``_cache_lock`` and ``dispose()`` closes sockets. Dispose - the result after releasing the lock. + Returns them instead of disposing: callers hold ``_cache_lock``, and + ``dispose()`` does I/O. Dispose after releasing. """ limit = _max_cached_engines() evicted: list[sa.Engine] = [] @@ -152,11 +125,9 @@ def _cache_key(datasource: DatasourceConfig, connection_string: str) -> EngineCa with it; a divergence between the two caches means a caller can get a client whose engine was built for different credentials. - ``datasource`` is read, not retained. Only ``get_engine`` snapshots it - first, because only there does a slow engine build sit between the key and - the credentials it claims to describe. Callers that compute a key and use - it immediately are safe: a config mutated underneath them yields a key that - simply misses and rebuilds. + Only ``get_engine`` snapshots ``datasource`` first — it is the one caller + with a slow build between the key and the credentials it describes. Callers + that use the key immediately just miss and rebuild. """ return ( connection_string, @@ -275,16 +246,13 @@ def get_engine(datasource: DatasourceConfig) -> sa.Engine: apply the wrong USE statements. The cap is re-applied on hits as well as inserts, so lowering - ``SLAYER_MAX_CACHED_ENGINES`` — or setting it to ``0`` to turn caching off - — takes effect on the next call instead of waiting for a miss. + ``SLAYER_MAX_CACHED_ENGINES`` takes effect on the next call, not the next + miss. """ - # Snapshot first: DatasourceConfig is mutable, and the key is computed long - # before _build_engine re-reads the credentials — engine construction sits - # between them. A rotation landing in that window would cache an engine - # under a fingerprint that doesn't describe what it authenticates as, which - # is the exact confusion the credential leg of the key exists to prevent. - # Shallow copy suffices: every field is a scalar, so this is already - # detached from later writes to the caller's object. + # DatasourceConfig is mutable and the build sits between the key and the + # dialect's second read of the credentials; a rotation in that window would + # cache an engine under a fingerprint that misdescribes it. Shallow is + # enough — every field is a scalar. snapshot = datasource.model_copy() connection_string = snapshot.get_connection_string() cache_key = _cache_key(datasource=snapshot, connection_string=connection_string) @@ -294,28 +262,25 @@ def get_engine(datasource: DatasourceConfig) -> sa.Engine: trimmed, reuse = [], False else: _engine_cache.move_to_end(cache_key) - # Re-apply the cap on the hit path too. A cache that only trims on - # insert stays oversized until the next miss, and a cap of 0 would - # keep serving entries admitted before it was set. + # Trim here too, else a lowered cap waits for the next miss. trimmed = _take_evictions_over_limit() - # A cap of 0 trims the entry we just touched, which is what "no - # caching" has to mean: fall through and build instead of handing - # back an engine we are about to dispose. + # A cap of 0 trims the entry we just touched: fall through and + # build rather than hand back an engine we are about to dispose. reuse = cache_key in _engine_cache for stale in trimmed: _dispose_quietly(engine=stale, reason="cache limit lowered") if reuse: return cached - # Built outside the lock — construction does real work (BigQuery mints an - # API client), and every other datasource would queue behind it. + # Built outside the lock: construction does I/O, and every other datasource + # would queue behind it. engine = _build_engine( datasource=snapshot, connection_string=connection_string, ) with _cache_lock: winner = _engine_cache.get(cache_key) if winner is not None: - # Another caller built the same engine while we were constructing. - # Converge on theirs so one key never backs two live pools. + # Someone else built it first; converge so one key never backs + # two live pools. _engine_cache.move_to_end(cache_key) loser, engine = engine, winner else: @@ -333,15 +298,11 @@ def invalidate_engine(datasource: DatasourceConfig) -> bool: """Drop and dispose the cached engine for ``datasource``. Returns whether one was cached. - Meant for the case where the *credentials* an engine was built with have - stopped working — a revoked OAuth grant, a rotated service-account key. - Those engines are poisoned for good: the credential object is baked in at - construction, so every retry through the cache fails identically until - someone throws the engine away. Retrying a transient network blip, by - contrast, wants the pool kept. - - ``get_engine`` rebuilds on the next call, picking up whatever credentials - the datasource now carries. + For credentials that have stopped working — revoked grant, rotated key. + Such engines are poisoned permanently (the credentials are baked in at + construction), so retrying through the cache fails identically until one is + thrown away; a transient blip, by contrast, wants the pool kept. + ``get_engine`` then rebuilds from whatever the datasource now carries. """ connection_string = datasource.get_connection_string() cache_key = _cache_key(datasource=datasource, connection_string=connection_string) @@ -356,10 +317,9 @@ def invalidate_engine(datasource: DatasourceConfig) -> bool: def reset_cache(*, dispose: bool = False) -> None: """Discard every cached engine. - ``dispose`` defaults to False to preserve the long-standing test-fixture - behaviour of dropping references without touching pools. Hosts tearing a - process down want ``dispose=True`` so server-side connections close - promptly instead of waiting on garbage collection. + ``dispose`` defaults to False, preserving the test-fixture behaviour of + dropping references only. Pass True when tearing a process down, so + server-side connections close promptly rather than at GC. """ with _cache_lock: dropped = list(_engine_cache.items()) diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py index 58aa83b3..4307bdd9 100644 --- a/tests/dialects/test_bigquery.py +++ b/tests/dialects/test_bigquery.py @@ -668,12 +668,9 @@ def test_build_engine_with_non_object_credentials_json_raises(payload: str) -> N # --------------------------------------------------------------------------- -# build_engine — per-end-user OAuth grant -# -# ``sqlalchemy-bigquery`` routes every credentials kwarg it has to -# ``service_account.Credentials``, so OAuth user grants have to go through -# its ``user_supplied_client`` escape hatch instead. These tests pin that -# wiring and the mutual exclusion with the service-account path. +# build_engine — per-end-user OAuth grant. Every credentials kwarg the driver +# has routes to service_account.Credentials, so grants go through its +# user_supplied_client escape hatch; these pin that wiring. # --------------------------------------------------------------------------- @@ -699,9 +696,8 @@ def _oauth_ds(name: str = "bq", **overrides) -> DatasourceConfig: def test_build_engine_oauth_uses_user_supplied_client() -> None: - """OAuth path must set the ``user_supplied_client`` URL flag AND pass the - client via ``connect_args``. Without the flag, ``sqlalchemy-bigquery`` - builds an ADC client first (and raises where no ADC exists).""" + """Needs both the ``user_supplied_client`` flag and the client in + ``connect_args``; without the flag the driver builds an ADC client first.""" dialect = BigqueryDialect() captured: dict = {} @@ -733,8 +729,8 @@ def fake_create_engine(url, **kwargs): def test_build_engine_oauth_without_project_raises() -> None: - """An OAuth grant carries no project, so omitting it from the connection - string is a config error rather than a confusing downstream 404.""" + """A grant carries no project, so omitting it is a config error rather than + a confusing downstream 404.""" dialect = BigqueryDialect() ds = _oauth_ds() with pytest.raises(ValueError, match="must be given in the connection string"): @@ -742,8 +738,7 @@ def test_build_engine_oauth_without_project_raises() -> None: def test_build_engine_oauth_falls_back_to_quota_project() -> None: - """A grant carrying ``quota_project_id`` supplies the project when the - connection string doesn't.""" + """``quota_project_id`` supplies the project when the URL doesn't.""" dialect = BigqueryDialect() with ( patch("slayer.sql.dialects.bigquery.sa.create_engine", return_value=object()), @@ -757,8 +752,8 @@ def test_build_engine_oauth_falls_back_to_quota_project() -> None: def test_build_engine_rejects_both_credential_kinds() -> None: - """Setting both is a config error, not a silent precedence win: guessing - is how a per-user query quietly runs as the shared service account.""" + """Guessing between the two is how a per-user query quietly runs as the + shared service account.""" ds = DatasourceConfig( name="bq", type="bigquery", @@ -771,8 +766,8 @@ def test_build_engine_rejects_both_credential_kinds() -> None: def test_build_engine_rejects_oauth_grant_in_credentials_json() -> None: - """An authorized-user grant in ``credentials_json`` cannot work — the - driver hands it to ``from_service_account_info``. Say so up front.""" + """The driver hands ``credentials_json`` to ``from_service_account_info``, + so a grant there cannot work. Say so up front.""" ds = DatasourceConfig( name="bq", type="bigquery", credentials_json=json.dumps(_oauth_info()), ) @@ -801,8 +796,7 @@ def test_build_engine_oauth_malformed_raises(payload: str, message: str) -> None def test_credential_fingerprint_empty_without_credentials() -> None: - """ADC datasources keep the empty fingerprint, so their cache key shape - is unchanged.""" + """ADC datasources keep the empty fingerprint.""" ds = DatasourceConfig(name="bq", type="bigquery", database="p") assert BigqueryDialect().credential_fingerprint(ds) == "" @@ -830,8 +824,7 @@ def test_credential_fingerprint_differs_between_oauth_and_service_account() -> N def test_credential_fingerprint_stable_across_token_refresh() -> None: - """A refreshed access token is the same user. Keying on it would mint — - and leak — a fresh engine on every refresh.""" + """Same user. Keying on the token would leak a fresh engine per refresh.""" dialect = BigqueryDialect() before = dialect.credential_fingerprint(_oauth_ds(token="access-1", expiry="2026-01-01")) after = dialect.credential_fingerprint(_oauth_ds(token="access-2", expiry="2026-01-02")) @@ -839,8 +832,8 @@ def test_credential_fingerprint_stable_across_token_refresh() -> None: def test_credential_fingerprint_keeps_token_when_no_refresh_token() -> None: - """Without a refresh token the access token IS the whole identity, so it - must stay in the digest or two users collide on one engine.""" + """With no refresh token the access token is the whole identity, so it must + stay in the digest or two users collide.""" dialect = BigqueryDialect() info = _oauth_info() info.pop("refresh_token") @@ -853,18 +846,15 @@ def ds_for(token: str) -> DatasourceConfig: def test_credential_fingerprint_leaks_no_secret_material() -> None: - """The fingerprint lands in an in-memory cache key and log lines; it must - not be reversible to the grant.""" + """It lands in cache keys and logs, so it must not be reversible.""" fp = BigqueryDialect().credential_fingerprint(_oauth_ds()) for secret in ("rtok-alice", "csecret", "access-token-1"): assert secret not in fp def test_credential_fingerprint_tolerates_malformed_oauth_json() -> None: - """The fingerprint runs on every cache-key lookup, so a stored grant that - won't parse has to yield a digest rather than raise — otherwise a bad - datasource breaks engine lookup instead of reaching ``build_engine``'s - clear error.""" + """Runs on every cache-key lookup, so an unparseable grant must yield a + digest rather than raise — ``build_engine`` is where it earns its error.""" ds = DatasourceConfig( name="bq", type="bigquery", oauth_credentials_json="not json at all", ) @@ -882,8 +872,8 @@ def ds_for(payload: str) -> DatasourceConfig: def test_build_engine_oauth_validates_before_importing_optional_driver() -> None: - """Config errors must surface as themselves even where the optional - 'bigquery' extra is absent, so validation precedes the google.* imports.""" + """Config errors must surface as themselves even without the optional + 'bigquery' extra, so validation precedes the google.* imports.""" ds = DatasourceConfig(name="bq", type="bigquery", oauth_credentials_json="not json") dialect = BigqueryDialect() with ( diff --git a/tests/test_engine_factory.py b/tests/test_engine_factory.py index 468b360d..e03976b9 100644 --- a/tests/test_engine_factory.py +++ b/tests/test_engine_factory.py @@ -240,9 +240,8 @@ def test_sql_client_uses_engine_factory_for_engine_creation(self) -> None: class TestCredentialKeying: - """The cache key's credential leg. Without it, two callers whose only - difference is *who they authenticate as* share one engine — and one - silently runs the other's queries under the wrong identity.""" + """The credential leg of the cache key. Without it, two callers who differ + only in *who they authenticate as* share one engine.""" @staticmethod def _bq(name: str, credentials_json: str | None) -> DatasourceConfig: @@ -280,9 +279,8 @@ def test_query_engine_agrees_with_factory(self) -> None: class TestCacheBounding: - """Per-identity keys make cache cardinality track *users*, not - datasources, so the cache has to be bounded and evictions must actually - release the pooled connections.""" + """Per-identity keys make cardinality track users, not datasources — so the + cache must be bounded and evictions must release their pools.""" @staticmethod def _lite(n: int) -> DatasourceConfig: @@ -301,8 +299,8 @@ def test_cache_evicts_least_recently_used_over_limit(self, monkeypatch) -> None: engine_factory.reset_cache() def test_reuse_refreshes_recency(self, monkeypatch) -> None: - """A hit must move the entry to the MRU end, otherwise the cap - degenerates into FIFO and evicts the hottest engine.""" + """A hit must move the entry to the MRU end, else the cap degenerates + into FIFO and evicts the hottest engine.""" engine_factory.reset_cache() monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "2") first, second, third = (self._lite(i) for i in range(3)) @@ -326,8 +324,7 @@ def test_eviction_disposes_the_engine(self, monkeypatch) -> None: engine_factory.reset_cache() def test_dispose_failure_does_not_break_caching(self, monkeypatch) -> None: - """A pool that refuses to close must not take the whole factory with - it — the new engine still has to reach the caller.""" + """A pool that refuses to close must not take the factory with it.""" engine_factory.reset_cache() monkeypatch.setenv(engine_factory.MAX_CACHED_ENGINES_ENV, "1") first, second = self._lite(0), self._lite(1) @@ -337,9 +334,8 @@ def test_dispose_failure_does_not_break_caching(self, monkeypatch) -> None: engine_factory.reset_cache() def test_lowering_the_limit_trims_on_the_next_hit(self, monkeypatch) -> None: - """A cache that only trims on insert stays oversized until the next - miss. Re-applying the cap on hits makes the new bound take effect on - the next call.""" + """Trimming only on insert leaves the cache oversized until the next + miss.""" engine_factory.reset_cache() sources = [self._lite(i) for i in range(4)] for ds in sources: @@ -358,8 +354,7 @@ def test_lowering_the_limit_trims_on_the_next_hit(self, monkeypatch) -> None: engine_factory.reset_cache() def test_hit_trim_disposes_outside_the_lock(self, monkeypatch) -> None: - """Trimmed engines are disposed after ``_cache_lock`` is released — - ``dispose()`` closes sockets and must not run under the lock.""" + """``dispose()`` does I/O, so it must not run under the lock.""" engine_factory.reset_cache() cold, hot = self._lite(0), self._lite(1) engine_factory.get_engine(cold) @@ -377,9 +372,8 @@ def record_lock_state(*, engine, reason): engine_factory.reset_cache() def test_zero_limit_bypasses_reuse_on_a_hit(self, monkeypatch) -> None: - """With caching off, a previously-cached key must not be served from - the cache — the entry is dropped and a fresh engine built, rather than - handing back one we are about to dispose.""" + """With caching off, an already-cached key must be dropped and rebuilt, + not handed back moments before we dispose it.""" engine_factory.reset_cache() ds = self._lite(0) first = engine_factory.get_engine(ds) @@ -404,9 +398,8 @@ def test_zero_limit_disables_caching(self, monkeypatch) -> None: class TestInvalidateEngine: - """Credentials baked into an engine can be revoked out from under it. - Those engines are poisoned permanently, so retrying through the cache - reproduces the failure forever unless something evicts them.""" + """Credentials baked into an engine can be revoked out from under it, and + retrying through the cache then fails forever unless something evicts.""" @staticmethod def _lite() -> DatasourceConfig: @@ -435,9 +428,8 @@ def test_next_get_engine_rebuilds(self) -> None: class TestLogSafety: - """Cache keys carry the connection string, which for username/password - dialects is rendered with the password in plaintext. None of it may reach - a log line.""" + """Cache keys carry the connection string, password and all. None of it may + reach a log line.""" @staticmethod def _pg_with_password(password: str) -> DatasourceConfig: @@ -458,8 +450,8 @@ def test_loggable_key_hides_the_connection_string(self) -> None: assert fragment not in rendered def test_loggable_key_is_stable_and_distinguishing(self) -> None: - """Useful for correlating log lines: same key -> same id, different - credentials -> different id.""" + """Same credentials -> same id, different -> different, so log lines + stay correlatable.""" alice = self._pg_with_password("alice-pw") # NOSONAR(S2068) — test fixture # A separate object carrying the same values: the id must follow the # credentials, not the identity of the config object. @@ -475,9 +467,8 @@ def log_id(ds): assert log_id(alice) != log_id(bob) def test_reset_disposal_reason_carries_no_credentials(self) -> None: - """``reset_cache(dispose=True)`` builds its reason string from the - cache key; ``_dispose_quietly`` writes that into a warning when - ``dispose()`` raises.""" + """``reset_cache(dispose=True)`` builds its reason from the cache key, + and ``_dispose_quietly`` logs that when ``dispose()`` raises.""" secret = "reset-time-secret" # NOSONAR(S2068) — test fixture engine_factory.reset_cache() engine_factory.get_engine(self._pg_with_password(secret)) @@ -509,7 +500,7 @@ def test_reset_disposes_only_when_asked(self) -> None: class TestCacheConcurrency: """Engines are reached from worker threads as well as the event loop, so - the cache's read/bump/insert sequences need one lock around them.""" + read/bump/insert needs one lock around it.""" @staticmethod def _lite(n: int) -> DatasourceConfig: @@ -519,10 +510,10 @@ def _lite(n: int) -> DatasourceConfig: def _join_all(threads: list[threading.Thread], *, timeout: float = 10.0) -> None: """Join every worker and fail if any is still running. - Without this, ``join(timeout=...)`` silently returns on a deadlocked - worker and the test passes — which is the one failure mode adding a - lock introduces. Workers are daemons so a genuine deadlock trips this - assertion instead of hanging the whole pytest process. + ``join(timeout=...)`` returns silently on a deadlocked worker, so + without this the test would pass on the one failure mode a lock + introduces. Workers are daemons, so a deadlock trips the assertion + rather than hanging pytest. """ for thread in threads: thread.join(timeout=timeout) @@ -530,8 +521,8 @@ def _join_all(threads: list[threading.Thread], *, timeout: float = 10.0) -> None assert not stuck, f"workers still running after {timeout}s (deadlock?): {stuck}" def test_concurrent_misses_yield_one_shared_engine(self) -> None: - """Two threads missing on the same key must converge on one pool, and - the losing engine must be disposed rather than orphaned.""" + """Two threads missing on one key must converge on a single pool, and + the loser must be disposed rather than orphaned.""" engine_factory.reset_cache() ds = self._lite(0) built: list[sa.Engine] = [] @@ -606,9 +597,8 @@ def hammer(fn): class TestConfigSnapshot: - """``get_engine`` snapshots the config before deriving anything from it. - A ``DatasourceConfig`` is mutable, and engine construction sits between the - cache key and the dialect's second read of the credentials.""" + """``DatasourceConfig`` is mutable, and the build sits between the cache key + and the dialect's second read of the credentials — hence the snapshot.""" @staticmethod def _oauth_ds(refresh_token: str) -> DatasourceConfig: @@ -623,9 +613,8 @@ def _oauth_ds(refresh_token: str) -> DatasourceConfig: ) def test_rotation_mid_build_cannot_desync_key_from_engine(self) -> None: - """A grant refresh landing while the engine is under construction must - not leave that engine cached under a fingerprint describing the *other* - set of credentials — the confusion the credential leg exists to stop.""" + """A refresh landing mid-build must not leave the engine cached under a + fingerprint describing the *other* credentials.""" engine_factory.reset_cache() ds = self._oauth_ds("before") key_for_before = engine_factory._cache_key( @@ -651,8 +640,8 @@ def rotate_then_build(*, datasource, connection_string): engine_factory.reset_cache() def test_caller_mutation_does_not_leak_into_the_cached_engine(self) -> None: - """Mutating the config after the call is self-correcting: the next - lookup keys off the new credentials, misses, and rebuilds.""" + """Mutating after the call is self-correcting: the next lookup keys off + the new credentials, misses, and rebuilds.""" engine_factory.reset_cache() ds = self._oauth_ds("before") with patch.object( @@ -667,9 +656,8 @@ def test_caller_mutation_does_not_leak_into_the_cached_engine(self) -> None: engine_factory.reset_cache() def test_snapshot_is_detached_from_the_callers_object(self) -> None: - """Sanity-check the copy depth: every field is scalar, so a shallow - model_copy already detaches. If a mutable field is ever added, this is - where the assumption breaks.""" + """Every field is scalar, so a shallow copy already detaches. Add a + mutable field and this is where the assumption breaks.""" ds = self._oauth_ds("before") snapshot = ds.model_copy() ds.oauth_credentials_json = "mutated" diff --git a/tests/test_sql_client.py b/tests/test_sql_client.py index f3093db5..73986151 100644 --- a/tests/test_sql_client.py +++ b/tests/test_sql_client.py @@ -450,9 +450,8 @@ def test_none_db_type_uses_limit(self) -> None: class TestIsAuthFailure: - """Credential rejection is classified separately from transient errors: - retrying it is pointless (the credentials are baked into the engine), and - the right response is to throw the engine away.""" + """Credential rejection is classified apart from transient errors: retrying + is pointless, so the engine gets thrown away instead.""" def test_oauth_invalid_grant_is_auth_failure(self) -> None: assert _is_auth_failure(Exception("('invalid_grant: Token has been expired or revoked.')")) @@ -461,7 +460,7 @@ def test_libpq_password_failure_is_auth_failure(self) -> None: assert _is_auth_failure(Exception('FATAL: password authentication failed for user "svc"')) def test_signal_found_through_sqlalchemy_orig(self) -> None: - """Drivers surface wrapped; the signal is a layer or two down.""" + """Drivers surface wrapped; the signal sits a layer or two down.""" inner = Exception("invalid_grant") wrapped = sqlalchemy.exc.OperationalError("SELECT 1", {}, inner) assert _is_auth_failure(wrapped) @@ -473,8 +472,8 @@ def test_signal_found_through_cause_chain(self) -> None: assert _is_auth_failure(outer) def test_google_refresh_error_matched_by_type_name(self) -> None: - """google-auth ships only with the optional 'bigquery' extra, so the - classifier matches on class name rather than importing it.""" + """google-auth ships only with the optional extra, so the classifier + matches on class name.""" class RefreshError(Exception): pass assert _is_auth_failure(RefreshError("bad news")) @@ -484,8 +483,7 @@ def test_transient_errors_are_not_auth_failures(self) -> None: assert not _is_auth_failure(Exception(message)), message def test_table_permission_denied_is_not_an_auth_failure(self) -> None: - """The credentials worked; the grant didn't. Evicting a healthy engine - over this is pure pool churn.""" + """The credentials worked; the grant didn't. Evicting is pool churn.""" assert not _is_auth_failure(Exception("permission denied for table orders")) def test_cyclic_cause_chain_terminates(self) -> None: @@ -530,8 +528,8 @@ async def test_non_auth_failure_keeps_the_engine(self) -> None: assert client._sync_engine is engine async def test_async_engine_is_disposed_too(self) -> None: - """Native-async dialects hold a second pool that ``invalidate_engine`` - knows nothing about; it has to go as well.""" + """Native-async dialects hold a second pool ``invalidate_engine`` knows + nothing about.""" client = self._client() async_engine = AsyncMock() client._async_engine = async_engine @@ -545,8 +543,8 @@ async def test_async_engine_is_disposed_too(self) -> None: assert client._async_engine is None async def test_get_column_types_also_discards(self) -> None: - """It runs its own SQL against the same cached engine, so it needs the - same cleanup ``execute`` gets.""" + """Runs its own SQL on the same cached engine, so it needs the same + cleanup ``execute`` gets.""" client = self._client() client._sync_engine = object() with ( @@ -562,8 +560,8 @@ async def test_get_column_types_also_discards(self) -> None: assert client._sync_engine is None def test_execute_sync_also_discards(self) -> None: - """The sync path shares the factory-cached engine. It cannot dispose an - async pool (no loop to do it on), so it only drops the sync one.""" + """Shares the factory-cached engine, but has no loop to dispose an async + pool on — so it drops only the sync one.""" client = self._client() client._sync_engine = object() with ( @@ -591,8 +589,7 @@ def test_execute_sync_keeps_engine_on_non_auth_failure(self) -> None: assert client._sync_engine is engine async def test_cleanup_failure_does_not_mask_the_original_error(self) -> None: - """The auth error is what the caller needs to see; a failed eviction - must not displace it.""" + """A failed eviction must not displace the auth error.""" client = self._client() with ( patch.object(sql_client.SlayerSQLClient, "_execute", side_effect=Exception("invalid_grant")), From 9698ae82dd86a8f6ddff7e41a7d2287bba251069 Mon Sep 17 00:00:00 2001 From: AivanF Date: Tue, 11 Aug 2026 13:25:38 +0300 Subject: [PATCH 5/5] Neater doc-strings --- slayer/sql/dialects/base.py | 26 ++++++-------------------- slayer/sql/engine_factory.py | 20 ++++++++------------ 2 files changed, 14 insertions(+), 32 deletions(-) diff --git a/slayer/sql/dialects/base.py b/slayer/sql/dialects/base.py index e5cfdadb..9cda160a 100644 --- a/slayer/sql/dialects/base.py +++ b/slayer/sql/dialects/base.py @@ -169,11 +169,8 @@ def _sqlglot_backslash_escapes(sqlglot_name: str) -> bool: def _digest(secret: str | None) -> str: - """Short, stable, non-reversible id for secret material used in cache keys. - - Truncated to 16 hex chars: collision risk is negligible for the number of - live credentials in a process, and short keys stay readable in logs. - """ + """Non-reversible id for secret material in cache keys. 16 hex chars keeps + it log-readable; collisions are negligible at this scale.""" if not secret: return "" return hashlib.sha256(secret.encode("utf-8")).hexdigest()[:16] @@ -612,21 +609,10 @@ def build_engine( def credential_fingerprint(self, datasource: "DatasourceConfig") -> str: """Opaque identity of the credentials this datasource authenticates with. - Engines are cached per ``(connection_string, runtime_fingerprint, - credential_fingerprint)``. Any dialect whose secret does **not** appear - in the connection string MUST override this, or two callers holding - different credentials for the same URL will silently share one engine — - the first caller's identity then serves everyone for the life of the - process. - - The default covers the common case safely: username/password dialects - embed their credentials in the URL, so the connection string already - distinguishes them and ``""`` adds nothing. ``credentials_json`` is - hashed here rather than left out, so a dialect that carries it - out-of-band is keyed correctly even before it overrides this. - - Return a digest or a stable subject id — **never** raw secret material, - since cache keys reach logs and error messages. + Part of the engine cache key: a dialect whose secret is *not* in the + connection string MUST override this, or callers with different + credentials share one engine. Return a digest, never raw secret — + keys reach logs. """ return _digest(datasource.credentials_json) diff --git a/slayer/sql/engine_factory.py b/slayer/sql/engine_factory.py index 87fed6b8..e5a99a3f 100644 --- a/slayer/sql/engine_factory.py +++ b/slayer/sql/engine_factory.py @@ -82,11 +82,10 @@ def _max_cached_engines() -> int: def loggable_key(key: EngineCacheKey) -> str: - """Short, stable, log-safe id for a cache key. + """Log-safe id for a cache key. - ``key[0]`` is the connection string, rendered with ``hide_password=False`` - — so the raw key must never reach a log line. The digest stays correlatable - across lines without being reversible. + ``key[0]`` renders with ``hide_password=False``, so the raw key must never + reach a log line. The digest stays correlatable but not reversible. """ return _digest("\x00".join(key)) @@ -119,15 +118,12 @@ def _take_evictions_over_limit() -> list[sa.Engine]: def _cache_key(datasource: DatasourceConfig, connection_string: str) -> EngineCacheKey: - """Cache identity for ``datasource``: URL + runtime fields + credentials. + """Cache identity: URL + runtime fields + credentials. - Kept in one place because ``query_engine._sql_client_cache_key`` must agree - with it; a divergence between the two caches means a caller can get a client - whose engine was built for different credentials. - - Only ``get_engine`` snapshots ``datasource`` first — it is the one caller - with a slow build between the key and the credentials it describes. Callers - that use the key immediately just miss and rebuild. + Shared with ``query_engine._sql_client_cache_key`` — if the two diverged, a + caller could get a client whose engine was built for other credentials. + Only ``get_engine`` snapshots ``datasource`` first; callers that use the key + immediately just miss and rebuild. """ return ( connection_string,