diff --git a/deploy/agami.env.example b/deploy/agami.env.example index 2d23dd04..bce775d6 100644 --- a/deploy/agami.env.example +++ b/deploy/agami.env.example @@ -42,6 +42,13 @@ AGAMI_ADMIN_PASSWORD= # Override the bundled-Postgres password if you like (it's never exposed, so the default is fine). # POSTGRES_PASSWORD= +# Optional: how long ONE statement may run, in whole SECONDS (unset = 30). This is the deployment's +# availability budget — a query that outruns it is cancelled on the warehouse and the caller is told +# to narrow it. Raise it if your warehouse is slow and your users would rather wait; lower it to keep +# a runaway query from holding a connection. The executor's other time bounds are derived from this +# one, so this is the single number to tune. +# AGAMI_SQL_TIMEOUT_S=30 # per-statement budget (default 30 seconds) + # Optional free registration key (best-effort; absent = full run). # AGAMI_LICENSE_KEY= diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 9a5b51d2..3374f44e 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -52,16 +52,18 @@ import argparse import configparser +import contextlib import csv import json import logging import os import stat import sys +import threading import urllib.parse import uuid -from collections.abc import Callable -from contextvars import ContextVar +from collections.abc import Callable, Iterator +from contextvars import ContextVar, copy_context from dataclasses import asdict, dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -70,6 +72,7 @@ from guardrail import ( RULE_MODEL_SAFETY, RULE_MODEL_UNAVAILABLE, + RULE_RESOURCE_LIMIT, Envelope, Failure, FailureKind, @@ -483,21 +486,70 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Postgres connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # Bound outside the try so the `finally` can close it whether or not the declare got that far. + cur = None try: + # `with conn` is the TRANSACTION, not the connection: leaving it by an exception rolls back, + # which is why it stays even though the cursor no longer lives inside a `with` of its own. with conn: + # `SET LOCAL statement_timeout` is the native server-side backstop. It is + # TRANSACTION-scoped, so it has to run on THIS transaction and before the named cursor + # declares — a regular client-side cursor, because a named cursor may only ever run the + # one query it was declared for. The libpq `options` startup parameter would be the + # other way to set it and is deliberately not used: a transaction-mode connection pooler + # can reject an unknown startup parameter, which would break the connect outright rather + # than bound the statement. + with conn.cursor() as bound_cur: + bound_cur.execute( + "SET LOCAL statement_timeout = %s", + ((timeout_s + _NATIVE_BOUND_SKEW_S) * 1000,), # the setting is in milliseconds + ) # A server-side (named) cursor so the row cap bounds TRANSFER, not just what we write: # psycopg2's default client-side cursor buffers the ENTIRE result before we can fetchmany, # so a runaway result would still be pulled whole. The named cursor streams from the - # server in bounded batches (ACE-038). Read-only SELECTs (the only thing the guard admits) + # server in bounded batches. Read-only SELECTs (the only thing the guard admits) # are exactly what a server-side cursor supports. - with conn.cursor(name="agami_bounded") as cur: - cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window - cur.execute(sql) - result = _collect_cursor(cur) + cur = conn.cursor(name="agami_bounded") + cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window + # `connection.cancel()` is psycopg2's cancel: it opens a second connection and sends the + # libpq cancel request, so it is safe to call from the watchdog thread while this one is + # blocked in the driver. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Postgres execution error: {e}", code=5) finally: - conn.close() + # The named cursor is closed HERE, by hand, rather than by a `with` around it. Closing one + # sends `CLOSE agami_bounded` to the server, and when we are unwinding on the timeout marker + # the transaction is already aborted, so that statement raises in turn — from inside + # `__exit__`, where a raised exception REPLACES the one being propagated. The refusal would + # reach the chokepoint as an ordinary driver error and be reported as a failure. Swallowing + # it costs nothing: by this point the transaction has ended and the server-side portal is + # gone with it, so the close is a courtesy rather than the thing that frees the resource. + if cur is not None: + try: + cur.close() + except Exception: + pass + # And the connection close is guarded for the same reason, one line down: an exception raised + # inside a `finally` REPLACES the one propagating through it, so an unguarded close here would + # destroy the marker the `except _ResourceLimit: raise` above just re-raised — and the caller + # would read a bound we imposed as an unclassified server break. + try: + conn.close() + except Exception: + pass return result @@ -520,14 +572,38 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"MySQL connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: with conn.cursor() as cur: - cur.execute(sql) - result = _collect_cursor(cur) + # pymysql's connection has NO `cancel()`. Its two cancel-shaped methods are `close()`, + # which sends COM_QUIT down the very socket the blocked statement owns — and so can + # itself block on a connection that is already stuck — and `_force_close()`, which + # closes the socket outright. Only the second one actually unblocks a statement in + # flight, so it is named here explicitly rather than reached for by shape. + with _deadline(conn._force_close, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: - conn.close() + # Guarded because an exception raised inside a `finally` REPLACES the one propagating through + # it, which would destroy the marker re-raised just above. MySQL is the most exposed of the + # ten: its cancel is `_force_close()`, which deliberately destroys the socket that this + # `close()` then tries to write COM_QUIT to — so on exactly the timeout path this is the close + # most likely to raise. + try: + conn.close() + except Exception: + pass return result @@ -548,11 +624,20 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: "Add one to /local/credentials.", code=2, ) + # Resolved before the connect, because the native backstop is a SESSION parameter and so has to + # be handed to the connect call itself. + timeout_s = _resolve_timeout_s() conn_kwargs: dict[str, Any] = { "account": creds["account"], "user": creds["user"], "client_session_keep_alive": False, "login_timeout": 15, + # Snowflake's native server-side bound, set behind our watchdog by the usual skew so the + # watchdog wins and the flag stays the only classification. On a warehouse the backstop is + # worth more than elsewhere: a statement nobody is waiting for still bills credits. + "session_parameters": { + "STATEMENT_TIMEOUT_IN_SECONDS": timeout_s + _NATIVE_BOUND_SKEW_S, + }, } for k in ("password", "warehouse", "database", "schema", "role", "authenticator"): if creds.get(k): @@ -563,8 +648,32 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"Snowflake connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + + def cancel_snowflake_statement() -> None: + """Ask Snowflake to abort the statement this cursor is running. + + Neither the connection nor the cursor has a `cancel()` — verified against the connector's + own source, where the only public abort is `SnowflakeCursor.abort_query(qid)` and the + connection's cancel helpers are private. The query id is what identifies the statement to + abort, and it exists only once the statement has been submitted; before that there is + nothing to abort and the session parameter above is what bounds the call. + """ + qid = cur.sfqid + if qid: + cur.abort_query(qid) + + with _deadline(cancel_snowflake_statement, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Snowflake execution error: {e}", code=5) finally: @@ -634,9 +743,20 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"BigQuery client init failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # `job_timeout_ms` is BigQuery's native server-side bound, and on this engine it is the ONLY + # bound: there is no watchdog cancel here, deliberately. BigQuery hands out no connection or + # cursor to cancel, and the call that blocks is `job.result()`, which is reached only AFTER + # `client.query()` has returned — so at the instant a watchdog would fire there is nothing in + # hand to stop. Stated plainly, and accepted rather than papered over: on BigQuery a client-side + # stall comes back as a `failed` envelope, not a `resource_limit` refusal. The query itself is + # still stopped by the bound below, which is what keeps a runaway from scanning on unattended. + # The usual skew is kept so this engine's number matches every other engine's. + job_config_kwargs: dict[str, Any] = { + "job_timeout_ms": (timeout_s + _NATIVE_BOUND_SKEW_S) * 1000, + } # If `dataset` was set, prefix unqualified table references via the # default_dataset job config so the SQL can omit `..` - job_config_kwargs: dict[str, Any] = {} if creds.get("dataset"): try: job_config_kwargs["default_dataset"] = f"{project}.{creds['dataset']}" @@ -645,15 +765,17 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: cap = _resolve_row_cap() try: - if job_config_kwargs: - job_config = bigquery.QueryJobConfig(**job_config_kwargs) - job = client.query(sql, job_config=job_config) - else: - job = client.query(sql) + job_config = bigquery.QueryJobConfig(**job_config_kwargs) + job = client.query(sql, job_config=job_config) # BigQuery has no DB-API cursor, so it can't funnel through `_collect_cursor`; apply the # same bounded-fetch cap here. `max_results=cap+1` bounds what the API returns (transfer), # and the (cap+1)th row flags truncation — the never-silent guarantee holds for BigQuery too. results = job.result(max_results=cap + 1) # waits for completion; raises on error + except _ResourceLimit: + # Nothing in this engine raises the marker today, and this clause is still not optional: the + # rule is that NO engine may relabel it as a driver error, and the clause below would. It is + # what makes the rule checkable across all ten engines rather than nine. + raise except Exception as e: raise ExecutorError(f"BigQuery execution error: {e}", code=5) @@ -672,6 +794,12 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: + """Tier-3 path for SQLite using the stdlib `sqlite3` module. + + The first engine to run under the per-statement deadline. `sqlite3.Connection.interrupt()` is a + genuine cancel — it stops the statement mid-scan from another thread — so this engine can prove + the whole refusal contract in-process, with no network and no fixture warehouse. + """ import sqlite3 # always available in stdlib _require(creds, "path") path = os.path.expanduser(creds["path"]) @@ -679,14 +807,48 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: conn = sqlite3.connect(path) except Exception as e: raise ExecutorError(f"SQLite connect failed: {e}", code=4) + # Resolved ONCE for the call, so the budget the watchdog enforces and the number the refusal + # quotes cannot be two different values. + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The deadline covers the FETCH as well as the execute. `_collect_cursor` pulls `cap + 1` + # rows in a single `fetchmany`, and on a cursor that streams its result that pull is where + # the scan actually happens — a clock that stopped at `execute` would bound the cheap half + # of the work and leave the expensive half unbounded. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + # A cancelled statement raises `sqlite3.OperationalError("interrupted")`, whose text + # is not a classification we would want to key on. The FLAG is the classification, + # and it is the ONLY one: `_deadline` sets it before the cancel lands, so an error + # raised while it is unset is the database's own — however late it arrives. + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + # Checked after the watchdog is disarmed, so the flag is final. A cancel can also land + # between the two calls above, or just as the second returns, and leave nothing to raise: + # the budget still elapsed, so the outcome is still a refusal rather than a result gathered + # past it. + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + # Ahead of the catch-all below on purpose: our own marker must reach `execute_guarded` + # intact. Wrapped in an `ExecutorError` it would become a `failed`/`syntax` envelope, which + # is the database's outcome rather than the bound we imposed. + raise except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: - conn.close() + # Guarded like every other engine's: an exception raised inside a `finally` REPLACES the one + # propagating through it, so an unguarded close would silently convert the refusal above into + # an unclassified failure. + try: + conn.close() + except Exception: + pass return result @@ -705,10 +867,26 @@ def _run_sqlserver(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"SQL Server connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # pymssql's DB-API `Connection` has no `cancel()` — verified against pymssql 2.3, whose + # connection exposes only close/commit/cursor/rollback/bulk_copy, and whose cursor exposes + # none either. The cancel lives one layer down, on the `_mssql.MSSQLConnection` that + # connection wraps, reachable only as the private `_conn`; that object's `cancel()` sends the + # TDS attention packet, which is the thing that actually stops a statement in flight. + with _deadline(conn._conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"SQL Server execution error: {e}", code=5) finally: @@ -735,10 +913,23 @@ def _run_oracle(creds: dict[str, str], sql: str) -> ExecResult: conn = oracledb.connect(user=creds["user"], password=creds["password"], dsn=dsn) except Exception as e: raise ExecutorError(f"Oracle connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # `oracledb.Connection.cancel()` breaks out of the call in progress on that connection — + # verified against python-oracledb 3.x, where it is on the CONNECTION and not on the cursor. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Oracle execution error: {e}", code=5) finally: @@ -766,10 +957,24 @@ def _run_databricks(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Databricks connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The cancel is on the CURSOR here, not the connection — verified against + # databricks-sql-connector 4.x, whose `Cursor.cancel()` posts a cancel for the operation that + # cursor is running while its connection has none. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Databricks execution error: {e}", code=5) finally: @@ -798,10 +1003,23 @@ def _run_trino(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Trino connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # Trino's cancel is on the CURSOR — verified against the trino client, whose + # `Cursor.cancel()` sends the coordinator a DELETE for the query that cursor started. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Trino execution error: {e}", code=5) finally: @@ -823,9 +1041,25 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: conn = duckdb.connect(path, read_only=True) except Exception as e: raise ExecutorError(f"DuckDB open failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: - cur = conn.execute(sql) - result = _collect_cursor(cur) + # `conn.interrupt()` is DuckDB's cancel, and it is on the connection — which is just as well, + # because on this engine there is no cursor to reach for until the execute has returned: + # `conn.execute` HANDS BACK THE CONNECTION ITSELF, so `cur` below IS `conn`. The deadline + # therefore has to be armed around the execute as well as the fetch, and on an in-process + # engine the execute is where the scan happens. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur = conn.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"DuckDB execution error: {e}", code=5) finally: @@ -836,7 +1070,7 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: return result -_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation (ACE-038) +_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation # Per-call cap from --max-rows (ACE-044). A ContextVar, not a plain global, so it is REQUEST-SCOPED # once the HTTP server runs execution in-process (ACE-028): concurrent handlers run in worker threads # (`run_blocking` copies the context per call, like `_current_org_ctx`), so each request's cap is @@ -860,10 +1094,225 @@ def _resolve_row_cap() -> int: return cap +_DEFAULT_TIMEOUT_S = 30 # wall-clock seconds one statement may run before the watchdog cancels it +# How far BEHIND our own watchdog a NATIVE server-side bound is set, on the three engines that have +# one. The skew is the whole point: our watchdog fires at the budget, the engine's own bound five +# seconds later, so the watchdog always wins the race and its Event stays the SOLE classification +# signal. Reverse the order and a server-side kill would beat us to the statement, the flag would be +# clear, and the refusal would come back as an ordinary database failure. The native bound is a +# backstop for the one case the watchdog cannot cover — our process dying with a statement in +# flight, which would otherwise leave the engine scanning for nobody. +_NATIVE_BOUND_SKEW_S = 5 +# How far behind the watchdog the OUTER bound sits — the one `execute_guarded` puts around +# `executor.execute` itself, so an INJECTED executor is bounded too. The four bounds are one ordered +# family resolved from the same budget, innermost first: +# +# watchdog (timeout_s) < native (+5) < outer (+10) < supervisor (+60) +# +# and the order is the contract. The inner watchdog must always win, because it is the only layer +# that can attribute the stop to the statement and hand the caller the precise refusal; every layer +# behind it is a backstop for a failure the layer in front of it cannot see. Collapse the order and +# a bound we imposed comes back as something else — a native server-side kill reads as an ordinary +# database error, and a supervisor kill reads as `failed`/`timeout` naming nothing the caller can +# act on. Both are strictly worse answers to the same event. +_OUTER_BOUND_SKEW_S = 10 +# The OUTERMOST bound: how far behind the watchdog the subprocess supervisor in `tools` stops waiting +# for a forked child. It lives here, next to its three siblings, so the whole family is derived from +# one resolver and one place — the supervisor was a hardcoded 240s, which for any statement budget +# approaching it became the FIRST bound to fire and inverted the order. Its slack is large because it +# bounds a whole process (interpreter start, model load, credential resolution, connect) rather than +# a statement. +_SUPERVISOR_SKEW_S = 60 +# How many worker threads `_execute_bounded` may have given up on at once, process-wide. That layer +# holds nothing it can cancel, so every expiry abandons a thread which still occupies a pooled +# connection and, server-side, a running statement. On the built-in path the inner watchdog fires ten +# seconds earlier, so reaching the outer bound means a cancel already failed — rare, and an honest +# cost. On an INJECTED executor there is no inner watchdog at all, so every slow statement abandons +# one, and because the caller's own thread is freed at the bound the abandonments accumulate with no +# ceiling of their own: the anyio worker limiter used to supply that ceiling by blocking the caller, +# and bounding the wait removed it. The cap restores one explicitly. 8 is deliberately far below both +# that former limiter and a typical warehouse connection pool, so a saturated executor loses a +# minority of the pool to work nobody is waiting for rather than taking the datasource away from the +# whole organization. +_MAX_ABANDONED_WORKERS = 8 +_abandoned_lock = threading.Lock() +_abandoned_workers = 0 + + +def _resolve_timeout_s() -> int: + """Effective per-statement timeout, in whole seconds. `AGAMI_SQL_TIMEOUT_S` is the + operator-configurable DEPLOYMENT budget (default 30 when unset) — an operator owns their + availability tradeoff and may set it higher OR lower than 30. A missing or non-positive value + falls back to the default. + + **The environment is the ONLY source, deliberately.** A request-scoped override would outrank it + in the parent and be invisible to a forked child, which re-resolves from `os.environ` alone — so + the supervisor bound the parent derives could sit BELOW the budget the child actually enforces + and fire first, inverting the ordered family the whole design rests on. One source, readable on + both sides of the fork, makes that inversion unrepresentable rather than merely unlikely. + + Unlike `_resolve_row_cap`, a value that is PRESENT and does not survive to become the budget is + logged at warning before the fallback. That covers `45.5` and `30s`, which cannot be read at all, + and equally `-5` and `0`, which can be read and are then declined: an operator who wrote either + asked for something specific, and a deployment quietly running 30 instead is exactly the + invisible degradation the warning exists against. The warning goes to the module logger and never + to stderr, because the subprocess transport parses stderr and an extra line there would break + that contract.""" + raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() + digits = raw[1:] if raw.startswith("-") else raw # a leading minus is a value, not a typo + # `isdecimal`, not `isdigit`: the latter admits `²` and `①`, which `int()` then refuses — turning + # a misconfigured deployment into a ValueError raised out of this resolver, at a call site (the + # fork path's supervisor bound) that sits outside any handler. + written = int(raw) if digits.isdecimal() else None + timeout_s = written if written is not None and written > 0 else _DEFAULT_TIMEOUT_S + if raw and timeout_s != written: + _LOG.warning( + "AGAMI_SQL_TIMEOUT_S=%r is not a usable whole number of seconds; falling back to %ds.", + raw, + _DEFAULT_TIMEOUT_S, + ) + return timeout_s + + +class _ResourceLimit(Exception): + """Raised when our own watchdog fired, so a cancelled statement unwinds the engine function the + same way any other failure does and the surrounding transaction rolls back. It is an internal + marker only: it is always caught and translated inside this module and never crosses the tool + boundary.""" + + +class _OuterBoundExpired(_ResourceLimit): + """Raised when the OUTER bound expired: the executor never returned and we stopped waiting. + + A subclass, so the one `except _ResourceLimit` handler in `execute_guarded` produces exactly one + refusal for either layer and no second one can be minted. It stays distinguishable because the + two events are not the same event: the watchdog CANCELLED a statement it holds a connection to, + while this layer only abandoned its wait — so the sentence the caller reads differs, and saying + "cancelled" here would be false. + """ + + +class _ExecutorSaturated(_ResourceLimit): + """Raised INSTEAD of starting a statement, when the abandoned-worker cap is already reached. + + A sibling of `_OuterBoundExpired` under the same parent for the same reason — one handler, one + refusal — and distinguishable for the same reason too: this statement did not run long, it did + not run at all. What we observed is the executor, not the statement, so the sentence the caller + reads must not blame the query it just sent. + """ + + +# The marker's message, single-sourced because every engine raises it. It is diagnostic text, not +# caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own +# detail, so nothing a caller reads comes from here. +_OUTLIVED_BUDGET = "the statement outlived its per-statement budget" +# The outer marker's message, and diagnostic in the same way — it names the executor rather than the +# statement, because at this layer the statement is not the thing we observed. +_OUTLIVED_OUTER_BOUND = "the executor outlived the outer bound around it" +# The saturation marker's message, diagnostic in the same way again. +_EXECUTOR_SATURATED = "the executor already has its limit of abandoned calls outstanding" + +# The remediation names only what would make THIS statement executable: on the served path the caller +# is an assistant with no shell and no deployment, so naming the environment variable would be advice +# it cannot take, addressed to someone who is not reading. +_NARROW_IT = ("Narrow the time range, reduce the grouping, or add a selective filter, " + "then run it again.") + + +def _resource_limit_refusal(exc: _ResourceLimit) -> Refusal: + """The ONE `resource_limit` refusal, built from whichever of the three bounds raised. + + Single-sourced because two entries into the engines lead here — the guarded chokepoint and the + subprocess/CLI adapter — and a caller must not be able to tell which one it came through. + + The budget is re-resolved rather than carried on the marker: the marker stays a plain exception, + and nothing between the engine call and here can change the environment the resolver reads, so it + is the same number the watchdog used. The configured number belongs in the detail — it is a + deployment setting, not a data value, and a bound the caller cannot see is one it cannot plan + around. + """ + timeout_s = _resolve_timeout_s() + if isinstance(exc, _ExecutorSaturated): + # Nothing ran, so nothing about THIS statement is the finding. Saying "your query was too + # slow" here would be false and would send the caller off simplifying a statement that is + # very possibly fine. + detail = ( + "The executor is saturated: too many earlier calls have not returned, so this " + "statement was not started." + ) + remediation = "Wait for the calls already in flight to finish, then run it again." + elif isinstance(exc, _OuterBoundExpired): + # The outer layer holds no connection, so it stopped WAITING rather than stopping the + # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the + # driver call and the statement may still be running. The caller is told what we actually + # know, and the number quoted is the bound that actually elapsed. + detail = ( + f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " + "and the query was abandoned." + ) + remediation = _NARROW_IT + else: + detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." + remediation = _NARROW_IT + return refuse(RULE_RESOURCE_LIMIT, detail=detail, remediation=remediation) + + +@contextlib.contextmanager +def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threading.Event]: + """Arm a watchdog that calls `cancel` if the wrapped block outlives `timeout_s`, and yield the + `threading.Event` that says whether it fired. + + The event is set BEFORE `cancel` runs. Order matters: whoever catches the driver error that the + cancellation provokes must be able to read an already-set flag and attribute the failure to us + rather than to the database. A `cancel` that raises is swallowed and logged, because some drivers + raise when cancelled from a thread other than the one running the statement, and an exception + escaping a timer thread is both unhandleable by the caller and invisible in the result. + + **The disarm is a JOIN, not a request.** `threading.Timer.cancel()` only sets the timer's internal + `finished` Event, so a `fire` that already passed its own check runs to completion regardless — + setting the flag and delivering a cancel AFTER this block has returned. Both consequences are + real: every engine re-reads the flag once the block has exited and would read a stale `False`, and + a cancel arriving late lands on a connection the engine has moved on from, which on a pooled one + is by then someone else's statement. The lock plus the flag close that: a `fire` that loses the + race neither sets the Event nor calls `cancel`, and a `fire` that wins holds the lock across the + cancel, so the disarm waits for it. The wait is bounded by the layer outside this one — a cancel + that hangs is what the outer bound around the whole executor call exists to survive.""" + fired = threading.Event() + lock = threading.Lock() + # Assigned only below, in the enclosing scope, and only read inside `fire` — so the closure sees + # the current value with no `nonlocal` and no mutable box. + disarmed = False + + def fire() -> None: + with lock: + if disarmed: + return # the block already returned; there is nothing of ours left to stop + fired.set() + try: + cancel() + except Exception as exc: + _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) + + timer = threading.Timer(timeout_s, fire) + timer.daemon = True # a hung cancel must never hold the interpreter open at shutdown + timer.start() + try: + yield fired + finally: + # Claim the race under the lock FIRST, then ask the timer to stand down. `Timer.cancel()` is + # not a join, so doing it the other way round leaves a window between this block returning + # and the flag being set, in which a `fire` that has already expired can still take the lock + # and mark a statement that actually completed. Taking the lock first closes that window: a + # `fire` already holding the lock is waited for, and every other `fire` observes `disarmed`. + with lock: + disarmed = True + timer.cancel() # belt and braces, for a timer that has not started running yet + + def _flag_truncated(cap: int) -> None: """Signal a bounded-fetch truncation to the caller — a non-error `{"truncated": …}` marker on stderr (distinct from the guards' `{"error": …}`), so a truncated result is never mistaken for a - complete one (ACE-038/044). Shared by every engine's materialization path. One write so the + complete one (ACE-044). Shared by every engine's materialization path. One write so the marker is always a single line, even if other notices surround it.""" sys.stderr.write(json.dumps({"truncated": {"row_cap": cap}}) + "\n") @@ -876,7 +1325,7 @@ def _collect_cursor(cur: Any) -> ExecResult: executor path share, so the row cap is enforced once, identically, for every caller. Fetch FIRST, then read ``cur.description``: a psycopg2 **server-side (named) cursor** — which the - Postgres/Redshift path uses to bound transfer (ACE-038) — reports ``description is None`` until the + Postgres/Redshift path uses to bound transfer — reports ``description is None`` until the first fetch, so reading it beforehand would drop EVERY row of a real Postgres result. Client-side cursors (sqlite/mysql/…) set ``description`` at execute, so fetch-first is equally correct there.""" cap = _resolve_row_cap() @@ -969,7 +1418,8 @@ def _resolve_guard_model(profile: str): def _write_refusal(refusal: Refusal) -> None: """Write a guard refusal to stderr in the ONE shape every caller parses — the single wire-writer, - called only from ``main``. + called from ``main`` and from the per-engine CSV adapter, which is the other entry into the + engines. One JSON object, on one line: ``{"refusal": {reason, rule, detail, remediation}}`` — the wire shape S2 established, which ``tools._stderr_refusal`` rebuilds through ``Refusal`` on the parent @@ -1115,9 +1565,18 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa def _emit_or_err(run: Callable[[], ExecResult]) -> int: """Subprocess/CLI adapter over a ``_run_`` function: write its result to stdout as CSV and return exit code 0, or translate an ``ExecutorError`` into the stderr message + exit code the CLI - contract documents (byte-identical to what the old ``_execute_`` emitted).""" + contract documents (byte-identical to what the old ``_execute_`` emitted). + + The watchdog marker gets an arm of its own, AHEAD of the classified one, because this is the + second entry into the engine functions and they raise it. Without it a per-statement timeout + reached through here would escape as a traceback rather than as the exit code this docstring + promises. It renders exactly what ``main`` renders for a refusal — the one JSON object on stderr + and exit 1 — so the two entries cannot disagree about what a bound we imposed looks like.""" try: _emit_result_csv(run()) + except _ResourceLimit as exc: + _write_refusal(_resource_limit_refusal(exc)) + return 1 except ExecutorError as e: return _err(e.msg, code=e.code) return 0 @@ -1250,6 +1709,105 @@ def _envelope( ) +def _execute_bounded( + executor: Executor, sql: str, creds: dict[str, str], *, profile: str +) -> ExecResult: + """Run ``executor.execute`` under the OUTER bound and return what it returned. + + This is the layer that makes "no executor escapes the limit" true. The inner deadline lives + inside the BUILT-IN executor's engine functions, so before this an INJECTED executor — the + hosted connection-reuse path — ran with no per-statement bound at all, and even on the built-in + path BigQuery has no watchdog (it has no connection object to cancel), so a client-side stall + there was unbounded too. Both are bounded here, and deliberately by the same layer: applying + this only to non-built-in executors would leave the BigQuery hole open while looking closed. + + **The mechanism is a worker thread, because this layer holds nothing it can cancel.** An + arbitrary ``Executor`` exposes one method and no connection, so the only thing this code owns is + its own WAIT — it starts the call on a daemon thread and joins with the budget. + + **The cost is an abandoned worker, and it is real.** On expiry the thread is still inside the + driver call, and it stays there: nothing here cancels it, and it may hold a database connection + (and, on the server side, a running statement) until that call returns on its own. It is a + daemon so it cannot hold the interpreter open at exit. This is a bound on how long a CALLER + waits, not a promise that the work stopped — which is exactly why the inner watchdog, the layer + that really can cancel, is set to fire first. + + **And the cost is CAPPED, because bounding the wait removed the ceiling that used to bound it.** + Before this layer existed the same slow call blocked the caller's own worker thread, so the + host's thread limiter capped how much abandoned work could be in flight and applied backpressure + to everything behind it. Returning at the bound frees that slot, so on the injected path — where + there is no inner watchdog and EVERY slow statement abandons one — the abandonments would + otherwise accumulate without limit until the pool, and with it the datasource, belonged entirely + to work nobody is waiting for. So the count is capped by ``_MAX_ABANDONED_WORKERS`` and checked + BEFORE a worker is started: at the cap this refuses immediately rather than starting abandonment + N+1, and says the executor is saturated rather than blaming a statement that never ran. The slot + is released when the abandoned worker finally returns, which is the moment the leak actually + ends. + + Exceptions cross the thread boundary with their ORIGINAL type, re-raised here. That is what + keeps the handlers in ``execute_guarded`` correct: ``_ResourceLimit`` still reaches the refusal + branch and ``ExecutorError`` still reaches the classified one, rather than every executor + failure arriving as a worker that finished having produced nothing. ``BaseException`` is caught + rather than ``Exception`` for the same reason: a driver's deep ``sys.exit`` raises ``SystemExit``, + which ``tools._run_in_process`` nets one layer up, and swallowing it in a worker thread would + turn that fail-closed answer into a silent hang until the bound expired. + + The call runs inside a copy of the CALLER's context. A new thread starts with an empty one, so + without this the request-scoped ``_max_rows_override`` would read as unset inside the worker and + every in-process call would silently fall back to the deployment default. + """ + global _abandoned_workers + + timeout_s = _resolve_timeout_s() + outcome: dict[str, Any] = {} + ctx = copy_context() + + with _abandoned_lock: + if _abandoned_workers >= _MAX_ABANDONED_WORKERS: + # Fail closed, and before the work starts: the cheapest moment to say no, and the only + # one at which saying no still prevents anything. + raise _ExecutorSaturated(_EXECUTOR_SATURATED) + + def call() -> None: + global _abandoned_workers + try: + outcome["result"] = ctx.run(executor.execute, sql, creds, profile=profile) + except BaseException as exc: + outcome["error"] = exc + finally: + # `finished` and the slot release are set under the same lock the abandonment takes, so + # the SLOT ACCOUNTING cannot double-count: a worker that finishes in the sliver between + # the join timing out and the lock being taken decrements the slot the caller just took. + # + # The OUTCOME is not decided that cleanly, and the difference is a known defect — see + # issue #177. Whichever thread reaches this lock first decides: if the caller wins it + # reads `finished` unset and raises, so a call whose work had in fact completed is + # refused. It fails closed, and the built-in path is shielded by the per-statement + # watchdog firing a full skew earlier, so it is reachable only through an injected + # executor returning in that same instant. The repair is to signal the end of the WORK + # separately from this accounting, so the flag the caller reads cannot be delayed by + # lock contention. Do not "fix" it with `worker.is_alive()`: a worker parked here in + # its `finally` is still alive, and reads as abandoned in exactly the losing ordering. + with _abandoned_lock: + if outcome.get("abandoned"): + _abandoned_workers -= 1 + outcome["finished"] = True + + worker = threading.Thread(target=call, name="agami-bounded-execute", daemon=True) + worker.start() + worker.join(timeout_s + _OUTER_BOUND_SKEW_S) + with _abandoned_lock: + abandoned = not outcome.get("finished") + if abandoned: + _abandoned_workers += 1 + outcome["abandoned"] = True + if abandoned: + raise _OuterBoundExpired(_OUTLIVED_OUTER_BOUND) + if "error" in outcome: + raise outcome["error"] + return outcome["result"] + + def execute_guarded( sql: str, profile: str, @@ -1269,11 +1827,12 @@ def execute_guarded( **Every path returns exactly one ``Envelope``, and this function is TOTAL** — nothing is raised out of here for a caller to interpret, so the subprocess ``main`` and the in-process MCP handler - cannot disagree about what happened. The six outcomes: + cannot disagree about what happened. The seven outcomes: * the read-only gate refuses -> ``refused`` carrying that gate's ``Refusal`` * ``_model_safety`` returns a Refusal -> ``refused`` carrying it verbatim * ``_model_safety`` returns an int -> ``refused`` carrying the interim ``model_safety`` + * either time bound fired -> ``refused`` carrying ``resource_limit`` * ``executor.execute`` raises -> ``failed`` carrying a classified ``Failure`` * anything else raises -> ``failed``/``other``, generic message, raw to the log * the statement ran -> ``ok`` carrying the ``ExecResult`` @@ -1322,11 +1881,26 @@ def execute_guarded( "query.", )) creds = _load_credentials(profile, org_id or "local") - result = executor.execute(sql, creds, profile=profile) + # Bounded at the CHOKEPOINT, so the limit reaches every executor rather than only the + # built-in one whose engines carry the inner watchdog. See `_execute_bounded` for the + # mechanism and for the leaked worker it costs on expiry. + result = _execute_bounded(executor, sql, creds, profile=profile) # Inside the try on purpose: an executor that returns `None` (or anything else the contract # does not accept) fails the present-iff check in `Envelope.__post_init__`, and that is a # broken adapter, not a reason for the chokepoint to raise at its caller. return _envelope("ok", data=result) + except _ResourceLimit as exc: + # AHEAD of both handlers below: `_ResourceLimit` is an `ExecutorError` sibling, not a + # subclass, but the catch-all would swallow it and report a bound WE imposed as an + # unclassified server break. This is the per-statement bound the contract reserves + # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a + # fix we can honestly name, unlike the supervisor's kill of a child that never returned. + # + # ONE handler for all three bounds, and so exactly one refusal per call however many layers + # were armed. When the inner watchdog fired, its marker is what arrives here — the outer + # layer's join returned long before its own budget and has nothing to add, so the inner + # refusal is the answer and cannot be overwritten by a later one. + return _envelope("refused", refusal=_resource_limit_refusal(exc)) except ExecutorError as exc: # The classified branch: `msg` is authored by this module (a missing driver, a connect # failure, the credential-resolution remediation naming DATASOURCE_URL) or relayed from the diff --git a/packages/agami-core/src/guardrail.py b/packages/agami-core/src/guardrail.py index 0e7e4376..585d4ecb 100644 --- a/packages/agami-core/src/guardrail.py +++ b/packages/agami-core/src/guardrail.py @@ -50,11 +50,15 @@ RULE_COLUMN_SCOPE = "column_scope" RULE_SELECT_STAR = "select_star" RULE_MODEL_UNAVAILABLE = "model_unavailable" -# Declared and pinned below, with NO producer today. The contract reserves it for a **per-statement** -# timeout the guard imposes — a bound whose subject is the statement, so "narrow the query" is a fix -# we can honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this -# rule: it stops a child that never returned, without knowing what the child was doing when it -# stopped, so it is a `failed`/`timeout` (see `FailureKind` below and contract §3). +# Produced by the **per-statement** timeout the executor imposes: a watchdog cancels a statement that +# outlives the configured budget, and `execute_sql.execute_guarded` turns that into this refusal. Its +# subject is the statement, which is what earns it a rule at all — "narrow the query" is a fix we can +# honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this rule: it +# stops a child that never returned, without knowing what the child was doing when it stopped, so it +# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Every engine the executor speaks +# to is wired, with one recorded residual: BigQuery has no connection to cancel, so there the bound is +# server-side only and a client-side stall comes back as the executor never returning rather than as a +# statement we stopped. RULE_RESOURCE_LIMIT = "resource_limit" RULE_UNPARSEABLE = "unparseable" @@ -90,8 +94,8 @@ RULE_SELECT_STAR: "out_of_scope", RULE_MODEL_UNAVAILABLE: "undetermined", # A bound we imposed, not a property of the statement: neither unsafe nor out of scope — we - # simply did not determine the answer within the bound. Pinned here while unproduced so the gate - # that eventually imposes a per-statement timeout fills a constant rather than inventing one. + # simply did not determine the answer within the bound. Pinned before it had a producer, so the + # gate that imposes the per-statement timeout filled a constant rather than inventing one. RULE_RESOURCE_LIMIT: "undetermined", RULE_UNPARSEABLE: "undetermined", RULE_MODEL_SAFETY: "undetermined", @@ -165,9 +169,12 @@ def refuse(rule: str, *, detail: str, remediation: str) -> Refusal: in connect, credential resolution or model load, where "narrow the query" is the wrong fix. So an unresponsive executor is `failed` / `timeout`, and its message says only that we stopped waiting (guardrail contract §3). A **per-statement** timeout is the other case — its subject IS the -statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT` — and nothing imposes one yet, which is -why that rule is pinned with no producer. Driver-level connect/login timeouts fold into the connect -failure the executor already reports as `auth` (exit 4), because that is what the driver raises. +statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT`, and the executor now imposes one: a +watchdog cancels the statement through the driver and the outcome leaves the chokepoint on the +refusal channel rather than this one. The two therefore coexist and stay distinguishable in the audit +trail, which is the whole point of splitting them. Driver-level connect/login timeouts fold into the +connect failure the executor already reports as `auth` (exit 4), because that is what the driver +raises. `column_not_found`, `table_not_found`, `permission` and `network` are DECLARED BUT UNREACHABLE: producing them means parsing driver text, and sanitizing driver text belongs to the error-hardening diff --git a/packages/agami-core/src/mcp_http.py b/packages/agami-core/src/mcp_http.py index f032d14a..1704ecb1 100644 --- a/packages/agami-core/src/mcp_http.py +++ b/packages/agami-core/src/mcp_http.py @@ -428,9 +428,10 @@ async def _call_tool(name: str, arguments: dict) -> list: raised = False try: # Run the tool handler OFF the event loop. The heavy handlers block for the whole query — - # execute_sql shells out to a subprocess (up to 240 s) and hits the warehouse — so on the - # loop a single slow query would freeze every other in-flight request. This completes - # ACE-048 (which off-loaded the KDF/OIDC/audit calls but left the handler on the loop). + # execute_sql runs it under the executor's own bounds (the per-statement budget, plus the + # supervisor's slack on the fork path) and hits the warehouse — so on the loop a single + # slow query would freeze every other in-flight request. This completes ACE-048 (which + # off-loaded the KDF/OIDC/audit calls but left the handler on the loop). # `run_blocking` (anyio.to_thread) copies the request context into the worker thread, so # the org-scoped model cache (ACE-045, read via `_current_org_ctx`) stays tenant-correct. result_text = await run_blocking(meta["handler"], arguments or {}) diff --git a/packages/agami-core/src/tools.py b/packages/agami-core/src/tools.py index 36ba72f3..a7f97456 100644 --- a/packages/agami-core/src/tools.py +++ b/packages/agami-core/src/tools.py @@ -1063,7 +1063,7 @@ def _child_failure_message(returncode: int, stderr: str | None) -> str: def _executor_truncated(stderr: str | None) -> bool: - """True if execute_sql flagged a bounded-fetch truncation (ACE-038/044). The executor emits a + """True if execute_sql flagged a bounded-fetch truncation (ACE-044). The executor emits a non-error `{"truncated": {"row_cap": N}}` line on stderr alongside any other notices; scan for it.""" for line in (stderr or "").splitlines(): line = line.strip() @@ -1205,7 +1205,7 @@ def _emit( columns = list(env.data.columns) rows = [["" if v is None else str(v) for v in row] for row in env.data.rows] truncated = env.data.truncated - # Backstop only: the executor already caps at the source (ACE-038/044) and flags it. This + # Backstop only: the executor already caps at the source (ACE-044) and flags it. This # catches a result that slipped past that bound, and marks it truncated rather than # presenting a trimmed result as complete. if max_rows is not None and len(rows) > max_rows: @@ -1414,13 +1414,32 @@ def tool_execute_sql(args: dict[str, Any]) -> str: # materializes the whole result — not a client-side trim after the fact. cmd += ["--max-rows", str(max_rows)] + # The supervisor's bound is the OUTERMOST of the four time bounds, and it is DERIVED from the + # same resolver every inner layer reads rather than being a number of its own. A fixed 240s + # inverted that order for any statement budget approaching it: the supervisor fired FIRST, so a + # statement we could have cancelled and refused precisely came back instead as a + # `failed`/`timeout` naming nothing the caller can act on. Imported lazily for the same + # reason `_run_in_process` does it. + # + # Resolved HERE and enforced on a child that re-resolves for itself, which only works because the + # resolver reads the environment and nothing else: the child inherits `os.environ` (no `env=` + # below) and therefore reaches the identical number. A request-scoped override would be the one + # thing that could break that — it would outrank the environment on this side of the fork and be + # invisible on the other, so a parent bound of 65s could sit against a child budget of 300s and + # fire first, inverting the order this whole family exists to hold. There is deliberately no such + # override; `_resolve_timeout_s` documents why, and a test pins that the budget keeps exactly one + # configuration surface. + import execute_sql + + supervisor_timeout_s = execute_sql._resolve_timeout_s() + execute_sql._SUPERVISOR_SKEW_S + started = time.monotonic() try: proc = subprocess.run( cmd, capture_output=True, text=True, - timeout=240, + timeout=supervisor_timeout_s, ) except subprocess.TimeoutExpired: # A `failed`/`timeout`, NOT a `refused`/`resource_limit` — and the reason is what we can @@ -1428,8 +1447,10 @@ def tool_execute_sql(args: dict[str, Any]) -> str: # back; it does not say the STATEMENT ran long. The child may have hung in connect, in # credential resolution or in loading the model, and on any of those "narrow the query" is # advice pointing at the wrong thing. A refusal must name a fix, so a kill we cannot attribute - # is not one. (Guardrail contract §3: an unresponsive executor is `failed`. The per-statement - # timeout a later gate imposes IS a refusal, and `RULE_RESOURCE_LIMIT` is reserved for it.) + # is not one. (Guardrail contract §3: an unresponsive executor is `failed`. The executor's own + # per-statement deadline IS a refusal and carries `RULE_RESOURCE_LIMIT`; it can attribute the + # cancel to the statement because it is the thing it cancelled. The two bounds coexist, and + # this one must not borrow the other's rule.) return _emit( _envelope("failed", failure=Failure( kind="timeout", @@ -1466,7 +1487,7 @@ def tool_execute_sql(args: dict[str, Any]) -> str: env, sql=sql, execution_ms=execution_ms, profile=profile, args=args, ) - # Parse the RFC-4180 CSV emitted on stdout. The executor caps at the source (ACE-038/044) and + # Parse the RFC-4180 CSV emitted on stdout. The executor caps at the source (ACE-044) and # flags it on stderr; carry that flag so a truncated result is never presented as complete. The # `max_rows` backstop is applied by `_emit`, the same one the in-process path gets. reader = csv.reader(io.StringIO(proc.stdout)) diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 9a5b51d2..3374f44e 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -52,16 +52,18 @@ import argparse import configparser +import contextlib import csv import json import logging import os import stat import sys +import threading import urllib.parse import uuid -from collections.abc import Callable -from contextvars import ContextVar +from collections.abc import Callable, Iterator +from contextvars import ContextVar, copy_context from dataclasses import asdict, dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -70,6 +72,7 @@ from guardrail import ( RULE_MODEL_SAFETY, RULE_MODEL_UNAVAILABLE, + RULE_RESOURCE_LIMIT, Envelope, Failure, FailureKind, @@ -483,21 +486,70 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Postgres connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # Bound outside the try so the `finally` can close it whether or not the declare got that far. + cur = None try: + # `with conn` is the TRANSACTION, not the connection: leaving it by an exception rolls back, + # which is why it stays even though the cursor no longer lives inside a `with` of its own. with conn: + # `SET LOCAL statement_timeout` is the native server-side backstop. It is + # TRANSACTION-scoped, so it has to run on THIS transaction and before the named cursor + # declares — a regular client-side cursor, because a named cursor may only ever run the + # one query it was declared for. The libpq `options` startup parameter would be the + # other way to set it and is deliberately not used: a transaction-mode connection pooler + # can reject an unknown startup parameter, which would break the connect outright rather + # than bound the statement. + with conn.cursor() as bound_cur: + bound_cur.execute( + "SET LOCAL statement_timeout = %s", + ((timeout_s + _NATIVE_BOUND_SKEW_S) * 1000,), # the setting is in milliseconds + ) # A server-side (named) cursor so the row cap bounds TRANSFER, not just what we write: # psycopg2's default client-side cursor buffers the ENTIRE result before we can fetchmany, # so a runaway result would still be pulled whole. The named cursor streams from the - # server in bounded batches (ACE-038). Read-only SELECTs (the only thing the guard admits) + # server in bounded batches. Read-only SELECTs (the only thing the guard admits) # are exactly what a server-side cursor supports. - with conn.cursor(name="agami_bounded") as cur: - cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window - cur.execute(sql) - result = _collect_cursor(cur) + cur = conn.cursor(name="agami_bounded") + cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window + # `connection.cancel()` is psycopg2's cancel: it opens a second connection and sends the + # libpq cancel request, so it is safe to call from the watchdog thread while this one is + # blocked in the driver. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Postgres execution error: {e}", code=5) finally: - conn.close() + # The named cursor is closed HERE, by hand, rather than by a `with` around it. Closing one + # sends `CLOSE agami_bounded` to the server, and when we are unwinding on the timeout marker + # the transaction is already aborted, so that statement raises in turn — from inside + # `__exit__`, where a raised exception REPLACES the one being propagated. The refusal would + # reach the chokepoint as an ordinary driver error and be reported as a failure. Swallowing + # it costs nothing: by this point the transaction has ended and the server-side portal is + # gone with it, so the close is a courtesy rather than the thing that frees the resource. + if cur is not None: + try: + cur.close() + except Exception: + pass + # And the connection close is guarded for the same reason, one line down: an exception raised + # inside a `finally` REPLACES the one propagating through it, so an unguarded close here would + # destroy the marker the `except _ResourceLimit: raise` above just re-raised — and the caller + # would read a bound we imposed as an unclassified server break. + try: + conn.close() + except Exception: + pass return result @@ -520,14 +572,38 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"MySQL connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: with conn.cursor() as cur: - cur.execute(sql) - result = _collect_cursor(cur) + # pymysql's connection has NO `cancel()`. Its two cancel-shaped methods are `close()`, + # which sends COM_QUIT down the very socket the blocked statement owns — and so can + # itself block on a connection that is already stuck — and `_force_close()`, which + # closes the socket outright. Only the second one actually unblocks a statement in + # flight, so it is named here explicitly rather than reached for by shape. + with _deadline(conn._force_close, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: - conn.close() + # Guarded because an exception raised inside a `finally` REPLACES the one propagating through + # it, which would destroy the marker re-raised just above. MySQL is the most exposed of the + # ten: its cancel is `_force_close()`, which deliberately destroys the socket that this + # `close()` then tries to write COM_QUIT to — so on exactly the timeout path this is the close + # most likely to raise. + try: + conn.close() + except Exception: + pass return result @@ -548,11 +624,20 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: "Add one to /local/credentials.", code=2, ) + # Resolved before the connect, because the native backstop is a SESSION parameter and so has to + # be handed to the connect call itself. + timeout_s = _resolve_timeout_s() conn_kwargs: dict[str, Any] = { "account": creds["account"], "user": creds["user"], "client_session_keep_alive": False, "login_timeout": 15, + # Snowflake's native server-side bound, set behind our watchdog by the usual skew so the + # watchdog wins and the flag stays the only classification. On a warehouse the backstop is + # worth more than elsewhere: a statement nobody is waiting for still bills credits. + "session_parameters": { + "STATEMENT_TIMEOUT_IN_SECONDS": timeout_s + _NATIVE_BOUND_SKEW_S, + }, } for k in ("password", "warehouse", "database", "schema", "role", "authenticator"): if creds.get(k): @@ -563,8 +648,32 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"Snowflake connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + + def cancel_snowflake_statement() -> None: + """Ask Snowflake to abort the statement this cursor is running. + + Neither the connection nor the cursor has a `cancel()` — verified against the connector's + own source, where the only public abort is `SnowflakeCursor.abort_query(qid)` and the + connection's cancel helpers are private. The query id is what identifies the statement to + abort, and it exists only once the statement has been submitted; before that there is + nothing to abort and the session parameter above is what bounds the call. + """ + qid = cur.sfqid + if qid: + cur.abort_query(qid) + + with _deadline(cancel_snowflake_statement, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Snowflake execution error: {e}", code=5) finally: @@ -634,9 +743,20 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"BigQuery client init failed: {e}", code=4) + timeout_s = _resolve_timeout_s() + # `job_timeout_ms` is BigQuery's native server-side bound, and on this engine it is the ONLY + # bound: there is no watchdog cancel here, deliberately. BigQuery hands out no connection or + # cursor to cancel, and the call that blocks is `job.result()`, which is reached only AFTER + # `client.query()` has returned — so at the instant a watchdog would fire there is nothing in + # hand to stop. Stated plainly, and accepted rather than papered over: on BigQuery a client-side + # stall comes back as a `failed` envelope, not a `resource_limit` refusal. The query itself is + # still stopped by the bound below, which is what keeps a runaway from scanning on unattended. + # The usual skew is kept so this engine's number matches every other engine's. + job_config_kwargs: dict[str, Any] = { + "job_timeout_ms": (timeout_s + _NATIVE_BOUND_SKEW_S) * 1000, + } # If `dataset` was set, prefix unqualified table references via the # default_dataset job config so the SQL can omit `..` - job_config_kwargs: dict[str, Any] = {} if creds.get("dataset"): try: job_config_kwargs["default_dataset"] = f"{project}.{creds['dataset']}" @@ -645,15 +765,17 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: cap = _resolve_row_cap() try: - if job_config_kwargs: - job_config = bigquery.QueryJobConfig(**job_config_kwargs) - job = client.query(sql, job_config=job_config) - else: - job = client.query(sql) + job_config = bigquery.QueryJobConfig(**job_config_kwargs) + job = client.query(sql, job_config=job_config) # BigQuery has no DB-API cursor, so it can't funnel through `_collect_cursor`; apply the # same bounded-fetch cap here. `max_results=cap+1` bounds what the API returns (transfer), # and the (cap+1)th row flags truncation — the never-silent guarantee holds for BigQuery too. results = job.result(max_results=cap + 1) # waits for completion; raises on error + except _ResourceLimit: + # Nothing in this engine raises the marker today, and this clause is still not optional: the + # rule is that NO engine may relabel it as a driver error, and the clause below would. It is + # what makes the rule checkable across all ten engines rather than nine. + raise except Exception as e: raise ExecutorError(f"BigQuery execution error: {e}", code=5) @@ -672,6 +794,12 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: + """Tier-3 path for SQLite using the stdlib `sqlite3` module. + + The first engine to run under the per-statement deadline. `sqlite3.Connection.interrupt()` is a + genuine cancel — it stops the statement mid-scan from another thread — so this engine can prove + the whole refusal contract in-process, with no network and no fixture warehouse. + """ import sqlite3 # always available in stdlib _require(creds, "path") path = os.path.expanduser(creds["path"]) @@ -679,14 +807,48 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: conn = sqlite3.connect(path) except Exception as e: raise ExecutorError(f"SQLite connect failed: {e}", code=4) + # Resolved ONCE for the call, so the budget the watchdog enforces and the number the refusal + # quotes cannot be two different values. + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The deadline covers the FETCH as well as the execute. `_collect_cursor` pulls `cap + 1` + # rows in a single `fetchmany`, and on a cursor that streams its result that pull is where + # the scan actually happens — a clock that stopped at `execute` would bound the cheap half + # of the work and leave the expensive half unbounded. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + # A cancelled statement raises `sqlite3.OperationalError("interrupted")`, whose text + # is not a classification we would want to key on. The FLAG is the classification, + # and it is the ONLY one: `_deadline` sets it before the cancel lands, so an error + # raised while it is unset is the database's own — however late it arrives. + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + # Checked after the watchdog is disarmed, so the flag is final. A cancel can also land + # between the two calls above, or just as the second returns, and leave nothing to raise: + # the budget still elapsed, so the outcome is still a refusal rather than a result gathered + # past it. + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + # Ahead of the catch-all below on purpose: our own marker must reach `execute_guarded` + # intact. Wrapped in an `ExecutorError` it would become a `failed`/`syntax` envelope, which + # is the database's outcome rather than the bound we imposed. + raise except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: - conn.close() + # Guarded like every other engine's: an exception raised inside a `finally` REPLACES the one + # propagating through it, so an unguarded close would silently convert the refusal above into + # an unclassified failure. + try: + conn.close() + except Exception: + pass return result @@ -705,10 +867,26 @@ def _run_sqlserver(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"SQL Server connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # pymssql's DB-API `Connection` has no `cancel()` — verified against pymssql 2.3, whose + # connection exposes only close/commit/cursor/rollback/bulk_copy, and whose cursor exposes + # none either. The cancel lives one layer down, on the `_mssql.MSSQLConnection` that + # connection wraps, reachable only as the private `_conn`; that object's `cancel()` sends the + # TDS attention packet, which is the thing that actually stops a statement in flight. + with _deadline(conn._conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"SQL Server execution error: {e}", code=5) finally: @@ -735,10 +913,23 @@ def _run_oracle(creds: dict[str, str], sql: str) -> ExecResult: conn = oracledb.connect(user=creds["user"], password=creds["password"], dsn=dsn) except Exception as e: raise ExecutorError(f"Oracle connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # `oracledb.Connection.cancel()` breaks out of the call in progress on that connection — + # verified against python-oracledb 3.x, where it is on the CONNECTION and not on the cursor. + with _deadline(conn.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Oracle execution error: {e}", code=5) finally: @@ -766,10 +957,24 @@ def _run_databricks(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Databricks connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # The cancel is on the CURSOR here, not the connection — verified against + # databricks-sql-connector 4.x, whose `Cursor.cancel()` posts a cancel for the operation that + # cursor is running while its connection has none. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Databricks execution error: {e}", code=5) finally: @@ -798,10 +1003,23 @@ def _run_trino(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Trino connect failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # Trino's cancel is on the CURSOR — verified against the trino client, whose + # `Cursor.cancel()` sends the coordinator a DELETE for the query that cursor started. + with _deadline(cur.cancel, timeout_s) as expired: + try: + cur.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Trino execution error: {e}", code=5) finally: @@ -823,9 +1041,25 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: conn = duckdb.connect(path, read_only=True) except Exception as e: raise ExecutorError(f"DuckDB open failed: {e}", code=4) + timeout_s = _resolve_timeout_s() try: - cur = conn.execute(sql) - result = _collect_cursor(cur) + # `conn.interrupt()` is DuckDB's cancel, and it is on the connection — which is just as well, + # because on this engine there is no cursor to reach for until the execute has returned: + # `conn.execute` HANDS BACK THE CONNECTION ITSELF, so `cur` below IS `conn`. The deadline + # therefore has to be armed around the execute as well as the fetch, and on an in-process + # engine the execute is where the scan happens. + with _deadline(conn.interrupt, timeout_s) as expired: + try: + cur = conn.execute(sql) + result = _collect_cursor(cur) + except Exception: + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + raise + if expired.is_set(): + raise _ResourceLimit(_OUTLIVED_BUDGET) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"DuckDB execution error: {e}", code=5) finally: @@ -836,7 +1070,7 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: return result -_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation (ACE-038) +_DEFAULT_MAX_ROWS = 1000 # rows materialized per result before truncation # Per-call cap from --max-rows (ACE-044). A ContextVar, not a plain global, so it is REQUEST-SCOPED # once the HTTP server runs execution in-process (ACE-028): concurrent handlers run in worker threads # (`run_blocking` copies the context per call, like `_current_org_ctx`), so each request's cap is @@ -860,10 +1094,225 @@ def _resolve_row_cap() -> int: return cap +_DEFAULT_TIMEOUT_S = 30 # wall-clock seconds one statement may run before the watchdog cancels it +# How far BEHIND our own watchdog a NATIVE server-side bound is set, on the three engines that have +# one. The skew is the whole point: our watchdog fires at the budget, the engine's own bound five +# seconds later, so the watchdog always wins the race and its Event stays the SOLE classification +# signal. Reverse the order and a server-side kill would beat us to the statement, the flag would be +# clear, and the refusal would come back as an ordinary database failure. The native bound is a +# backstop for the one case the watchdog cannot cover — our process dying with a statement in +# flight, which would otherwise leave the engine scanning for nobody. +_NATIVE_BOUND_SKEW_S = 5 +# How far behind the watchdog the OUTER bound sits — the one `execute_guarded` puts around +# `executor.execute` itself, so an INJECTED executor is bounded too. The four bounds are one ordered +# family resolved from the same budget, innermost first: +# +# watchdog (timeout_s) < native (+5) < outer (+10) < supervisor (+60) +# +# and the order is the contract. The inner watchdog must always win, because it is the only layer +# that can attribute the stop to the statement and hand the caller the precise refusal; every layer +# behind it is a backstop for a failure the layer in front of it cannot see. Collapse the order and +# a bound we imposed comes back as something else — a native server-side kill reads as an ordinary +# database error, and a supervisor kill reads as `failed`/`timeout` naming nothing the caller can +# act on. Both are strictly worse answers to the same event. +_OUTER_BOUND_SKEW_S = 10 +# The OUTERMOST bound: how far behind the watchdog the subprocess supervisor in `tools` stops waiting +# for a forked child. It lives here, next to its three siblings, so the whole family is derived from +# one resolver and one place — the supervisor was a hardcoded 240s, which for any statement budget +# approaching it became the FIRST bound to fire and inverted the order. Its slack is large because it +# bounds a whole process (interpreter start, model load, credential resolution, connect) rather than +# a statement. +_SUPERVISOR_SKEW_S = 60 +# How many worker threads `_execute_bounded` may have given up on at once, process-wide. That layer +# holds nothing it can cancel, so every expiry abandons a thread which still occupies a pooled +# connection and, server-side, a running statement. On the built-in path the inner watchdog fires ten +# seconds earlier, so reaching the outer bound means a cancel already failed — rare, and an honest +# cost. On an INJECTED executor there is no inner watchdog at all, so every slow statement abandons +# one, and because the caller's own thread is freed at the bound the abandonments accumulate with no +# ceiling of their own: the anyio worker limiter used to supply that ceiling by blocking the caller, +# and bounding the wait removed it. The cap restores one explicitly. 8 is deliberately far below both +# that former limiter and a typical warehouse connection pool, so a saturated executor loses a +# minority of the pool to work nobody is waiting for rather than taking the datasource away from the +# whole organization. +_MAX_ABANDONED_WORKERS = 8 +_abandoned_lock = threading.Lock() +_abandoned_workers = 0 + + +def _resolve_timeout_s() -> int: + """Effective per-statement timeout, in whole seconds. `AGAMI_SQL_TIMEOUT_S` is the + operator-configurable DEPLOYMENT budget (default 30 when unset) — an operator owns their + availability tradeoff and may set it higher OR lower than 30. A missing or non-positive value + falls back to the default. + + **The environment is the ONLY source, deliberately.** A request-scoped override would outrank it + in the parent and be invisible to a forked child, which re-resolves from `os.environ` alone — so + the supervisor bound the parent derives could sit BELOW the budget the child actually enforces + and fire first, inverting the ordered family the whole design rests on. One source, readable on + both sides of the fork, makes that inversion unrepresentable rather than merely unlikely. + + Unlike `_resolve_row_cap`, a value that is PRESENT and does not survive to become the budget is + logged at warning before the fallback. That covers `45.5` and `30s`, which cannot be read at all, + and equally `-5` and `0`, which can be read and are then declined: an operator who wrote either + asked for something specific, and a deployment quietly running 30 instead is exactly the + invisible degradation the warning exists against. The warning goes to the module logger and never + to stderr, because the subprocess transport parses stderr and an extra line there would break + that contract.""" + raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() + digits = raw[1:] if raw.startswith("-") else raw # a leading minus is a value, not a typo + # `isdecimal`, not `isdigit`: the latter admits `²` and `①`, which `int()` then refuses — turning + # a misconfigured deployment into a ValueError raised out of this resolver, at a call site (the + # fork path's supervisor bound) that sits outside any handler. + written = int(raw) if digits.isdecimal() else None + timeout_s = written if written is not None and written > 0 else _DEFAULT_TIMEOUT_S + if raw and timeout_s != written: + _LOG.warning( + "AGAMI_SQL_TIMEOUT_S=%r is not a usable whole number of seconds; falling back to %ds.", + raw, + _DEFAULT_TIMEOUT_S, + ) + return timeout_s + + +class _ResourceLimit(Exception): + """Raised when our own watchdog fired, so a cancelled statement unwinds the engine function the + same way any other failure does and the surrounding transaction rolls back. It is an internal + marker only: it is always caught and translated inside this module and never crosses the tool + boundary.""" + + +class _OuterBoundExpired(_ResourceLimit): + """Raised when the OUTER bound expired: the executor never returned and we stopped waiting. + + A subclass, so the one `except _ResourceLimit` handler in `execute_guarded` produces exactly one + refusal for either layer and no second one can be minted. It stays distinguishable because the + two events are not the same event: the watchdog CANCELLED a statement it holds a connection to, + while this layer only abandoned its wait — so the sentence the caller reads differs, and saying + "cancelled" here would be false. + """ + + +class _ExecutorSaturated(_ResourceLimit): + """Raised INSTEAD of starting a statement, when the abandoned-worker cap is already reached. + + A sibling of `_OuterBoundExpired` under the same parent for the same reason — one handler, one + refusal — and distinguishable for the same reason too: this statement did not run long, it did + not run at all. What we observed is the executor, not the statement, so the sentence the caller + reads must not blame the query it just sent. + """ + + +# The marker's message, single-sourced because every engine raises it. It is diagnostic text, not +# caller-facing: the refusal `execute_guarded` builds re-resolves the budget and writes its own +# detail, so nothing a caller reads comes from here. +_OUTLIVED_BUDGET = "the statement outlived its per-statement budget" +# The outer marker's message, and diagnostic in the same way — it names the executor rather than the +# statement, because at this layer the statement is not the thing we observed. +_OUTLIVED_OUTER_BOUND = "the executor outlived the outer bound around it" +# The saturation marker's message, diagnostic in the same way again. +_EXECUTOR_SATURATED = "the executor already has its limit of abandoned calls outstanding" + +# The remediation names only what would make THIS statement executable: on the served path the caller +# is an assistant with no shell and no deployment, so naming the environment variable would be advice +# it cannot take, addressed to someone who is not reading. +_NARROW_IT = ("Narrow the time range, reduce the grouping, or add a selective filter, " + "then run it again.") + + +def _resource_limit_refusal(exc: _ResourceLimit) -> Refusal: + """The ONE `resource_limit` refusal, built from whichever of the three bounds raised. + + Single-sourced because two entries into the engines lead here — the guarded chokepoint and the + subprocess/CLI adapter — and a caller must not be able to tell which one it came through. + + The budget is re-resolved rather than carried on the marker: the marker stays a plain exception, + and nothing between the engine call and here can change the environment the resolver reads, so it + is the same number the watchdog used. The configured number belongs in the detail — it is a + deployment setting, not a data value, and a bound the caller cannot see is one it cannot plan + around. + """ + timeout_s = _resolve_timeout_s() + if isinstance(exc, _ExecutorSaturated): + # Nothing ran, so nothing about THIS statement is the finding. Saying "your query was too + # slow" here would be false and would send the caller off simplifying a statement that is + # very possibly fine. + detail = ( + "The executor is saturated: too many earlier calls have not returned, so this " + "statement was not started." + ) + remediation = "Wait for the calls already in flight to finish, then run it again." + elif isinstance(exc, _OuterBoundExpired): + # The outer layer holds no connection, so it stopped WAITING rather than stopping the + # statement. "Cancelled" would be a claim we cannot make: the worker is still inside the + # driver call and the statement may still be running. The caller is told what we actually + # know, and the number quoted is the bound that actually elapsed. + detail = ( + f"The executor did not return within the {timeout_s + _OUTER_BOUND_SKEW_S}s limit " + "and the query was abandoned." + ) + remediation = _NARROW_IT + else: + detail = f"The statement ran longer than the {timeout_s}s limit and was cancelled." + remediation = _NARROW_IT + return refuse(RULE_RESOURCE_LIMIT, detail=detail, remediation=remediation) + + +@contextlib.contextmanager +def _deadline(cancel: Callable[[], None], timeout_s: float) -> Iterator[threading.Event]: + """Arm a watchdog that calls `cancel` if the wrapped block outlives `timeout_s`, and yield the + `threading.Event` that says whether it fired. + + The event is set BEFORE `cancel` runs. Order matters: whoever catches the driver error that the + cancellation provokes must be able to read an already-set flag and attribute the failure to us + rather than to the database. A `cancel` that raises is swallowed and logged, because some drivers + raise when cancelled from a thread other than the one running the statement, and an exception + escaping a timer thread is both unhandleable by the caller and invisible in the result. + + **The disarm is a JOIN, not a request.** `threading.Timer.cancel()` only sets the timer's internal + `finished` Event, so a `fire` that already passed its own check runs to completion regardless — + setting the flag and delivering a cancel AFTER this block has returned. Both consequences are + real: every engine re-reads the flag once the block has exited and would read a stale `False`, and + a cancel arriving late lands on a connection the engine has moved on from, which on a pooled one + is by then someone else's statement. The lock plus the flag close that: a `fire` that loses the + race neither sets the Event nor calls `cancel`, and a `fire` that wins holds the lock across the + cancel, so the disarm waits for it. The wait is bounded by the layer outside this one — a cancel + that hangs is what the outer bound around the whole executor call exists to survive.""" + fired = threading.Event() + lock = threading.Lock() + # Assigned only below, in the enclosing scope, and only read inside `fire` — so the closure sees + # the current value with no `nonlocal` and no mutable box. + disarmed = False + + def fire() -> None: + with lock: + if disarmed: + return # the block already returned; there is nothing of ours left to stop + fired.set() + try: + cancel() + except Exception as exc: + _LOG.warning("Cancelling the statement after its timeout expired failed: %s", exc) + + timer = threading.Timer(timeout_s, fire) + timer.daemon = True # a hung cancel must never hold the interpreter open at shutdown + timer.start() + try: + yield fired + finally: + # Claim the race under the lock FIRST, then ask the timer to stand down. `Timer.cancel()` is + # not a join, so doing it the other way round leaves a window between this block returning + # and the flag being set, in which a `fire` that has already expired can still take the lock + # and mark a statement that actually completed. Taking the lock first closes that window: a + # `fire` already holding the lock is waited for, and every other `fire` observes `disarmed`. + with lock: + disarmed = True + timer.cancel() # belt and braces, for a timer that has not started running yet + + def _flag_truncated(cap: int) -> None: """Signal a bounded-fetch truncation to the caller — a non-error `{"truncated": …}` marker on stderr (distinct from the guards' `{"error": …}`), so a truncated result is never mistaken for a - complete one (ACE-038/044). Shared by every engine's materialization path. One write so the + complete one (ACE-044). Shared by every engine's materialization path. One write so the marker is always a single line, even if other notices surround it.""" sys.stderr.write(json.dumps({"truncated": {"row_cap": cap}}) + "\n") @@ -876,7 +1325,7 @@ def _collect_cursor(cur: Any) -> ExecResult: executor path share, so the row cap is enforced once, identically, for every caller. Fetch FIRST, then read ``cur.description``: a psycopg2 **server-side (named) cursor** — which the - Postgres/Redshift path uses to bound transfer (ACE-038) — reports ``description is None`` until the + Postgres/Redshift path uses to bound transfer — reports ``description is None`` until the first fetch, so reading it beforehand would drop EVERY row of a real Postgres result. Client-side cursors (sqlite/mysql/…) set ``description`` at execute, so fetch-first is equally correct there.""" cap = _resolve_row_cap() @@ -969,7 +1418,8 @@ def _resolve_guard_model(profile: str): def _write_refusal(refusal: Refusal) -> None: """Write a guard refusal to stderr in the ONE shape every caller parses — the single wire-writer, - called only from ``main``. + called from ``main`` and from the per-engine CSV adapter, which is the other entry into the + engines. One JSON object, on one line: ``{"refusal": {reason, rule, detail, remediation}}`` — the wire shape S2 established, which ``tools._stderr_refusal`` rebuilds through ``Refusal`` on the parent @@ -1115,9 +1565,18 @@ def _model_safety(sql: str, profile: str, area: str | None) -> tuple[str, Refusa def _emit_or_err(run: Callable[[], ExecResult]) -> int: """Subprocess/CLI adapter over a ``_run_`` function: write its result to stdout as CSV and return exit code 0, or translate an ``ExecutorError`` into the stderr message + exit code the CLI - contract documents (byte-identical to what the old ``_execute_`` emitted).""" + contract documents (byte-identical to what the old ``_execute_`` emitted). + + The watchdog marker gets an arm of its own, AHEAD of the classified one, because this is the + second entry into the engine functions and they raise it. Without it a per-statement timeout + reached through here would escape as a traceback rather than as the exit code this docstring + promises. It renders exactly what ``main`` renders for a refusal — the one JSON object on stderr + and exit 1 — so the two entries cannot disagree about what a bound we imposed looks like.""" try: _emit_result_csv(run()) + except _ResourceLimit as exc: + _write_refusal(_resource_limit_refusal(exc)) + return 1 except ExecutorError as e: return _err(e.msg, code=e.code) return 0 @@ -1250,6 +1709,105 @@ def _envelope( ) +def _execute_bounded( + executor: Executor, sql: str, creds: dict[str, str], *, profile: str +) -> ExecResult: + """Run ``executor.execute`` under the OUTER bound and return what it returned. + + This is the layer that makes "no executor escapes the limit" true. The inner deadline lives + inside the BUILT-IN executor's engine functions, so before this an INJECTED executor — the + hosted connection-reuse path — ran with no per-statement bound at all, and even on the built-in + path BigQuery has no watchdog (it has no connection object to cancel), so a client-side stall + there was unbounded too. Both are bounded here, and deliberately by the same layer: applying + this only to non-built-in executors would leave the BigQuery hole open while looking closed. + + **The mechanism is a worker thread, because this layer holds nothing it can cancel.** An + arbitrary ``Executor`` exposes one method and no connection, so the only thing this code owns is + its own WAIT — it starts the call on a daemon thread and joins with the budget. + + **The cost is an abandoned worker, and it is real.** On expiry the thread is still inside the + driver call, and it stays there: nothing here cancels it, and it may hold a database connection + (and, on the server side, a running statement) until that call returns on its own. It is a + daemon so it cannot hold the interpreter open at exit. This is a bound on how long a CALLER + waits, not a promise that the work stopped — which is exactly why the inner watchdog, the layer + that really can cancel, is set to fire first. + + **And the cost is CAPPED, because bounding the wait removed the ceiling that used to bound it.** + Before this layer existed the same slow call blocked the caller's own worker thread, so the + host's thread limiter capped how much abandoned work could be in flight and applied backpressure + to everything behind it. Returning at the bound frees that slot, so on the injected path — where + there is no inner watchdog and EVERY slow statement abandons one — the abandonments would + otherwise accumulate without limit until the pool, and with it the datasource, belonged entirely + to work nobody is waiting for. So the count is capped by ``_MAX_ABANDONED_WORKERS`` and checked + BEFORE a worker is started: at the cap this refuses immediately rather than starting abandonment + N+1, and says the executor is saturated rather than blaming a statement that never ran. The slot + is released when the abandoned worker finally returns, which is the moment the leak actually + ends. + + Exceptions cross the thread boundary with their ORIGINAL type, re-raised here. That is what + keeps the handlers in ``execute_guarded`` correct: ``_ResourceLimit`` still reaches the refusal + branch and ``ExecutorError`` still reaches the classified one, rather than every executor + failure arriving as a worker that finished having produced nothing. ``BaseException`` is caught + rather than ``Exception`` for the same reason: a driver's deep ``sys.exit`` raises ``SystemExit``, + which ``tools._run_in_process`` nets one layer up, and swallowing it in a worker thread would + turn that fail-closed answer into a silent hang until the bound expired. + + The call runs inside a copy of the CALLER's context. A new thread starts with an empty one, so + without this the request-scoped ``_max_rows_override`` would read as unset inside the worker and + every in-process call would silently fall back to the deployment default. + """ + global _abandoned_workers + + timeout_s = _resolve_timeout_s() + outcome: dict[str, Any] = {} + ctx = copy_context() + + with _abandoned_lock: + if _abandoned_workers >= _MAX_ABANDONED_WORKERS: + # Fail closed, and before the work starts: the cheapest moment to say no, and the only + # one at which saying no still prevents anything. + raise _ExecutorSaturated(_EXECUTOR_SATURATED) + + def call() -> None: + global _abandoned_workers + try: + outcome["result"] = ctx.run(executor.execute, sql, creds, profile=profile) + except BaseException as exc: + outcome["error"] = exc + finally: + # `finished` and the slot release are set under the same lock the abandonment takes, so + # the SLOT ACCOUNTING cannot double-count: a worker that finishes in the sliver between + # the join timing out and the lock being taken decrements the slot the caller just took. + # + # The OUTCOME is not decided that cleanly, and the difference is a known defect — see + # issue #177. Whichever thread reaches this lock first decides: if the caller wins it + # reads `finished` unset and raises, so a call whose work had in fact completed is + # refused. It fails closed, and the built-in path is shielded by the per-statement + # watchdog firing a full skew earlier, so it is reachable only through an injected + # executor returning in that same instant. The repair is to signal the end of the WORK + # separately from this accounting, so the flag the caller reads cannot be delayed by + # lock contention. Do not "fix" it with `worker.is_alive()`: a worker parked here in + # its `finally` is still alive, and reads as abandoned in exactly the losing ordering. + with _abandoned_lock: + if outcome.get("abandoned"): + _abandoned_workers -= 1 + outcome["finished"] = True + + worker = threading.Thread(target=call, name="agami-bounded-execute", daemon=True) + worker.start() + worker.join(timeout_s + _OUTER_BOUND_SKEW_S) + with _abandoned_lock: + abandoned = not outcome.get("finished") + if abandoned: + _abandoned_workers += 1 + outcome["abandoned"] = True + if abandoned: + raise _OuterBoundExpired(_OUTLIVED_OUTER_BOUND) + if "error" in outcome: + raise outcome["error"] + return outcome["result"] + + def execute_guarded( sql: str, profile: str, @@ -1269,11 +1827,12 @@ def execute_guarded( **Every path returns exactly one ``Envelope``, and this function is TOTAL** — nothing is raised out of here for a caller to interpret, so the subprocess ``main`` and the in-process MCP handler - cannot disagree about what happened. The six outcomes: + cannot disagree about what happened. The seven outcomes: * the read-only gate refuses -> ``refused`` carrying that gate's ``Refusal`` * ``_model_safety`` returns a Refusal -> ``refused`` carrying it verbatim * ``_model_safety`` returns an int -> ``refused`` carrying the interim ``model_safety`` + * either time bound fired -> ``refused`` carrying ``resource_limit`` * ``executor.execute`` raises -> ``failed`` carrying a classified ``Failure`` * anything else raises -> ``failed``/``other``, generic message, raw to the log * the statement ran -> ``ok`` carrying the ``ExecResult`` @@ -1322,11 +1881,26 @@ def execute_guarded( "query.", )) creds = _load_credentials(profile, org_id or "local") - result = executor.execute(sql, creds, profile=profile) + # Bounded at the CHOKEPOINT, so the limit reaches every executor rather than only the + # built-in one whose engines carry the inner watchdog. See `_execute_bounded` for the + # mechanism and for the leaked worker it costs on expiry. + result = _execute_bounded(executor, sql, creds, profile=profile) # Inside the try on purpose: an executor that returns `None` (or anything else the contract # does not accept) fails the present-iff check in `Envelope.__post_init__`, and that is a # broken adapter, not a reason for the chokepoint to raise at its caller. return _envelope("ok", data=result) + except _ResourceLimit as exc: + # AHEAD of both handlers below: `_ResourceLimit` is an `ExecutorError` sibling, not a + # subclass, but the catch-all would swallow it and report a bound WE imposed as an + # unclassified server break. This is the per-statement bound the contract reserves + # `resource_limit` for — its subject IS the statement, so "narrow it and run it again" is a + # fix we can honestly name, unlike the supervisor's kill of a child that never returned. + # + # ONE handler for all three bounds, and so exactly one refusal per call however many layers + # were armed. When the inner watchdog fired, its marker is what arrives here — the outer + # layer's join returned long before its own budget and has nothing to add, so the inner + # refusal is the answer and cannot be overwritten by a later one. + return _envelope("refused", refusal=_resource_limit_refusal(exc)) except ExecutorError as exc: # The classified branch: `msg` is authored by this module (a missing driver, a connect # failure, the credential-resolution remediation naming DATASOURCE_URL) or relayed from the diff --git a/plugins/agami/lib/guardrail.py b/plugins/agami/lib/guardrail.py index 0e7e4376..585d4ecb 100644 --- a/plugins/agami/lib/guardrail.py +++ b/plugins/agami/lib/guardrail.py @@ -50,11 +50,15 @@ RULE_COLUMN_SCOPE = "column_scope" RULE_SELECT_STAR = "select_star" RULE_MODEL_UNAVAILABLE = "model_unavailable" -# Declared and pinned below, with NO producer today. The contract reserves it for a **per-statement** -# timeout the guard imposes — a bound whose subject is the statement, so "narrow the query" is a fix -# we can honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this -# rule: it stops a child that never returned, without knowing what the child was doing when it -# stopped, so it is a `failed`/`timeout` (see `FailureKind` below and contract §3). +# Produced by the **per-statement** timeout the executor imposes: a watchdog cancels a statement that +# outlives the configured budget, and `execute_sql.execute_guarded` turns that into this refusal. Its +# subject is the statement, which is what earns it a rule at all — "narrow the query" is a fix we can +# honestly name. The subprocess supervisor's kill is NOT that bound and must not borrow this rule: it +# stops a child that never returned, without knowing what the child was doing when it stopped, so it +# is a `failed`/`timeout` (see `FailureKind` below and contract §3). Every engine the executor speaks +# to is wired, with one recorded residual: BigQuery has no connection to cancel, so there the bound is +# server-side only and a client-side stall comes back as the executor never returning rather than as a +# statement we stopped. RULE_RESOURCE_LIMIT = "resource_limit" RULE_UNPARSEABLE = "unparseable" @@ -90,8 +94,8 @@ RULE_SELECT_STAR: "out_of_scope", RULE_MODEL_UNAVAILABLE: "undetermined", # A bound we imposed, not a property of the statement: neither unsafe nor out of scope — we - # simply did not determine the answer within the bound. Pinned here while unproduced so the gate - # that eventually imposes a per-statement timeout fills a constant rather than inventing one. + # simply did not determine the answer within the bound. Pinned before it had a producer, so the + # gate that imposes the per-statement timeout filled a constant rather than inventing one. RULE_RESOURCE_LIMIT: "undetermined", RULE_UNPARSEABLE: "undetermined", RULE_MODEL_SAFETY: "undetermined", @@ -165,9 +169,12 @@ def refuse(rule: str, *, detail: str, remediation: str) -> Refusal: in connect, credential resolution or model load, where "narrow the query" is the wrong fix. So an unresponsive executor is `failed` / `timeout`, and its message says only that we stopped waiting (guardrail contract §3). A **per-statement** timeout is the other case — its subject IS the -statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT` — and nothing imposes one yet, which is -why that rule is pinned with no producer. Driver-level connect/login timeouts fold into the connect -failure the executor already reports as `auth` (exit 4), because that is what the driver raises. +statement, so it is a refusal carrying `RULE_RESOURCE_LIMIT`, and the executor now imposes one: a +watchdog cancels the statement through the driver and the outcome leaves the chokepoint on the +refusal channel rather than this one. The two therefore coexist and stay distinguishable in the audit +trail, which is the whole point of splitting them. Driver-level connect/login timeouts fold into the +connect failure the executor already reports as `auth` (exit 4), because that is what the driver +raises. `column_not_found`, `table_not_found`, `permission` and `network` are DECLARED BUT UNREACHABLE: producing them means parsing driver text, and sanitizing driver text belongs to the error-hardening diff --git a/plugins/agami/shared/readonly-grants.md b/plugins/agami/shared/readonly-grants.md index ce363302..9b0a419d 100644 --- a/plugins/agami/shared/readonly-grants.md +++ b/plugins/agami/shared/readonly-grants.md @@ -11,7 +11,7 @@ agami only ever runs **read-only SELECT** queries against your datasource: query ## What the role does and does not guarantee -The read-only role is the **primary, non-bypassable** guarantee of **integrity and confinement**: SELECT-only means **no write / DDL / `COPY` / file access / server-side network call / cross-database reach** — and that holds even if the app-layer guard were bypassed. It does **not** bound a **runaway** query (a recursive CTE or cartesian join), and it does **not** stop schema/metadata **recon**. Those are the **app layer's** job, not the role's — and only partly done today: the executor caps the **result-row count** (`AGAMI_SQL_MAX_ROWS`, default 1000) and a supervisor stops an executor that stops responding, but a **per-statement query timeout and error-text/recon hardening are not in place yet**. So don't read "read-only" as "time-bounded" — the role confines *what* a query can touch, and bounding *how much* it consumes is app-side. Until that bound ships, an operator who needs one should set a `statement_timeout` (or the engine's equivalent) on the role itself. +The read-only role is the **primary, non-bypassable** guarantee of **integrity and confinement**: SELECT-only means **no write / DDL / `COPY` / file access / server-side network call / cross-database reach** — and that holds even if the app-layer guard were bypassed. It does **not** bound a **runaway** query (a recursive CTE or cartesian join), and it does **not** stop schema/metadata **recon**. Those are the **app layer's** job, not the role's. The executor now bounds **how long** a statement runs (`AGAMI_SQL_TIMEOUT_S`, default 30 seconds — the statement is cancelled on the server and the caller gets a structured refusal, not a partial result) as well as **how many rows** it returns (`AGAMI_SQL_MAX_ROWS`, default 1000), with a supervisor behind both for an executor that stops responding altogether. **Error-text and recon hardening are still not in place.** Two residuals are worth knowing about the time bound: on **BigQuery** it is the server-side job timeout alone, because a BigQuery query is a job with no connection to cancel; and a *per-statement* bound is not a *concurrency* bound, so enough simultaneous queries can still load a database. An operator who wants a bound that does not depend on agami at all can still set `statement_timeout` (or the engine's equivalent) on the role itself. ## Creating the role @@ -134,8 +134,8 @@ bq add-iam-policy-binding --member="serviceAccount:agami-ro@.iam.gservi --role="roles/bigquery.dataViewer" : ``` -Here the two IAM roles above **are** the whole guarantee — BigQuery has no SQL-level role to scope further and no role-level statement timeout. Confinement comes from `dataViewer` (read-only); runaway bounding is app-side (and you can add a BigQuery custom quota / maximum-bytes-billed as extra defense). +Here the two IAM roles above **are** the whole guarantee — BigQuery has no SQL-level role to scope further and no role-level statement timeout. Confinement comes from `dataViewer` (read-only); runaway bounding is app-side, and BigQuery is the one engine where that bound is **server-side only** — a query here is a job that outlives the client, so agami sets the job's own timeout rather than cancelling a connection it does not have. A custom quota / maximum-bytes-billed is worthwhile extra defense on this engine in particular. ## SQLite / DuckDB -File-based — there's no user or role, so the **read-only file / read-only open is the whole guarantee** at this layer. Safety comes from agami's **read-only SQL guard** (it refuses anything that isn't a `SELECT`) plus **filesystem permissions**. DuckDB files are additionally opened in read-only mode; SQLite is not, so for a hard guarantee point agami at a **read-only copy** of the file, or mark the file read-only for the account agami runs as. (Runaway-query bounding is still app-side, as above — the file mode only stops writes.) +File-based — there's no user or role, so the **read-only file / read-only open is the whole guarantee** at this layer. Safety comes from agami's **read-only SQL guard** (it refuses anything that isn't a `SELECT`) plus **filesystem permissions**. DuckDB files are additionally opened in read-only mode; SQLite is not, so for a hard guarantee point agami at a **read-only copy** of the file, or mark the file read-only for the account agami runs as. (Runaway-query bounding is app-side, as above — the file mode only stops writes. On these two the cancel is genuine and in-process, so the time bound is as strong here as anywhere.) diff --git a/plugins/agami/skills/agami-deploy/bundle/agami.env.example b/plugins/agami/skills/agami-deploy/bundle/agami.env.example index 7d7bc235..88466443 100644 --- a/plugins/agami/skills/agami-deploy/bundle/agami.env.example +++ b/plugins/agami/skills/agami-deploy/bundle/agami.env.example @@ -37,6 +37,13 @@ AGAMI_ADMIN_PASSWORD= # Override the bundled-Postgres password if you like (it's never exposed, so the default is fine). # POSTGRES_PASSWORD= +# Optional: how long ONE statement may run, in whole SECONDS (unset = 30). This is the deployment's +# availability budget — a query that outruns it is cancelled on the warehouse and the caller is told +# to narrow it. Raise it if your warehouse is slow and your users would rather wait; lower it to keep +# a runaway query from holding a connection. The executor's other time bounds are derived from this +# one, so this is the single number to tune. +# AGAMI_SQL_TIMEOUT_S=30 # per-statement budget (default 30 seconds) + # Only for the `tunnel` profile: # CLOUDFLARE_TUNNEL_TOKEN= diff --git a/tests/test_ace035_guardrail_audit.py b/tests/test_ace035_guardrail_audit.py index e62c9d46..d81d7ab7 100644 --- a/tests/test_ace035_guardrail_audit.py +++ b/tests/test_ace035_guardrail_audit.py @@ -694,9 +694,10 @@ def test_killing_an_unresponsive_executor_is_a_failure_we_cannot_attribute(env, authority of a guardrail decision. `failed` / `timeout` claims exactly what we know — we stopped waiting — and the value-free message says only that. - `RULE_RESOURCE_LIMIT` stays declared and pinned in `REASON_FOR_RULE` for the per-statement bound - a later gate imposes, which IS a refusal because its subject is the statement. It has no producer - today, and this branch must not become one. + `RULE_RESOURCE_LIMIT` belongs to the per-statement bound the executor now imposes, which IS a + refusal because its subject is the statement. That the rule finally has a producer is precisely + why this branch is worth holding: the two bounds now coexist, and the one that cannot say what it + stopped must not drift into borrowing the vocabulary of the one that can. """ def _timed_out(cmd, **kwargs): raise subprocess.TimeoutExpired(cmd, kwargs.get("timeout", 240)) @@ -720,13 +721,15 @@ def _timed_out(cmd, **kwargs): assert row["datasource"] == PROFILE and row["question"] == QUESTION -def test_no_gate_produces_the_resource_limit_refusal(env, monkeypatch): - """`RULE_RESOURCE_LIMIT` is declared with NO producer, and the supervisor kill is where it would - creep back in — so drive that branch and assert the rule is absent from what the caller gets. +def test_the_supervisor_kill_does_not_borrow_the_resource_limit_refusal(env, monkeypatch): + """The supervisor kill is where `RULE_RESOURCE_LIMIT` would creep in — so drive that branch and + assert the rule is absent from what the caller gets, and from the row recording it. - The pin is on `REASON_FOR_RULE`, not on the absence: the rule stays in the contract's table so - the gate that eventually imposes a per-statement timeout fills a constant instead of inventing a - string. What must not happen is a branch borrowing it because the word "timeout" fits. + Renamed rather than retired: the rule DOES have a producer now (the executor's per-statement + deadline), so "no gate produces it" is no longer the claim. The claim that survives is the sharper + half and always was — a branch must not borrow this rule merely because the word "timeout" fits. + The pin is on `REASON_FOR_RULE`, not on the absence, which is why this test needed no weakening + when the producer landed: it asserts what this branch says, not what the contract contains. """ assert guardrail.RULE_RESOURCE_LIMIT in guardrail.REASON_FOR_RULE diff --git a/tests/test_ace035_no_enumeration.py b/tests/test_ace035_no_enumeration.py index 8bcb31cd..7c347c8b 100644 --- a/tests/test_ace035_no_enumeration.py +++ b/tests/test_ace035_no_enumeration.py @@ -15,18 +15,20 @@ schema-listing endpoint, reachable by anyone who can send one deliberately-wrong statement — which is the recon surface a later slice exists to close. -**This file is lock-in, not a fix.** All five gates are echo-only today; that was verified when they -were converted, and it is asserted here so a later reword cannot quietly turn a refusal into a -listing. Unlike the rest of this spec's tests, the property this one pins is expected to hold on -`446cc20` (the pre-conversion base) as well — and it does. It was checked by running this file's -model, canaries, vectors and scanner against a materialized `446cc20` tree, driving all five rules -through all four routes below and scanning the 20 refusal bodies the base produced: all clean. (The -file cannot be *collected* at that commit — `guardrail` does not exist there and the tool edge -speaks `{"error": {kind, remediation}}` rather than the Envelope — so the base run drops only the -two shape assertions, `status == "refused"` and `refusal.rule == …`, and keeps the scanner -verbatim.) A future failure here is therefore a live disclosure bug, not a conversion regression. - -**What this file covers.** Every field of the serialized tool-edge body, for the five refusal rules +**This file is lock-in, not a fix.** All five gates were echo-only when they were converted, and the +per-statement timeout that has since joined them is too; it is asserted here so a later reword cannot +quietly turn a refusal into a listing. Unlike the rest of this spec's tests, the property this one +pins is expected to hold on `446cc20` (the pre-conversion base) as well — and it does. It was checked +by running this file's model, canaries, vectors and scanner against a materialized `446cc20` tree, +driving the five rules that existed then through all four routes below and scanning the 20 refusal +bodies the base produced: all clean. (The `resource_limit` vector came later and has no counterpart +there — nothing imposed a per-statement bound at that commit. The file cannot be *collected* at that +commit either — `guardrail` does not exist there and the tool edge speaks +`{"error": {kind, remediation}}` rather than the Envelope — so the base run drops only the two shape +assertions, `status == "refused"` and `refusal.rule == …`, and keeps the scanner verbatim.) A future +failure here is therefore a live disclosure bug, not a conversion regression. + +**What this file covers.** Every field of the serialized tool-edge body, for the six refusal rules and for a `failed` — across both surfaces and both execution paths. `failure.message` is in scope because it is part of that body: it was the one field the scanner never saw, and it is where a PostgreSQL `HINT: Perhaps you meant to reference the column "orders.internal_ref".` arrives, which @@ -178,6 +180,11 @@ def declared(tmp_path, monkeypatch): monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(artifacts)) monkeypatch.setenv("DATASOURCE_URL__ACME", f"sqlite:///{warehouse}") + # Deliberately NOT a one-second budget for every test that takes this fixture. Only the + # `resource_limit` vector wants one, and it sets it for itself; the `failed` vector reaches the + # same warehouse under it, so a stall on a loaded runner would turn an expected database + # rejection into a timeout refusal and fail a test that has nothing to do with the clock. + monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) # Local, not hosted: the disk model is the one the gates use. (`model_unavailable` needs the # hosted signal and gets its own fixture.) monkeypatch.delenv("AGAMI_DB_URL", raising=False) @@ -297,18 +304,34 @@ def _route_http(sql: str, profile: str = PROFILE) -> dict: # The sentinel # --------------------------------------------------------------------------- +# A statement that passes every gate and then runs for minutes: a recursive CTE with two billion +# iterations, filtered on a subquery over the declared table so the vector is the shape a real +# runaway query has rather than a synthetic spin. `orders` and `id` are therefore the caller's own, +# and echoing them back is legitimate; the canaries stay unsent, which is what makes a refusal that +# reached for the model a failing assertion here. +_RUNAWAY_SQL = ( + "WITH RECURSIVE burn(n) AS (" + "SELECT 1 UNION ALL SELECT n + 1 FROM burn WHERE n < 2000000000" + ") SELECT count(n) AS c FROM burn WHERE n > (SELECT count(id) FROM orders)" +) + # One statement per rule, chosen so the rule under test is the FIRST gate to fire: -# read_only — a denied keyword, refused at the tool edge before a profile is resolved -# table_scope — an undeclared table (runs before the star ban and the column gate) -# select_star — a declared table, projected with `*` -# column_scope — a declared table, an undeclared column +# read_only — a denied keyword, refused at the tool edge before a profile is resolved +# table_scope — an undeclared table (runs before the star ban and the column gate) +# select_star — a declared table, projected with `*` +# column_scope — a declared table, an undeclared column +# resource_limit — the odd one out: every gate PASSES it, and the fixture's one-second budget stops +# it inside the executor. So this row scans the one refusal that is produced after +# the model has been loaded and consulted — the state in which a "helpful" message +# has the declared surface closest to hand. # `audit_trail` and `ref_no` are the caller's own inventions, so echoing them back is legitimate; -# nothing the MODEL declares appears in any of them except `orders`. +# nothing the MODEL declares appears in any of them except `orders` and `id`. VECTORS = ( (guardrail.RULE_READ_ONLY, "DELETE FROM orders"), (guardrail.RULE_TABLE_SCOPE, "SELECT ref_no FROM audit_trail"), (guardrail.RULE_SELECT_STAR, "SELECT * FROM orders"), (guardrail.RULE_COLUMN_SCOPE, "SELECT ref_no FROM orders"), + (guardrail.RULE_RESOURCE_LIMIT, _RUNAWAY_SQL), ) UNAVAILABLE_SQL = "SELECT id FROM orders" @@ -347,8 +370,10 @@ def _assert_echo_only(body: dict, sql: str) -> None: @pytest.mark.parametrize(("rule", "sql", "route"), _MATRIX, ids=_MATRIX_IDS) -def test_no_declared_name_the_caller_did_not_send_reaches_a_refusal(declared, rule, sql, route): - """Five rules — four here, `model_unavailable` below — across both surfaces and both execution +def test_no_declared_name_the_caller_did_not_send_reaches_a_refusal( + declared, rule, sql, route, monkeypatch +): + """Six rules — five here, `model_unavailable` below — across both surfaces and both execution paths. The fork column is not redundant with in-process: the child serializes the refusal to stderr and @@ -357,6 +382,13 @@ def test_no_declared_name_the_caller_did_not_send_reaches_a_refusal(declared, ru by construction; the child's own read-only refusal crossing the wire is covered by `tests/test_ace035_read_only_refusal.py::test_parent_reconstructs_the_child_refusal`.) """ + if rule == guardrail.RULE_RESOURCE_LIMIT: + # The smallest per-statement budget the resolver accepts, scoped to the one vector that needs + # it: this is the only statement here that passes every gate, reaches the executor, and has + # to be stopped there. The forked and stdio routes inherit `os.environ`, so setting it here + # reaches them too. + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "1") + body = ROUTES[route](sql) assert body["status"] == "refused", body @@ -706,13 +738,6 @@ def test_echoing_the_callers_own_identifier_is_allowed(): "than refusing, so nothing constructs this rule yet; turning that fail-open into a refusal " "is the unparseable-statement slice's job, and the refusal it introduces needs a vector here." ), - guardrail.RULE_RESOURCE_LIMIT: ( - "No producer. It is reserved for the per-statement timeout a later gate imposes, and the " - "one branch that might borrow it — the subprocess supervisor killing an unresponsive " - "executor — is a `failed`/`timeout` instead, pinned by " - "test_ace035_guardrail_audit.py::test_no_gate_produces_the_resource_limit_refusal. A rule " - "nothing emits cannot leak; the day something emits it, it needs a vector here." - ), guardrail.RULE_MODEL_SAFETY: ( "Reachable, but its detail is authored as static prose at a single construction site in " "`execute_sql.execute_guarded` — the unconverted `_model_safety` branches hand back a bare " @@ -744,7 +769,7 @@ def test_every_rule_and_every_route_is_covered(): assert not covered & set(_NO_VECTOR) assert all(reason.strip() for reason in _NO_VECTOR.values()) assert set(ROUTES) == {"in_process", "fork", "stdio", "http"} - assert len(_MATRIX) == len(VECTORS) * len(ROUTES) == 16 + assert len(_MATRIX) == len(VECTORS) * len(ROUTES) == 20 # The `failed` channel is covered by its own matrix rather than this one, because it has no # rule. Its vector must not be one of these: a statement that a gate refuses would report on the # refusal channel a second time and leave `failure.message` unscanned again, which is exactly diff --git a/tests/test_ace038_timeout.py b/tests/test_ace038_timeout.py new file mode 100644 index 00000000..dcf29c7b --- /dev/null +++ b/tests/test_ace038_timeout.py @@ -0,0 +1,2119 @@ +"""Per-statement timeout — config resolution, the deadline primitive, and the refusal it produces. + +`_resolve_timeout_s` answers "how long may one statement run", from `AGAMI_SQL_TIMEOUT_S` and a 30s +default and from nothing else — one configuration surface, so the parent and the forked child reach +the same number. Unlike the row cap it complains out loud whenever the budget it returns is not the +one the operator wrote, whether the text was unreadable or merely declined. `_deadline` is the +watchdog those seconds feed: it fires an Event and calls a cancel callable when a block outlives its +budget, and disarms cleanly when it does not — where "cleanly" means the disarm is a JOIN, so a +watchdog that loses the race lands nowhere at all. + +The second half of this file proves the whole contract end to end on ONE engine — SQLite, chosen +because it is in-process, needs no network, and `sqlite3.Connection.interrupt()` is a genuine cancel +rather than a polite request. A statement that outlives its budget is cancelled, unwinds on the +internal `_ResourceLimit` marker, and leaves `execute_guarded` as a `refused` Envelope carrying +`resource_limit` — with no partial data, a detail that quotes the configured budget, and a +remediation addressed to whoever can actually act on it. + +What no test here can show is that the cancel reached a SERVER: SQLite is in-process, and on a +client/server engine an abandoned statement produces the identical Envelope while the backend keeps +running. That assertion needs a live database and a second connection, and lives in +`test_postgres_timeout_integration.py`, which skips unless one is configured. + +**The classification is the FLAG, and only the flag.** A cancelled SQLite statement raises +`OperationalError("interrupted")`, so neither the error text nor the elapsed clock can be the test: +both are properties an ordinary database error can have by coincidence, and reading either one would +mean an unlucky query gets told to narrow itself when nothing timed out. `_deadline` sets its Event +*before* the cancel lands, so "did WE stop this?" is answerable without inference — and it is +asserted here in both directions. + +The last section is the OUTER bound, which is what makes the limit apply to every executor rather +than to the built-in one alone. An injected executor (the hosted connection-reuse path) carries none +of the engine watchdogs, and even the built-in BigQuery path has no cancel to arm — so +`execute_guarded` bounds `executor.execute` itself, on a worker thread it can stop waiting for. The +four bounds are one ordered family (watchdog < native < outer < supervisor) resolved from one budget, +and the order is asserted as a single fact so no future change can let them drift apart. +""" + +from __future__ import annotations + +import contextlib +import json +import logging +import os +import sqlite3 +import subprocess +import sys +import threading +import time +import types +from contextvars import ContextVar +from pathlib import Path +from types import SimpleNamespace + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +PKG_SRC = REPO_ROOT / "packages" / "agami-core" / "src" +if str(PKG_SRC) not in sys.path: + sys.path.insert(0, str(PKG_SRC)) + +import execute_sql # noqa: E402 +import guardrail # noqa: E402 + +# Short enough that the whole file stays sub-second, long enough that a loaded CI runner still +# schedules the timer thread before the assertion runs. +_TINY = 0.05 + + +@pytest.fixture(autouse=True) +def _reset_row_cap(): + # The row cap is a request-scoped ContextVar; isolate every test from it. It matters more than it + # used to because the outer bound copies the caller's whole context into its worker, so a cap + # left set by one test is now visible to the next test's executor as well. + execute_sql._max_rows_override.set(None) + yield + execute_sql._max_rows_override.set(None) + + +@pytest.fixture(autouse=True) +def _drain_abandoned_workers(): + """Wait for any worker a test abandoned to finish, so its slot is back before the next test runs. + + The counter is process-wide and released by the abandoned worker itself, which returns shortly + after the test releases it. Draining here rather than zeroing the counter keeps the accounting + honest: setting it to 0 while a worker was still running would let that worker decrement past + zero and hand a later test a cap it has not actually got. + """ + yield + deadline = time.monotonic() + 5 + while execute_sql._abandoned_workers and time.monotonic() < deadline: + time.sleep(0.01) + + +@pytest.fixture(autouse=True) +def _clear_env(monkeypatch): + # The suite must not inherit an operator's real budget from the ambient environment. + monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) + + +@pytest.fixture(autouse=True) +def _no_injected_executor(): + """`_INJECTED_EXECUTOR` is a process global, and the tests below install one both directly and + (via `create_app`) as a side effect of building the HTTP app. Reset it around every test so + neither leaks into the next one and quietly moves it onto the other execution path.""" + try: + import tools + except Exception: + yield + return + tools.set_injected_executor(None) + yield + tools.set_injected_executor(None) + + +# -------------------------------------------------------------------------------------------- +# _resolve_timeout_s +# -------------------------------------------------------------------------------------------- + + +def test_an_absent_env_var_yields_the_default_budget(): + assert execute_sql._resolve_timeout_s() == 30 + assert execute_sql._DEFAULT_TIMEOUT_S == 30 + + +@pytest.mark.parametrize("raw", ["1", "5", "45", "600"]) +def test_a_valid_env_value_is_honoured(monkeypatch, raw): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + assert execute_sql._resolve_timeout_s() == int(raw) + + +def test_surrounding_whitespace_does_not_defeat_a_valid_value(monkeypatch): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", " 45 ") + assert execute_sql._resolve_timeout_s() == 45 + + +@pytest.mark.parametrize("raw", ["0", "00", "-5"]) +def test_a_non_positive_value_falls_back_to_the_default(monkeypatch, raw): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + assert execute_sql._resolve_timeout_s() == 30 + + +# `²` and `①` are the cases `str.isdigit()` waves through and `int()` then refuses. They are not a +# curiosity: the resolver is called at the fork path's supervisor bound, outside any handler, so a +# ValueError raised out of it escapes the tool edge as a traceback rather than as a budget. +@pytest.mark.parametrize("raw", ["6O", "30s", "45.5", "thirty", "1e3", "²", "①"]) +def test_an_unreadable_value_falls_back_and_says_so(monkeypatch, caplog, raw): + """The row cap falls back silently; this one must not. An operator who typed a capital O for a + zero has to be able to find out why their budget is not what they configured.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + assert execute_sql._resolve_timeout_s() == 30 + + warnings = [r for r in caplog.records if r.levelno == logging.WARNING] + assert warnings, f"no warning emitted for the rejected value {raw!r}" + assert any(raw in r.getMessage() for r in warnings), ( + f"the warning must name the rejected text {raw!r}; got {[r.getMessage() for r in warnings]}" + ) + + +@pytest.mark.parametrize("raw", ["0", "00", "-5", "-0"]) +def test_a_value_we_read_and_declined_says_so_too(monkeypatch, caplog, raw): + """The warning is about the OUTCOME, not about the parse. + + `-5` reads perfectly and is then declined, so it used to slip past silently while `6O` warned — + and a deployment quietly running 30s when its operator wrote something else is precisely the + invisible degradation the warning exists against. Whether we could read the text is not the + question the operator has; whether they got what they asked for is. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + assert execute_sql._resolve_timeout_s() == 30 + + warnings = [r for r in caplog.records if r.levelno == logging.WARNING] + assert warnings, f"no warning emitted for the declined value {raw!r}" + assert any(raw in r.getMessage() for r in warnings), ( + f"the warning must name the declined text {raw!r}; got {[r.getMessage() for r in warnings]}" + ) + + +@pytest.mark.parametrize("raw", ["", "45", "045", " 45 ", "600"]) +def test_a_value_the_operator_actually_got_stays_quiet(monkeypatch, caplog, raw): + """The other direction, so the warning cannot become noise. A leading zero or surrounding + whitespace still yields exactly the number that was written, and warning on it would train + operators to ignore the log.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", raw) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + execute_sql._resolve_timeout_s() + assert [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] == [] + + +def test_the_budget_has_exactly_one_configuration_surface(): + """The environment is the only source, and that is load-bearing rather than tidy. + + A request-scoped override outranks the environment in THIS process and cannot cross a fork: the + child re-resolves from `os.environ` alone. So a parent that derived the supervisor's bound from + an override would compute a bound BELOW the budget the child actually enforces, the supervisor + would fire first, and the ordered family the whole design rests on would invert — turning a + precise refusal into a `failed`/`timeout` that names nothing. Asserted structurally, because the + hazard is the existence of a second surface, not any particular value in one. + """ + context_vars = { + name for name, value in vars(execute_sql).items() if isinstance(value, ContextVar) + } + assert context_vars == {"_max_rows_override"}, ( + "a second, higher-precedence configuration surface for the budget cannot cross the fork; " + f"found {sorted(context_vars)}" + ) + + +def test_the_supervisor_bound_exceeds_the_budget_a_real_child_resolves(monkeypatch): + """The inversion, driven across the actual process boundary rather than reasoned about. + + The parent computes the supervisor's bound; a real forked interpreter, inheriting this + environment exactly as `subprocess.run` gives it one, resolves its own budget. The first must + exceed the second, or the outermost bound fires before the innermost. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "300") + parent_bound = execute_sql._resolve_timeout_s() + execute_sql._SUPERVISOR_SKEW_S + + child = subprocess.run( + [sys.executable, "-c", "import execute_sql; print(execute_sql._resolve_timeout_s())"], + capture_output=True, text=True, timeout=60, + env={**os.environ, "PYTHONPATH": str(PKG_SRC)}, + ) + assert child.returncode == 0, child.stderr + child_budget = int(child.stdout.strip()) + + assert child_budget == 300, child.stdout + assert parent_bound > child_budget, ( + f"the supervisor stops waiting at {parent_bound}s while the child is still inside a " + f"{child_budget}s budget" + ) + + +# -------------------------------------------------------------------------------------------- +# _deadline +# -------------------------------------------------------------------------------------------- + + +class _RecordingCancel: + """A stand-in for a driver's `cancel()`, recording the calls and the order relative to the flag.""" + + def __init__(self, fails: bool = False): + self.calls = 0 + self.fails = fails + self.done = threading.Event() + + def __call__(self) -> None: + self.calls += 1 + try: + if self.fails: + raise RuntimeError("driver refused to cancel from another thread") + finally: + self.done.set() + + +def test_an_overrunning_block_sets_the_event_and_cancels(): + cancel = _RecordingCancel() + with execute_sql._deadline(cancel, _TINY) as fired: + assert cancel.done.wait(2.0), "the watchdog never ran" + # The flag must already be readable by the time the cancel lands, so a caller catching the + # resulting driver error can attribute it to us rather than to the database. + assert fired.is_set() + assert cancel.calls == 1 + + +def test_a_block_that_finishes_first_neither_fires_nor_cancels(): + cancel = _RecordingCancel() + with execute_sql._deadline(cancel, 30) as fired: + pass + assert not fired.is_set() + assert cancel.calls == 0 + + +def test_the_timer_is_disarmed_on_exit_so_no_late_cancel_arrives(): + """The watchdog must not outlive its block: a cancel landing after the statement finished would + kill whatever the connection is doing next.""" + cancel = _RecordingCancel() + with execute_sql._deadline(cancel, _TINY) as fired: + pass + time.sleep(_TINY * 4) # comfortably past when an un-disarmed timer would have fired + assert cancel.calls == 0 + assert not fired.is_set() + + +def test_the_timer_is_disarmed_even_when_the_block_raises(): + cancel = _RecordingCancel() + with pytest.raises(ValueError): + with execute_sql._deadline(cancel, _TINY): + raise ValueError("the statement blew up on its own") + time.sleep(_TINY * 4) + assert cancel.calls == 0 + + +class _ManualTimer: + """A `threading.Timer` stand-in whose `cancel()` does nothing and whose callback the test fires. + + Not a caricature of the real one — a faithful model of its worst case. `Timer.cancel()` only sets + the timer's internal `finished` Event, so a timer thread that has ALREADY passed its own + `if not self.finished.is_set()` check runs the callback regardless. This makes that interleaving + — disarm first, `fire` second — happen on every run rather than once in a thousand. + """ + + def __init__(self, interval, function): + self.interval = interval + self.function = function + self.started = False + + def start(self) -> None: + self.started = True + + def cancel(self) -> None: + pass + + +def test_a_watchdog_that_loses_the_race_to_the_disarm_lands_nowhere(monkeypatch): + """The disarm has to be a JOIN, not a request, and `Timer.cancel()` is only ever a request. + + Two things go wrong when a `fire` that lost the race still runs. Every engine re-reads the flag + once the block has exited — "checked after the watchdog is disarmed, so the flag is final" — and + a flag set afterwards makes a completed statement look like a refusal. And the cancel itself + lands on a connection the engine has moved on from; on a pooled one that is somebody else's + statement being killed by our watchdog. + """ + timers: list = [] + + def _timer(interval, function): + timers.append(_ManualTimer(interval, function)) + return timers[-1] + + monkeypatch.setattr(execute_sql.threading, "Timer", _timer) + cancel = _RecordingCancel() + + with execute_sql._deadline(cancel, _TINY) as fired: + pass + + assert timers and timers[0].started, "no watchdog was armed" + timers[0].function() # the timer thread had already committed; `fire` runs anyway + + assert cancel.calls == 0, "a cancel landed after the block the deadline was bounding had exited" + assert not fired.is_set(), "the flag every engine re-reads after the block was not final" + + +def test_a_watchdog_that_fires_while_the_timer_is_being_disarmed_lands_nowhere(monkeypatch): + """The narrowest interleaving of the same race: `fire` lands DURING the disarm itself. + + `Timer.cancel()` is not a join, so the disarm has two steps that are not one atomic act: stand the + timer down, and claim the race under the lock. Do them in that order and the gap between them is + a window in which an already-expired `fire` takes the lock first, sets the flag, and marks a + statement that in fact completed. Claiming the race first closes it. This timer fires its + callback from inside `cancel()`, which puts a `fire` in exactly that gap on every run. + """ + cancel = _RecordingCancel() + + class _FiresWhileBeingCancelled(_ManualTimer): + def cancel(self) -> None: + self.function() + + monkeypatch.setattr( + execute_sql.threading, "Timer", lambda interval, function: _FiresWhileBeingCancelled(interval, function) + ) + + with execute_sql._deadline(cancel, _TINY) as fired: + pass + + assert cancel.calls == 0, "a cancel landed in the gap between standing the timer down and disarming" + assert not fired.is_set(), "a statement that completed was flagged by a watchdog firing mid-disarm" + + +def test_the_disarm_waits_for_a_cancel_already_in_flight(): + """The other half of the same guarantee: a `fire` that WON the race is finished with before the + block is allowed to return. + + Otherwise "no late cancel arrives" is only true of the cancels that had not started yet, and a + driver cancel that takes a moment still reaches into whatever the connection does next. + """ + entered = threading.Event() + release = threading.Event() + exited = threading.Event() + + def cancel() -> None: + entered.set() + release.wait(5) + + def run() -> None: + with execute_sql._deadline(cancel, _TINY): + assert entered.wait(5), "the watchdog never ran" + exited.set() + + runner = threading.Thread(target=run, daemon=True) + runner.start() + assert entered.wait(5), "the watchdog never ran" + + assert not exited.wait(0.3), "the block returned while its own cancel was still running" + release.set() + runner.join(5) + assert exited.is_set() + + +def _wait_for_warning(caplog, needle: str, timeout_s: float = 5.0) -> bool: + """Poll the captured records for a warning containing `needle`, up to a deadline.""" + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if any(needle in r.getMessage() + for r in caplog.records if r.levelno == logging.WARNING): + return True + time.sleep(0.005) + return False + + +def test_a_cancel_that_raises_does_not_escape_the_timer_thread(caplog): + """Some drivers raise when cancelled from a thread other than the one running the statement. That + must be logged and swallowed: an exception escaping a timer thread is unhandleable by the caller + and would be lost to threading's excepthook. + + Polled rather than slept on. `cancel.done` is set in a `finally` that runs BEFORE the raise + reaches `fire`'s handler, so nothing this test can wait on is synchronized with the log write, + and a fixed sleep is a wager on how the runner happened to schedule two threads. + """ + cancel = _RecordingCancel(fails=True) + with caplog.at_level(logging.WARNING, logger=execute_sql._LOG.name): + with execute_sql._deadline(cancel, _TINY) as fired: + assert cancel.done.wait(2.0), "the watchdog never ran" + assert _wait_for_warning(caplog, "driver refused to cancel"), ( + f"the failed cancel was not logged; got {[r.getMessage() for r in caplog.records]}" + ) + + assert fired.is_set() # the timeout still counts as fired even though the cancel failed + assert cancel.calls == 1 + + +def test_the_watchdog_thread_is_a_daemon_and_does_not_hold_the_process_open(): + """A hung cancel must never keep the interpreter alive at shutdown.""" + live_before = {t.ident for t in threading.enumerate()} + with execute_sql._deadline(_RecordingCancel(), 300) as fired: + new = [t for t in threading.enumerate() if t.ident not in live_before] + assert new, "no watchdog thread was started" + assert all(t.daemon for t in new) + assert not fired.is_set() + + +# -------------------------------------------------------------------------------------------- +# _ResourceLimit +# -------------------------------------------------------------------------------------------- + + +def test_the_resource_limit_marker_is_a_plain_exception(): + """It has to unwind an engine function like any other error so the transaction rolls back.""" + assert issubclass(execute_sql._ResourceLimit, Exception) + with pytest.raises(execute_sql._ResourceLimit): + raise execute_sql._ResourceLimit("statement exceeded its budget") + + +# -------------------------------------------------------------------------------------------- +# The deadline, wired into SQLite — the refusal, end to end +# -------------------------------------------------------------------------------------------- + +PROFILE = "analytics" + +# The smallest budget the resolver accepts, since it deals in whole seconds. Every end-to-end test +# below therefore costs about a second of wall clock, which is the price of driving a real cancel +# through a real driver rather than asserting on a stub. +_BUDGET_S = 1 + +# A recursive CTE with no physical table and no termination in reach: two billion iterations, which +# on this machine is minutes of pure CPU. Sized that way on purpose — "bounded rather than hanging" +# is only proved if the unbounded run would take far longer than the assertion allows. +# +# It also has to reach the executor, so it is written to pass every gate ahead of it: it opens with +# WITH…SELECT (read-only), names no physical table (`burn` is a CTE, which the table-scope gate +# excludes by construction), projects no star, and its one column binds to a CTE rather than to a +# declared table. +_RUNAWAY_SQL = ( + "WITH RECURSIVE burn(n) AS (" + "SELECT 1 UNION ALL SELECT n + 1 FROM burn WHERE n < 2000000000" + ") SELECT count(n) AS c FROM burn" +) + + +@pytest.fixture +def warehouse(tmp_path, monkeypatch): + """A real SQLite warehouse reachable as profile `analytics`, and nothing else configured. + + `no_safety=True` on the direct calls below skips the semantic-model pass, so this fixture only + has to satisfy credential resolution. The audit test further down needs the fuller install and + builds it itself. + """ + path = tmp_path / "warehouse.db" + con = sqlite3.connect(path) + con.execute("CREATE TABLE orders (id INTEGER)") + con.commit() + con.close() + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", f"sqlite:///{path}") + monkeypatch.delenv("AGAMI_DB_URL", raising=False) + monkeypatch.delenv("APP_DATABASE_URL", raising=False) + return path + + +def _guarded(sql: str) -> object: + """`execute_guarded` over the built-in executor — the single chokepoint, driven directly.""" + return execute_sql.execute_guarded( + sql, PROFILE, None, executor=execute_sql.BUILTIN_EXECUTOR, no_safety=True + ) + + +def test_a_runaway_statement_is_cancelled_rather_than_left_to_run(warehouse, monkeypatch): + """The headline: a statement that would run for minutes is stopped at its budget and comes back + as a refusal naming the rule the contract reserves for a bound we imposed. + + The elapsed assertion is the one that would still fail if the deadline were never armed — without + it a test that merely waited out the query would look identical and pass in several minutes. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + started = time.monotonic() + env = _guarded(_RUNAWAY_SQL) + elapsed = time.monotonic() - started + + assert elapsed < 20, f"the statement ran {elapsed:.1f}s against a {_BUDGET_S}s budget" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + # Neither unsafe nor out of scope: we simply did not determine the answer within the bound. + assert env.refusal.reason == "undetermined" + assert env.refusal.reason == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + + +def test_a_cancelled_statement_yields_no_partial_data(warehouse, monkeypatch): + """Whatever rows the engine had gathered when the watchdog fired are not an answer. + + A truncated result presented as a result is the failure mode the bounded-fetch work already + guards against on the row axis; on the time axis the answer is stronger — there is no data at + all, and `Envelope.__post_init__` enforces that a refusal cannot carry any. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + env = _guarded(_RUNAWAY_SQL) + + assert env.status == "refused" + assert env.data is None + assert env.failure is None + + +class _ResourceLimitExecutor: + """A `ports.Executor` that raises the internal marker, standing in for an engine whose watchdog + fired. Used where the subject is the REFUSAL TEXT rather than the cancel, so the assertion does + not have to pay a second of wall clock to reach it.""" + + def execute(self, vetted_sql: str, creds: dict, *, profile: str): + raise execute_sql._ResourceLimit("the statement outlived its per-statement budget") + + +def test_the_detail_quotes_the_configured_budget(warehouse, monkeypatch): + """A bound the caller cannot see is one it cannot plan around, so the number is in the detail. + + The configured value is not a data value: it is a deployment setting, and stating it discloses + nothing about the database or its contents. Asserted against a distinctive budget rather than the + default, so a hard-coded `30s` in the message cannot pass. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(7)) + + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_ResourceLimitExecutor(), no_safety=True, + ) + + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert "7s" in env.refusal.detail, env.refusal.detail + + +def test_the_remediation_names_no_deployment_environment_variable(warehouse, monkeypatch): + """The remediation has to be addressed to whoever is reading it. + + On the served path that is an assistant holding a statement, with no shell, no deployment and no + way to set an environment variable — so "raise AGAMI_SQL_TIMEOUT_S" is advice aimed past the + caller at an operator who is not in the conversation, and it reads as a fix while being + unfollowable. What is left has to be something that would make THIS statement executable. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_ResourceLimitExecutor(), no_safety=True, + ) + + authored = f"{env.refusal.detail} {env.refusal.remediation}" + assert "AGAMI_SQL_TIMEOUT_S" not in authored, authored + assert "AGAMI_" not in authored, authored # no sibling deployment var either + assert "environ" not in authored and "env var" not in authored.lower(), authored + # And it is still actionable — it names something to change about the statement. + assert env.refusal.remediation.strip() + + +# -------------------------------------------------------------------------------------------- +# The clock covers the fetch, not just the execute +# -------------------------------------------------------------------------------------------- + + +class _SlowFetchCursor: + """A cursor whose `execute` returns at once and whose single `fetchmany` blocks until cancelled. + + That is the shape of a streaming / server-side cursor, where the scan happens on the PULL rather + than on the call: `_collect_cursor` issues one `fetchmany(cap + 1)`, and on such a cursor that + one call is the whole query. A deadline that stopped at `execute` would bound the cheap half. + """ + + def __init__(self, cancelled: threading.Event): + self.description = [("c",)] + self._cancelled = cancelled + self.executed: str | None = None + + def execute(self, sql: str) -> None: + self.executed = sql + + def fetchmany(self, n: int): + # Bounded so a deadline that never reaches the fetch fails this test rather than hanging it. + if not self._cancelled.wait(30): + raise AssertionError("the fetch was never cancelled — the clock stopped at execute") + raise sqlite3.OperationalError("interrupted") + + +class _FakeConnection: + """The driver surface `_run_sqlite` uses: a cursor, a real `interrupt`, and a close.""" + + def __init__(self, cursor_factory): + self.cancelled = threading.Event() + self.cursor_obj = cursor_factory(self.cancelled) + self.closed = False + + def cursor(self): + return self.cursor_obj + + def interrupt(self) -> None: + self.cancelled.set() + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def fake_sqlite(monkeypatch): + """Replace `sqlite3.connect` so a test can choose exactly where the time goes. + + A real slow query cannot separate execute from fetch — sqlite decides that — and driving one + would also make these tests as slow as the thing they measure. `_run_sqlite` does its own + `import sqlite3`, which resolves through `sys.modules`, so patching the module attribute is + enough to reach it. + """ + def _install(cursor_factory, *, connect_delay_s: float = 0.0): + conn = _FakeConnection(cursor_factory) + + def _connect(path, *a, **kw): + if connect_delay_s: + time.sleep(connect_delay_s) + return conn + + monkeypatch.setattr(sqlite3, "connect", _connect) + return conn + + return _install + + +def test_a_slow_fetch_is_bounded_even_when_the_execute_returned_at_once( + warehouse, fake_sqlite, monkeypatch +): + """The clock covers the whole statement, fetch included — an explicit criterion, not a bonus. + + Bounding only `execute` would leave the common streaming shape unbounded: the driver returns + immediately and the engine scans while the caller pulls. The cancel has to land on the fetch, and + the refusal has to be the same one a slow execute produces. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + conn = fake_sqlite(_SlowFetchCursor) + + started = time.monotonic() + env = _guarded("SELECT c FROM orders") + elapsed = time.monotonic() - started + + assert conn.cursor_obj.executed == "SELECT c FROM orders" # the execute really did run, and fast + assert elapsed < 20, f"the fetch ran {elapsed:.1f}s against a {_BUDGET_S}s budget" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert conn.closed # the connection is still released on the refusing path + + +# -------------------------------------------------------------------------------------------- +# The classification is the flag alone — asserted in both directions +# -------------------------------------------------------------------------------------------- +# +# `interrupt()` makes the in-flight statement raise `OperationalError("interrupted")`. That text is +# therefore NOT evidence of a timeout — a database is free to raise it for its own reasons, and an +# error that merely arrives late is not one we caused. Both vectors below carry that exact text with +# the flag unset, so anything keying on the message or on the clock would turn them green as +# refusals; only the flag distinguishes them. + + +class _ImmediateErrorCursor: + """Raises the very error a cancel provokes, straight away and unprompted.""" + + def __init__(self, cancelled: threading.Event): + self.description = None + self._cancelled = cancelled + + def execute(self, sql: str) -> None: + raise sqlite3.OperationalError("interrupted") + + def fetchmany(self, n: int): # pragma: no cover - execute raises first + raise AssertionError("unreachable") + + +def test_a_database_error_with_the_flag_unset_is_a_failure_not_a_refusal( + warehouse, fake_sqlite, monkeypatch +): + """Direction (a): the watchdog never fired, so this is the database's outcome, not ours. + + A generous budget means the flag stays clear while the identical error text arrives. It must + unwind as an `ExecutorError` and leave the chokepoint as `failed`/`syntax` — a refusal here would + tell a caller to narrow a statement that never ran long at all. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(300)) + fake_sqlite(_ImmediateErrorCursor) + + env = _guarded("SELECT c FROM orders") + + assert env.status == "failed" + assert env.failure.kind == "syntax" + assert env.refusal is None + + +@contextlib.contextmanager +def _deadline_that_never_fires(cancel, timeout_s): + """A watchdog that does not fire, whatever the block does. + + Substituted so a block can genuinely outlive its budget with the flag clear — the one state a + real `_deadline` cannot be put into, and the one that separates "the clock ran out" from "we + cancelled it". + """ + yield threading.Event() + + +class _LateErrorCursor: + """Runs past the budget, then raises the error a cancel would have provoked.""" + + def __init__(self, cancelled: threading.Event): + self.description = None + self._cancelled = cancelled + + def execute(self, sql: str) -> None: + time.sleep(_BUDGET_S * 1.5) + raise sqlite3.OperationalError("interrupted") + + def fetchmany(self, n: int): # pragma: no cover - execute raises first + raise AssertionError("unreachable") + + +def test_an_error_that_merely_arrives_late_is_still_a_failure(warehouse, fake_sqlite, monkeypatch): + """Direction (b): elapsed time is not the classifier either. + + Everything a clock-based or text-based reading would need is present — the budget is exceeded and + the message is literally "interrupted" — and the flag is clear. The verdict must still be + `failed`, because we did not stop this statement; it stopped on its own, late. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_that_never_fires) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + fake_sqlite(_LateErrorCursor) + + started = time.monotonic() + env = _guarded("SELECT c FROM orders") + elapsed = time.monotonic() - started + + assert elapsed > _BUDGET_S, "the statement did not actually outlive its budget" + assert env.status == "failed" + assert env.failure.kind == "syntax" + assert env.refusal is None + + +# -------------------------------------------------------------------------------------------- +# The refusal is recorded +# -------------------------------------------------------------------------------------------- + +# A neutral hostname and a throwaway secret, only so the HTTP surface below can mint a token. +_BASE_URL = "https://your-host.example.com" +_SIGNING_SECRET = "x" * 40 + + +@pytest.fixture +def audited(tmp_path, monkeypatch): + """A complete single-datasource install: an app database to audit into, a semantic model on + disk, and the same real warehouse. + + Fuller than the `warehouse` fixture because this one goes through the tool edge rather than + straight to `execute_guarded`, so the model pass runs for real — and with an app database + configured the executor reads that as the hosted signal and fails closed without a model. + """ + pytest.importorskip("pydantic") + pytest.importorskip("sqlglot") + yaml = pytest.importorskip("yaml") + from store import Store + + app_db = "sqlite://" + str(tmp_path / "app.db") + store = Store.connect(app_db) + store.run_migrations() + store.close() + + root = tmp_path / "artifacts" / PROFILE + (root / "subject_areas" / "sales" / "tables").mkdir(parents=True) + (root / "datasource.yaml").write_text(yaml.safe_dump( + {"datasource": "Shop", "version": 1, "subject_areas": ["subject_areas/sales"]})) + (root / "subject_areas" / "sales" / "subject_area.yaml").write_text(yaml.safe_dump( + {"name": "sales", "tables": [ + {"storage_connection": "c", "schema": "public", "table": "orders"}]})) + (root / "subject_areas" / "sales" / "tables" / "orders.yaml").write_text(yaml.safe_dump({ + "name": "orders", "schema": "public", "storage_connection": "c", "grain": ["id"], + "description": "orders", + "columns": [{"name": "id", "type": "integer", "primary_key": True}], + })) + + path = tmp_path / "warehouse.db" + con = sqlite3.connect(path) + con.execute("CREATE TABLE orders (id INTEGER)") + con.commit() + con.close() + + monkeypatch.setenv("AGAMI_DB_URL", app_db) + monkeypatch.delenv("APP_DATABASE_URL", raising=False) + monkeypatch.delenv("AGAMI_ORG_ID", raising=False) + monkeypatch.setenv("AGAMI_ARTIFACTS_DIR", str(tmp_path / "artifacts")) + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", f"sqlite:///{path}") + # The HTTP surface needs both to mint and verify its bearer token; the stdio surface ignores + # them. Set here so one install serves both transports. + monkeypatch.setenv("PUBLIC_BASE_URL", _BASE_URL) + monkeypatch.setenv("AGAMI_SIGNING_SECRET", _SIGNING_SECRET) + return SimpleNamespace(app_db=app_db) + + +def test_the_refusal_is_written_to_the_audit_trail(audited, monkeypatch): + """A decision we made against a caller's statement is exactly the row a reviewer comes looking + for, so this outcome is audited like every other — and keyed by the id the caller was handed. + + Driven through the tool edge with the built-in executor injected, which is the in-process path + the hosted server runs: one real cancel, one real serializer, one real sink. + """ + import tools + from store import Store + + tools.set_injected_executor(execute_sql.BUILTIN_EXECUTOR) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + try: + body = json.loads(tools.tool_execute_sql({"sql": _RUNAWAY_SQL, "datasource": PROFILE, + "raw_query": "how many"})) + finally: + tools.set_injected_executor(None) + + assert body["status"] == "refused", body + assert body["refusal"]["rule"] == guardrail.RULE_RESOURCE_LIMIT, body + + store = Store.connect(audited.app_db) + try: + rows = store.query("SELECT id, status, reason, rule FROM query_executions") + finally: + store.close() + + assert len(rows) == 1, rows + (row,) = rows + assert row["id"] == body["audit_id"] # the answer and its record name the same id + assert row["status"] == "refused" + assert row["rule"] == guardrail.RULE_RESOURCE_LIMIT + assert row["reason"] == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + + +# -------------------------------------------------------------------------------------------- +# The outer bound — the limit reaches every executor, not only the built-in one +# -------------------------------------------------------------------------------------------- +# +# The engine watchdogs live INSIDE the built-in executor, so an injected one — which is what the +# hosted connection-reuse path supplies — ran with no per-statement bound at all, and BigQuery has no +# watchdog even on the built-in path. `execute_guarded` therefore bounds `executor.execute` itself. +# That layer holds no connection and can cancel nothing; it runs the call on a daemon worker and +# stops WAITING, which is why it must be the outer one and the watchdog must always win. + + +class _BlockingExecutor: + """A `ports.Executor` whose `execute` never returns on its own — the shape the outer bound + exists for. Nothing about it is cancellable from outside, which is exactly the point: an + arbitrary injected executor exposes one method and no handle on the work behind it.""" + + def __init__(self) -> None: + self.calls = 0 + self.entered = threading.Event() + self.release = threading.Event() + + def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> execute_sql.ExecResult: + self.calls += 1 + self.entered.set() + # Released by the test, so the worker this deliberately leaks does not outlive the test that + # made it. A real one would sit here until its driver call returned on its own. + self.release.wait(10) + return execute_sql.ExecResult(columns=["c"], rows=[(1,)], truncated=False) + + +def test_an_injected_executor_that_never_returns_is_still_bounded(warehouse, monkeypatch): + """The headline for this layer: an executor with no watchdog of its own, no cancel and no + intention of returning still yields a `resource_limit` refusal at a bound we set. + + Without the outer layer this call blocks in the executor and comes back `ok` with its rows, so + the assertion is on the WHOLE outcome — that it returned at all, in time, and as a refusal. + + The skew is patched to zero so the bound under test is the 1s budget rather than 11s of wall + clock. The +10 value it normally carries is not skipped: it is pinned by the ordering test below, + which needs no clock at all. + """ + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + executor = _BlockingExecutor() + + started = time.monotonic() + try: + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=executor, no_safety=True, + ) + elapsed = time.monotonic() - started + finally: + executor.release.set() + + assert executor.entered.is_set() # the executor really ran; this is not a gate refusing early + assert elapsed < 10, f"the call waited {elapsed:.1f}s on a {_BUDGET_S}s bound" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + # A bound WE imposed is a refusal, not a failure — and it carries no data, like every refusal. + assert env.data is None and env.failure is None + + +def test_the_outer_refusal_says_what_actually_happened(warehouse, monkeypatch): + """It must not claim the statement was cancelled. Nothing was: the worker is still inside the + driver call, may still hold a connection, and the statement may still be running on the server. + The sentence the caller reads names what we observed — the executor did not come back — and + quotes the bound that actually elapsed rather than the inner budget.""" + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + executor = _BlockingExecutor() + + try: + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=executor, no_safety=True, + ) + finally: + executor.release.set() + + assert "cancelled" not in env.refusal.detail, env.refusal.detail + assert "did not return" in env.refusal.detail, env.refusal.detail + # The remediation is the same one the inner bound gives: the fix that makes THIS statement + # runnable, addressed to a caller with no shell and no deployment. + authored = f"{env.refusal.detail} {env.refusal.remediation}" + assert "AGAMI_" not in authored, authored + + +def _guarded_with(executor) -> object: + return execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=executor, no_safety=True, + ) + + +def _drain(deadline_s: float = 5.0) -> None: + """Wait for the abandoned workers to return and give their slots back.""" + deadline = time.monotonic() + deadline_s + while execute_sql._abandoned_workers and time.monotonic() < deadline: + time.sleep(0.01) + + +def test_the_abandoned_workers_are_capped_and_the_cap_releases(warehouse, monkeypatch): + """The cost of the outer bound is an abandoned worker, and it has to have a ceiling. + + Before this layer existed, a call this slow blocked the host's own worker thread — so the thread + limiter capped how much abandoned work could exist and pushed back on everything behind it. + Returning at the bound frees that slot, and on the injected path (no inner watchdog, so EVERY + slow statement abandons one) the abandonments would accumulate at the rate callers arrive, each + still holding a pooled connection and a running statement, until the datasource belonged to work + nobody is waiting for. + + So at the cap the chokepoint refuses BEFORE starting one more — fast, fail-closed, and saying the + executor is saturated rather than blaming a statement that never ran. And the slot comes back + when the abandoned worker finally returns, which is the moment the leak actually ends. + """ + monkeypatch.setattr(execute_sql, "_OUTER_BOUND_SKEW_S", 0) + monkeypatch.setattr(execute_sql, "_MAX_ABANDONED_WORKERS", 2) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + executor = _BlockingExecutor() + + try: + for n in range(2): + env = _guarded_with(executor) + assert env.status == "refused", (n, env) + assert "did not return" in env.refusal.detail, (n, env.refusal.detail) + assert execute_sql._abandoned_workers == 2 + + started = time.monotonic() + env = _guarded_with(executor) + elapsed = time.monotonic() - started + + assert elapsed < _BUDGET_S, f"the capped call still waited out a bound ({elapsed:.2f}s)" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert "saturated" in env.refusal.detail, env.refusal.detail + # It must not blame a statement that never ran, and it must not claim a cancel. + assert "cancelled" not in env.refusal.detail, env.refusal.detail + assert executor.calls == 2, "the refused call started leak N+1 anyway" + finally: + executor.release.set() + + _drain() + assert execute_sql._abandoned_workers == 0, "the slots never came back" + assert _guarded_with(executor).status == "ok" # and the executor is usable again + + +def test_the_inner_refusal_is_the_one_the_caller_gets(warehouse, monkeypatch): + """When both layers are armed the INNER one wins, and the outer neither overwrites its refusal + nor mints a second. + + That ordering is what buys the caller a precise answer: only the watchdog holds the connection + it cancelled, so only it can say the STATEMENT was stopped. A real cancel is driven here (no + stub, no patched skew) with the outer bound left at its real +10, so the outer layer is genuinely + armed and genuinely loses the race. The detail is the assertion, because it is the one part of + the envelope the two layers word differently. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + started = time.monotonic() + env = _guarded(_RUNAWAY_SQL) + elapsed = time.monotonic() - started + + assert elapsed < 20, f"the statement ran {elapsed:.1f}s against a {_BUDGET_S}s budget" + assert env.status == "refused" + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert f"{_BUDGET_S}s" in env.refusal.detail, env.refusal.detail # the inner budget + assert "cancelled" in env.refusal.detail, env.refusal.detail # the inner sentence + assert "did not return" not in env.refusal.detail, env.refusal.detail # not the outer one + + +def test_the_four_bounds_are_ordered_from_one_resolved_budget(monkeypatch): + """One fact, asserted once: the four time bounds are derived from a single resolved budget and + stand in a fixed order — watchdog < native (+5) < outer (+10) < supervisor (+60). + + Asserted together rather than one skew per test because the ORDER is the property, and four + separate assertions would each stay green while the family drifted apart. Each layer behind the + watchdog is a backstop for something the layer in front cannot see, and every inversion trades a + precise refusal for a vaguer answer to the same event: a native kill reads as a database error, + an outer expiry cannot say the statement stopped, and a supervisor kill cannot say what hung. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "45") # distinctive: a hard-coded 30 cannot pass + timeout_s = execute_sql._resolve_timeout_s() + + watchdog = timeout_s + native = timeout_s + execute_sql._NATIVE_BOUND_SKEW_S + outer = timeout_s + execute_sql._OUTER_BOUND_SKEW_S + supervisor = timeout_s + execute_sql._SUPERVISOR_SKEW_S + + assert watchdog < native < outer < supervisor + assert (watchdog, native, outer, supervisor) == (45, 50, 55, 105) + + +class _RaisingExecutor: + """A `ports.Executor` that raises a given exception from inside the worker thread.""" + + def __init__(self, exc: BaseException) -> None: + self._exc = exc + + def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> execute_sql.ExecResult: + raise self._exc + + +def test_an_executor_error_keeps_its_type_across_the_worker(warehouse): + """The subtle half of running the call off-thread: an exception raised in the worker has to reach + the caller with its ORIGINAL type, or every handler in `execute_guarded` stops matching. + + `ExecutorError` is the one that would fail loudest — its handler is what turns a driver error + into a classified `failed` envelope carrying the message this module authored. Caught as anything + else it would land in the catch-all and the caller would get the generic string instead, so the + relayed message is asserted rather than just the status. + """ + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_RaisingExecutor(execute_sql.ExecutorError("no such column: nope", code=5)), + no_safety=True, + ) + + assert env.status == "failed" + assert env.failure.kind == "syntax" # the classified branch, i.e. `except ExecutorError` matched + assert env.failure.message == "no such column: nope" + assert env.failure.message != execute_sql.UNEXPECTED_FAILURE_MESSAGE + + +def test_a_plain_exception_still_reaches_the_catch_all_across_the_worker(warehouse): + """And the other direction: an exception nobody classified is still an unanticipated break, told + to the caller as the one fixed string. A worker that swallowed it would leave the call looking + like it produced nothing at all.""" + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, + executor=_RaisingExecutor(RuntimeError("a driver detail nobody has vetted")), + no_safety=True, + ) + + assert env.status == "failed" + assert env.failure.kind == "other" + assert env.failure.message == execute_sql.UNEXPECTED_FAILURE_MESSAGE + assert "driver detail" not in env.failure.message # the raw text goes to the log, not the caller + + +class _ContextProbeExecutor: + """Records the thread it ran on, the request-scoped row cap it read there, and the budget.""" + + def __init__(self) -> None: + self.thread: threading.Thread | None = None + self.seen_timeout: object = "not read" + self.seen_rows: object = "not read" + + def execute(self, vetted_sql: str, creds: dict, *, profile: str) -> execute_sql.ExecResult: + self.thread = threading.current_thread() + self.seen_timeout = execute_sql._resolve_timeout_s() + self.seen_rows = execute_sql._max_rows_override.get() + return execute_sql.ExecResult(columns=["c"], rows=[(1,)], truncated=False) + + +def test_the_request_scoped_row_cap_is_visible_inside_the_worker(warehouse, monkeypatch): + """A new thread starts with an EMPTY context, so without an explicit copy the caller's per-call + row cap would read as unset inside the worker and every in-process call would quietly fall back + to the deployment default — a bug with no symptom other than the wrong numbers. + + The budget is asserted alongside it for the opposite reason: it rides the environment, which a + new thread inherits for free, and the two must agree inside the worker as well as outside it. + + The thread identity is asserted first, because it is what makes the rest of this test mean + anything: if the call ever stopped running off-thread, the ContextVar assertion would pass + trivially and the copy could be deleted with the suite still green. + """ + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "7") + execute_sql._max_rows_override.set(5) + probe = _ContextProbeExecutor() + + env = execute_sql.execute_guarded( + "SELECT id FROM orders", PROFILE, None, executor=probe, no_safety=True, + ) + + assert env.status == "ok" + assert probe.thread is not threading.current_thread() # it really ran on a worker + assert probe.seen_timeout == 7 + assert probe.seen_rows == 5 + + +# -------------------------------------------------------------------------------------------- +# The supervisor's bound, derived rather than fixed +# -------------------------------------------------------------------------------------------- + + +def test_the_supervisor_bound_is_derived_from_the_statement_budget(warehouse, monkeypatch): + """The fork path's supervisor was a hardcoded 240s, which quietly became the FIRST bound to fire + on any deployment configuring a larger statement budget — turning a refusal that could name the + statement into a `failed`/`timeout` that names nothing. It is now `budget + 60`, from the same + resolver every inner layer reads. + + Asserted at a RAISED budget, which is the case the fixed number got wrong: at 300s the supervisor + must be 360, and 240 would have killed the child a full minute before its own statement bound. + The computed value is read off the `subprocess.run` call rather than waited out, for the obvious + reason. + """ + import tools + + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "300") + monkeypatch.setattr(tools, "resolve_profile", lambda ds: PROFILE) + + class _Proc: + returncode = 0 + stdout = "id\r\n1\r\n" + stderr = "" + + captured: dict = {} + + def _fake_run(cmd, **kwargs): + captured.update(kwargs) + return _Proc() + + monkeypatch.setattr(tools.subprocess, "run", _fake_run) + tools.set_injected_executor(None) # the fork path, which is the one with a supervisor + + tools.tool_execute_sql({"sql": "SELECT id FROM orders", "datasource": PROFILE}) + + assert captured["timeout"] == 300 + execute_sql._SUPERVISOR_SKEW_S == 360 + assert captured["timeout"] > 300, "the supervisor must fire AFTER the statement bound, not before" + assert captured["timeout"] != 240, "the fixed bound this replaces" + # And nothing was passed that would stop the child inheriting the same configuration. + assert "env" not in captured + + +def test_the_supervisors_verdict_is_unchanged(warehouse, monkeypatch): + """Only the DURATION moved. A child that never came back is still a `failed`/`timeout` and never + a `resource_limit`: the bound is ours, but all it tells us is that the child stopped responding — + it may have hung in connect, in credential resolution or in loading the model, and on any of + those "narrow the query" points at the wrong thing.""" + import tools + + monkeypatch.setattr(tools, "resolve_profile", lambda ds: PROFILE) + + def _timed_out(cmd, **kwargs): + raise subprocess.TimeoutExpired(cmd, kwargs.get("timeout")) + + monkeypatch.setattr(tools.subprocess, "run", _timed_out) + tools.set_injected_executor(None) + + body = json.loads(tools.tool_execute_sql({"sql": "SELECT id FROM orders", + "datasource": PROFILE})) + + assert body["status"] == "failed" + assert body["failure"]["kind"] == "timeout" + + +# -------------------------------------------------------------------------------------------- +# Both surfaces refuse identically +# -------------------------------------------------------------------------------------------- + + +def _stdio_refusal(sql: str) -> dict: + """`python -m mcp_harness` over JSON-RPC on stdin — the transport a desktop client launches, and + the one that forks `python -m execute_sql`. The child inherits this process's environment, so it + resolves the same budget, model and warehouse; nothing on this path is stubbed.""" + messages = [ + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}, + {"jsonrpc": "2.0", "method": "notifications/initialized"}, + {"jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "execute_sql", "arguments": {"sql": sql, "datasource": PROFILE, + "raw_query": "how many"}}}, + ] + proc = subprocess.run( + [sys.executable, "-m", "mcp_harness"], + input="".join(json.dumps(m) + "\n" for m in messages), + capture_output=True, text=True, timeout=180, env={**os.environ}, + ) + replies = { + m.get("id"): m + for m in (json.loads(line) for line in proc.stdout.splitlines() if line.strip()) + } + assert 2 in replies, proc.stderr + return json.loads(replies[2]["result"]["content"][0]["text"]) + + +def _http_refusal(sql: str) -> dict: + """The authenticated HTTP transport, which runs execution IN-PROCESS through `create_app()`'s + default adapters — the other of the two execution paths.""" + import mcp_http + import tools + from oauth_server import issue_jwt + from starlette.testclient import TestClient + + headers = { + "Authorization": f"Bearer {issue_jwt('jordan@example.com')}", + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + } + with TestClient(mcp_http.create_app()) as client: + assert tools._INJECTED_EXECUTOR is not None, ( + "create_app() no longer injects an executor, so this surface is now the fork path and " + "the in-process path this test believes it covers is uncovered" + ) + init = client.post("/mcp", headers=headers, json={ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": {"protocolVersion": "2025-06-18", "capabilities": {}, + "clientInfo": {"name": "t", "version": "1"}}}) + session = init.headers.get("mcp-session-id") + headers2 = {**headers, **({"mcp-session-id": session} if session else {})} + client.post("/mcp", headers=headers2, + json={"jsonrpc": "2.0", "method": "notifications/initialized"}) + resp = client.post("/mcp", headers=headers2, json={ + "jsonrpc": "2.0", "id": 2, "method": "tools/call", + "params": {"name": "execute_sql", "arguments": {"sql": sql, "datasource": PROFILE, + "raw_query": "how many"}}}) + assert resp.status_code == 200, resp.text + return json.loads(resp.json()["result"]["content"][0]["text"]) + + +def _audit_rows(app_db: str) -> list: + from store import Store + + store = Store.connect(app_db) + try: + return store.query("SELECT id, status, reason, rule FROM query_executions") + finally: + store.close() + + +def test_both_surfaces_refuse_a_runaway_statement_identically(audited, monkeypatch): + """The criterion in one test: the limit applies the same on stdio and on HTTP, no executor + escapes it, and BOTH leave an audit row. + + The two transports run DIFFERENT execution paths — stdio forks a child and reads its stderr back, + HTTP runs the built-in executor in-process behind the outer bound — so "identically" is asserted + on the whole refusal, not merely on the status. A difference in any field would mean the caller's + answer depends on which door it came through. + + The audit half is asserted here rather than only at the in-process tool edge for the same reason: + the row is written by the serializer, and the fork path reaches it having reconstructed the + refusal from a child's stderr. "A timeout refusal is recorded on both surfaces" is a claim about + both, and one surface's row does not evidence the other's. + """ + pytest.importorskip("starlette") + pytest.importorskip("mcp") + pytest.importorskip("jwt") + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + stdio = _stdio_refusal(_RUNAWAY_SQL) + after_stdio = _audit_rows(audited.app_db) + http = _http_refusal(_RUNAWAY_SQL) + after_http = _audit_rows(audited.app_db) + + assert stdio["status"] == http["status"] == "refused", (stdio, http) + assert stdio["refusal"]["rule"] == guardrail.RULE_RESOURCE_LIMIT + assert stdio["refusal"] == http["refusal"] + + # One row per call, keyed by the id that call handed back — counted after each transport so a + # single row could not stand in for both. + assert len(after_stdio) == 1, after_stdio + assert len(after_http) == 2, after_http + assert {r["id"] for r in after_http} == {stdio["audit_id"], http["audit_id"]} + assert {r["status"] for r in after_http} == {"refused"} + assert {r["rule"] for r in after_http} == {guardrail.RULE_RESOURCE_LIMIT} + assert {r["reason"] for r in after_http} == { + guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + } + + +# -------------------------------------------------------------------------------------------- +# Every engine, under one table +# -------------------------------------------------------------------------------------------- +# +# S2 proved the contract on SQLite. The other nine now run under the same deadline, and each names +# its OWN cancel — because there is no method a duck-typed probe could look for that is right +# everywhere. pymysql's connection has no `cancel()` at all and a probe would fall through to +# `close()`, which sends COM_QUIT down the socket the blocked statement owns; oracledb's connection +# DOES have `cancel()`, so a probe that tried `cancel` first would silently pick it for Oracle and +# never notice that Snowflake and Databricks put theirs on the cursor. The fakes below stand in for +# drivers this environment does not install, so the whole matrix runs on every machine. + + +def _module(name: str, **attrs: object): + """A stand-in driver module. The engine functions do their own `import `, which resolves + through `sys.modules`, so an entry there is enough to reach them.""" + mod = types.ModuleType(name) + for key, value in attrs.items(): + setattr(mod, key, value) + return mod + + +class _RecordingCursor: + """A DB-API cursor that logs every call, so a test can assert WHICH method a cancel reached.""" + + def __init__(self, log: list, *, name: str | None = None, on_execute=None, + close_raises: bool = False): + self.description = [("c",)] + self.name = name + self.itersize = 0 + self.sfqid = "01fake-0000-0000-0000-000000000000" # Snowflake sets a query id at submission + self._log = log + self._on_execute = on_execute + self._close_raises = close_raises + + def execute(self, sql: str, params=None) -> None: + self._log.append(("execute", sql, params)) + # Only the caller's statement is armed to fail. Postgres runs `SET LOCAL statement_timeout` + # on its own cursor first, and that one has to succeed — it runs before anything has gone + # wrong, and a test about the timeout path must not accidentally break the setup for it. + if self._on_execute is not None and not sql.startswith("SET "): + raise self._on_execute() + + def fetchmany(self, n: int): + self._log.append(("fetchmany", n)) + return [(1,)] + + def cancel(self) -> None: + self._log.append(("cursor.cancel",)) + + def abort_query(self, qid: str) -> bool: + self._log.append(("cursor.abort_query", qid)) + return True + + def close(self) -> None: + self._log.append(("cursor.close", self.name)) + if self._close_raises: + raise RuntimeError("current transaction is aborted, commands ignored") + + def __enter__(self): + return self + + def __exit__(self, *exc_info) -> bool: + self.close() + return False + + +class _RecordingConnection: + """Every cancel-shaped method a driver in this module might expose, each logging a DISTINCT + marker — so asserting the log pins down exactly which one the engine chose.""" + + def __init__(self, *, on_execute=None, named_close_raises: bool = False): + self.log: list = [] + self.connect_kwargs: dict = {} + self.job_config_kwargs: dict = {} + self._on_execute = on_execute + self._named_close_raises = named_close_raises + self.cursors: list[_RecordingCursor] = [] + self.executed_cursor = None + # pymssql's DB-API connection delegates its cancel to the `_mssql` connection it wraps. + self._conn = SimpleNamespace(cancel=lambda: self.log.append(("mssql.cancel",))) + + def cursor(self, name: str | None = None, **kwargs): + self.log.append(("cursor", name)) + cur = _RecordingCursor( + self.log, + name=name, + on_execute=self._on_execute, + close_raises=self._named_close_raises and name is not None, + ) + self.cursors.append(cur) + return cur + + # The per-driver cancels, each distinguishable in the log. + def cancel(self) -> None: + self.log.append(("conn.cancel",)) + + def interrupt(self) -> None: + self.log.append(("conn.interrupt",)) + + def _force_close(self) -> None: + self.log.append(("conn._force_close",)) + + def kill(self, thread_id: int) -> None: # pragma: no cover - present so a probe could find it + self.log.append(("conn.kill",)) + + def close(self) -> None: + self.log.append(("conn.close",)) + + # DuckDB's `execute` hands back the CONNECTION, so the connection is also a cursor there. + def execute(self, sql: str, params=None): + self.log.append(("execute", sql, params)) + self.executed_cursor = self + if self._on_execute is not None: + raise self._on_execute() + return self + + @property + def description(self): + return [("c",)] + + def fetchmany(self, n: int): + self.log.append(("fetchmany", n)) + return [(1,)] + + # `with conn` is the TRANSACTION on the Postgres path. + def __enter__(self): + self.log.append(("txn.enter",)) + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + self.log.append(("txn.exit", exc_type is not None)) + return False + + +class _FakeBigQueryResults: + def __init__(self): + self.schema = [SimpleNamespace(name="c")] + + def __iter__(self): + return iter([(1,)]) + + +def _install_simple(module_name: str, connect_attr: str = "connect"): + """Installer for a driver reached as `import ` + `.connect(...)`.""" + def _install(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute, named_close_raises=named_close_raises) + + def _connect(*args, **kwargs): + conn.connect_kwargs = kwargs + return conn + + monkeypatch.setitem(sys.modules, module_name, _module(module_name, **{connect_attr: _connect})) + return conn + + return _install + + +def _install_sqlite(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + monkeypatch.setattr(sqlite3, "connect", lambda *a, **kw: conn) + return conn + + +def _install_snowflake(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + def _connect(**kwargs): + conn.connect_kwargs = kwargs + return conn + + connector = _module("snowflake.connector", connect=_connect) + # `import snowflake.connector` short-circuits on `sys.modules` without the machinery ever setting + # the attribute, so the parent module has to carry it explicitly. + monkeypatch.setitem(sys.modules, "snowflake", _module("snowflake", connector=connector)) + monkeypatch.setitem(sys.modules, "snowflake.connector", connector) + return conn + + +def _install_databricks(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + def _connect(**kwargs): + conn.connect_kwargs = kwargs + return conn + + dbsql = _module("databricks.sql", connect=_connect) + monkeypatch.setitem(sys.modules, "databricks", _module("databricks", sql=dbsql)) + monkeypatch.setitem(sys.modules, "databricks.sql", dbsql) + return conn + + +def _install_trino(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + def _connect(**kwargs): + conn.connect_kwargs = kwargs + return conn + + monkeypatch.setitem( + sys.modules, "trino", _module("trino", dbapi=SimpleNamespace(connect=_connect)), + ) + return conn + + +def _install_bigquery(monkeypatch, *, on_execute=None, named_close_raises=False): + conn = _RecordingConnection(on_execute=on_execute) + + class _Job: + def result(self, max_results=None): + conn.log.append(("job.result", max_results)) + return _FakeBigQueryResults() + + class _Client: + def __init__(self, **kwargs): + conn.connect_kwargs = kwargs + + def query(self, sql, job_config=None): + conn.log.append(("execute", sql, None)) + if on_execute is not None: + raise on_execute() + return _Job() + + def _job_config(**kwargs): + conn.job_config_kwargs = kwargs + return SimpleNamespace(**kwargs) + + bigquery = _module("google.cloud.bigquery", Client=_Client, QueryJobConfig=_job_config) + oauth2 = _module("google.oauth2", service_account=_module("google.oauth2.service_account")) + monkeypatch.setitem(sys.modules, "google", _module("google")) + monkeypatch.setitem(sys.modules, "google.cloud", _module("google.cloud", bigquery=bigquery)) + monkeypatch.setitem(sys.modules, "google.cloud.bigquery", bigquery) + monkeypatch.setitem(sys.modules, "google.oauth2", oauth2) + monkeypatch.setitem(sys.modules, "google.oauth2.service_account", oauth2.service_account) + return conn + + +class _EngineCase: + """One engine's fake driver, its credentials, and the cancel it is required to reach.""" + + def __init__(self, fn, install, creds: dict, cancel: str | None): + self.fn = fn + self.install = install + self.creds = creds + self.cancel = cancel # the log marker the watchdog's cancel must produce; None = no watchdog + + def run(self, sql: str = "SELECT c FROM orders"): + return self.fn(self.creds, sql) + + +_SQL_CREDS = {"host": "db.example", "port": "5432", "user": "u", "password": "p", + "database": "shop"} + +ENGINE_CASES = { + "postgres": _EngineCase( + execute_sql._run_postgres, _install_simple("psycopg2"), dict(_SQL_CREDS), "conn.cancel"), + "mysql": _EngineCase( + execute_sql._run_mysql, _install_simple("pymysql"), dict(_SQL_CREDS), "conn._force_close"), + "snowflake": _EngineCase( + execute_sql._run_snowflake, _install_snowflake, + {"account": "acct", "user": "u", "password": "p"}, "cursor.abort_query"), + "bigquery": _EngineCase( + execute_sql._run_bigquery, _install_bigquery, {"project": "proj"}, None), + "sqlite": _EngineCase( + execute_sql._run_sqlite, _install_sqlite, {"path": "warehouse.db"}, "conn.interrupt"), + "sqlserver": _EngineCase( + execute_sql._run_sqlserver, _install_simple("pymssql"), + {"host": "db.example", "user": "u", "password": "p"}, "mssql.cancel"), + "oracle": _EngineCase( + execute_sql._run_oracle, _install_simple("oracledb"), + {"user": "u", "password": "p", "dsn": "db.example/shop"}, "conn.cancel"), + "databricks": _EngineCase( + execute_sql._run_databricks, _install_databricks, + {"host": "db.example", "http_path": "/sql/1", "token": "t"}, "cursor.cancel"), + "trino": _EngineCase( + execute_sql._run_trino, _install_trino, + {"host": "db.example", "user": "u"}, "cursor.cancel"), + "duckdb": _EngineCase( + execute_sql._run_duckdb, _install_simple("duckdb"), {"path": ":memory:"}, "conn.interrupt"), +} + + +def test_the_engine_table_covers_every_engine_this_module_has(): + """The guard that makes the table below a rule rather than a list. A new engine added without an + `except _ResourceLimit: raise` would otherwise ship silently reporting its refusals as driver + errors — this fails the moment a `_run_*` exists that the matrix does not name.""" + in_module = {n for n in dir(execute_sql) if n.startswith("_run_")} + assert in_module == {case.fn.__name__ for case in ENGINE_CASES.values()} + assert len(ENGINE_CASES) == 10 + + +def _raise_marker(): + return execute_sql._ResourceLimit(execute_sql._OUTLIVED_BUDGET) + + +@pytest.mark.parametrize("engine", sorted(ENGINE_CASES)) +def test_every_engine_re_raises_the_marker_instead_of_relabelling_it(engine, monkeypatch): + """The highest-value assertion in the slice, and the one a copy-paste omission fails. + + Each engine ends in `except Exception as e: raise ExecutorError(..., code=5)`. Without an + `except _ResourceLimit: raise` ahead of it, that catch-all swallows our own marker and the + chokepoint reports a bound WE imposed as the database's failure — a `failed` envelope telling the + caller their SQL is broken, when in fact it merely ran long. + """ + case = ENGINE_CASES[engine] + case.install(monkeypatch, on_execute=_raise_marker) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + +@pytest.mark.parametrize("engine", sorted(ENGINE_CASES)) +def test_no_engine_mistakes_an_ordinary_driver_error_for_the_marker(engine, monkeypatch): + """The other direction, on the same matrix: with the watchdog never fired, a driver error is the + database's outcome and has to stay one. An engine that raised the marker here would tell every + caller with a typo to narrow their query.""" + case = ENGINE_CASES[engine] + case.install(monkeypatch, on_execute=lambda: RuntimeError("relation does not exist")) + + with pytest.raises(execute_sql.ExecutorError) as exc: + case.run() + assert exc.value.code == 5 + + +@contextlib.contextmanager +def _deadline_already_fired(cancel, timeout_s): + """A watchdog that has already fired by the time the block runs — the state a real one reaches + only after a genuinely slow statement, reproduced here without paying for one.""" + fired = threading.Event() + fired.set() + yield fired + + +_WATCHDOG_ENGINES = sorted(e for e, c in ENGINE_CASES.items() if c.cancel is not None) + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_a_driver_error_under_a_fired_watchdog_becomes_the_marker(engine, monkeypatch): + """The conversion itself, on all nine engines that have a watchdog: the flag is set, so whatever + the driver raised is the wreckage of OUR cancel and unwinds as the marker.""" + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES[engine] + case.install(monkeypatch, on_execute=lambda: RuntimeError("connection reset by peer")) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_a_cancel_that_lands_without_raising_is_still_a_refusal(engine, monkeypatch): + """A cancel can land between the execute and the fetch, or just as the fetch returns, and leave + nothing to raise. The budget still elapsed, so the outcome is still a refusal rather than a + result gathered past it — which is why every engine re-reads the flag after the block.""" + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES[engine] + case.install(monkeypatch) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + +def test_the_cli_adapter_renders_the_marker_as_a_refusal_not_a_traceback(monkeypatch, capsys): + """The per-engine CSV wrappers are the SECOND entry into the engine functions, and the engines + raise the internal marker. + + That adapter caught only `ExecutorError`, so a per-statement timeout reached through it escaped + as a traceback rather than as the exit code its contract documents. It renders what `main` + renders for any refusal — the one JSON object on stderr and exit 1 — so the two entries cannot + disagree about what a bound we imposed looks like. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES["sqlite"] + case.install(monkeypatch, on_execute=lambda: RuntimeError("interrupted")) + + code = execute_sql._execute_sqlite(case.creds, "SELECT c FROM orders") + + assert code == 1 + body = json.loads(capsys.readouterr().err) # one JSON object, on one line, and nothing else + assert body["refusal"]["rule"] == guardrail.RULE_RESOURCE_LIMIT + assert body["refusal"]["reason"] == guardrail.REASON_FOR_RULE[guardrail.RULE_RESOURCE_LIMIT] + + +# -------------------------------------------------------------------------------------------- +# The named cancel is the one that actually runs +# -------------------------------------------------------------------------------------------- + + +class _CancelRecorder: + """Stands in for `_deadline` and keeps the callable it was handed, so a test can invoke exactly + what the watchdog would have invoked and see where it lands.""" + + def __init__(self, log: list | None = None): + self.cancels: list = [] + self.budgets: list = [] + # When handed the driver's own log, the arming and disarming are recorded IN it, so a test + # can assert where the deadline sits relative to the calls it is supposed to bound. + self._log = log + + @contextlib.contextmanager + def __call__(self, cancel, timeout_s): + self.cancels.append(cancel) + self.budgets.append(timeout_s) + if self._log is not None: + self._log.append(("deadline.arm",)) + try: + yield threading.Event() + finally: + if self._log is not None: + self._log.append(("deadline.disarm",)) + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_each_engine_arms_its_own_named_cancel(engine, monkeypatch): + """Not "a cancel ran" but "THE cancel ran". The fake connection exposes every cancel-shaped + method any of these drivers has — `cancel`, `interrupt`, `_force_close`, `kill`, `close`, a + cursor `cancel` and a cursor `abort_query` — each logging a distinct marker, so an engine that + reached for the wrong one lands on the wrong marker and fails here.""" + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert len(recorder.cancels) == 1, "exactly one deadline per call, resolved once" + before = len(conn.log) + recorder.cancels[0]() + landed = [entry[0] for entry in conn.log[before:]] + assert landed == [case.cancel], f"{engine} cancelled via {landed}, expected {[case.cancel]}" + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_the_deadline_covers_the_fetch_and_not_only_the_execute(engine, monkeypatch): + """An explicit criterion, and until now pinned on two engines out of nine. + + `_collect_cursor` pulls `cap + 1` rows in a single `fetchmany`, and on a cursor that streams its + result THAT is where the scan happens — a clock stopping at `execute` would bound the cheap half + of the work and leave the expensive half unbounded. Only SQLite (through a fake slow-fetch + cursor) and DuckDB (through this ordering assertion) actually held the line; on the other seven, + moving the fetch out of the deadline's block left the whole suite green. The ordering IS the + criterion, so it is asserted the same way on every engine that arms a watchdog. + """ + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch) + recorder = _CancelRecorder(conn.log) + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + kinds = [entry[0] for entry in conn.log] + # By value, not by first occurrence: Postgres runs `SET LOCAL statement_timeout` on its own + # cursor before the deadline is armed, and that is an `execute` too. + statement = next( + i for i, e in enumerate(conn.log) if e[0] == "execute" and e[1] == "SELECT c FROM orders" + ) + assert ( + kinds.index("deadline.arm") < statement < kinds.index("fetchmany") + < kinds.index("deadline.disarm") + ), conn.log + + +def test_the_mysql_cancel_forces_the_socket_shut_and_never_sends_quit(monkeypatch): + """Called out on its own because it is the case a duck-typed probe gets wrong and no test would + notice. `pymysql.Connection` has no `cancel()`, so a probe falls through to `close()` — which + writes COM_QUIT to the very socket the blocked statement owns and can block on it. Only + `_force_close()` shuts the socket outright, which is what unblocks the statement.""" + case = ENGINE_CASES["mysql"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + before = len(conn.log) + recorder.cancels[0]() + + during_cancel = [entry[0] for entry in conn.log[before:]] + assert during_cancel == ["conn._force_close"] + assert "conn.close" not in during_cancel + assert "conn.kill" not in during_cancel + + +def test_the_postgres_cancel_is_the_connections_own(monkeypatch): + """psycopg2's `connection.cancel()` opens a second connection and sends the libpq cancel request, + so it is safe to call while this thread is blocked inside the driver.""" + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + before = len(conn.log) + recorder.cancels[0]() + + assert [entry[0] for entry in conn.log[before:]] == ["conn.cancel"] + + +def test_the_snowflake_cancel_does_nothing_before_a_query_id_exists(monkeypatch): + """The abort is addressed to a query id, and Snowflake only issues one at submission. A cancel + firing in the sliver before that has nothing to abort, and must not raise on the way to finding + out — the session parameter is what bounds the statement in that window.""" + case = ENGINE_CASES["snowflake"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + conn.cursors[0].sfqid = None # rewind to "submitted nothing yet" + before = len(conn.log) + recorder.cancels[0]() + + assert conn.log[before:] == [] + + +@pytest.mark.parametrize("engine", _WATCHDOG_ENGINES) +def test_each_engine_arms_the_deadline_with_the_resolved_budget(engine, monkeypatch): + """One resolution per call, so the budget the watchdog enforces and the number the refusal quotes + cannot be two different values.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(11)) + case = ENGINE_CASES[engine] + case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert recorder.budgets == [11] + + +def test_duckdb_cancels_through_the_connection_even_though_the_cursor_is_the_connection(monkeypatch): + """DuckDB's `execute` hands back the CONNECTION rather than a cursor, so `cur` and `conn` are the + same object and a cursor-side cancel would be indistinguishable from a connection-side one by + inspection. The cancel that works is `interrupt()`, and it has to be armed around the execute: + on an in-process engine the execute is where the scan happens, not the fetch.""" + case = ENGINE_CASES["duckdb"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder(conn.log) + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert conn.executed_cursor is conn # the premise: execute returned the connection itself + # The deadline is armed BEFORE the execute and covers the fetch as well — on an in-process + # engine the execute is where the scan happens, so arming it after would bound nothing. + kinds = [entry[0] for entry in conn.log] + assert kinds[:1] == ["deadline.arm"] + assert kinds.index("execute") < kinds.index("fetchmany") < kinds.index("deadline.disarm") + before = len(conn.log) + recorder.cancels[0]() + assert [entry[0] for entry in conn.log[before:]] == ["conn.interrupt"] + + +# -------------------------------------------------------------------------------------------- +# Nothing in a `finally` may eat the marker +# -------------------------------------------------------------------------------------------- + + +class _EngineExecutor: + """Runs one engine's real `_run_` against its fake driver, so the assertion is about the + Envelope the chokepoint produces rather than about the exception the engine raises.""" + + def __init__(self, case: "_EngineCase"): + self._case = case + + def execute(self, vetted_sql: str, creds: dict, *, profile: str): + return self._case.fn(self._case.creds, vetted_sql) + + +# The three engines whose `finally` closed the connection unguarded. The other seven already wrapped +# it; these were the ones a refusal could still die in. +@pytest.mark.parametrize("engine", ["mysql", "postgres", "sqlite"]) +def test_a_connection_close_that_raises_does_not_eat_the_refusal(engine, warehouse, monkeypatch): + """The same structural bug as the named cursor above, one line further down. + + An exception raised inside a `finally` REPLACES the one propagating through it, and by the time + the `finally` runs the engine's `except _ResourceLimit: raise` has already re-raised the marker. + So a `close()` that throws destroys it, the catch-all one layer out wraps the replacement, and a + bound WE imposed reaches the caller as an unclassified server break telling them nothing. + + MySQL is the most exposed and the reason this is not theoretical: its cancel is `_force_close()`, + which deliberately destroys the socket — so on exactly the timeout path, `close()` is being asked + to write COM_QUIT to a socket that is already gone. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch, on_execute=lambda: RuntimeError("connection reset by peer")) + + def _close_raises() -> None: + conn.log.append(("conn.close",)) + raise RuntimeError("the socket the cancel destroyed is not there to close") + + monkeypatch.setattr(conn, "close", _close_raises) + + env = execute_sql.execute_guarded( + "SELECT c FROM orders", PROFILE, None, executor=_EngineExecutor(case), no_safety=True, + ) + + assert ("conn.close",) in conn.log, "the connection was never closed" + assert env.status == "refused", getattr(env, "failure", None) + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert env.failure is None + + +class _PostgresExecutor: + """Runs the real `_run_postgres` against the fake driver, so the assertion is about the Envelope + the chokepoint produces rather than about the exception the engine raises.""" + + def __init__(self, creds: dict): + self._creds = creds + + def execute(self, vetted_sql: str, creds: dict, *, profile: str): + return execute_sql._run_postgres(self._creds, vetted_sql) + + +def test_the_postgres_marker_survives_a_named_cursor_close_that_raises(warehouse, monkeypatch): + """The structural bug this slice had to fix before the deadline could work on Postgres at all. + + Closing a server-side cursor sends `CLOSE agami_bounded`. On the timeout path the transaction is + already aborted, so that statement raises in turn — and with the cursor inside a `with`, it + raises from `__exit__`, where a new exception REPLACES the one being propagated. The marker + vanishes, the engine's own catch-all wraps the replacement in an `ExecutorError`, and Postgres + alone reports every timeout as a failure. The fix is to close it by hand and swallow that. + """ + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES["postgres"] + conn = case.install( + monkeypatch, + on_execute=lambda: RuntimeError("canceling statement due to user request"), + named_close_raises=True, + ) + + env = execute_sql.execute_guarded( + "SELECT c FROM orders", PROFILE, None, + executor=_PostgresExecutor(case.creds), no_safety=True, + ) + + assert [e for e in conn.log if e[0] == "cursor.close"], "the named cursor was never closed" + assert env.status == "refused", getattr(env, "failure", None) + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert conn.log[-1] == ("conn.close",) # and the connection is still released + + +def test_the_postgres_transaction_still_rolls_back_when_the_deadline_fires(monkeypatch): + """`with conn` stayed for a reason: it is the TRANSACTION, and leaving it by an exception rolls + back. Taking the cursor out of its own `with` must not take that with it.""" + monkeypatch.setattr(execute_sql, "_deadline", _deadline_already_fired) + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch, on_execute=lambda: RuntimeError("canceling statement")) + + with pytest.raises(execute_sql._ResourceLimit): + case.run() + + assert ("txn.enter",) in conn.log + assert ("txn.exit", True) in conn.log # exited WITH an exception in flight, i.e. rolled back + + +# -------------------------------------------------------------------------------------------- +# The native server-side bounds — three engines, each one skew behind the watchdog +# -------------------------------------------------------------------------------------------- + +_BUDGET_FOR_NATIVE = 11 # distinctive, so a hard-coded 30 or 35 cannot pass + + +def test_the_skew_puts_the_native_bound_behind_the_watchdog(): + """The direction of the skew is the whole design. Ahead of the watchdog, a server-side kill would + win the race, the flag would be clear, and the refusal would arrive as a database failure.""" + assert execute_sql._NATIVE_BOUND_SKEW_S > 0 + + +def test_postgres_sets_the_native_bound_on_the_same_transaction_before_the_named_cursor(monkeypatch): + """`SET LOCAL` is transaction-scoped, so it is worthless unless it runs on the transaction the + statement will run in, and before it. Ordering is the assertion; the value is the other half.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch) + + case.run() + + kinds = [entry[0] for entry in conn.log] + set_local = next( + i for i, e in enumerate(conn.log) + if e[0] == "execute" and e[1].startswith("SET LOCAL statement_timeout") + ) + declared = conn.log.index(("cursor", "agami_bounded")) + statement = next( + i for i, e in enumerate(conn.log) if e[0] == "execute" and e[1] == "SELECT c FROM orders" + ) + + assert kinds.index("txn.enter") < set_local < declared < statement + # Same transaction: nothing committed or rolled back between the setting and the statement. + assert "txn.exit" not in kinds[:statement] + # The setting is in milliseconds, and it sits one skew behind our own budget. + assert conn.log[set_local][2] == ((_BUDGET_FOR_NATIVE + execute_sql._NATIVE_BOUND_SKEW_S) * 1000,) + + +def test_postgres_does_not_set_the_native_bound_through_the_connect_options(monkeypatch): + """The libpq `options` startup parameter is the other way to set `statement_timeout`, and it is + the wrong one here: a transaction-mode connection pooler can reject an unknown startup parameter, + which breaks the connect outright rather than bounding the statement.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) + case = ENGINE_CASES["postgres"] + conn = case.install(monkeypatch) + + case.run() + + assert "options" not in conn.connect_kwargs + + +def test_snowflake_sets_its_statement_timeout_as_a_session_parameter(monkeypatch): + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) + case = ENGINE_CASES["snowflake"] + conn = case.install(monkeypatch) + + case.run() + + session = conn.connect_kwargs["session_parameters"] + assert session["STATEMENT_TIMEOUT_IN_SECONDS"] == ( + _BUDGET_FOR_NATIVE + execute_sql._NATIVE_BOUND_SKEW_S + ) + + +def test_bigquery_sets_a_job_timeout_and_is_the_one_engine_with_no_cancel(monkeypatch): + """BigQuery's native bound is the ONLY bound it has, and that is a recorded residual rather than + an oversight: there is no connection to cancel, and the call that blocks is `job.result()`, which + is reached only after `client.query()` returns — so at the instant a watchdog would fire there is + nothing in hand to stop. A client-side stall here comes back `failed`, not `resource_limit`.""" + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_FOR_NATIVE)) + case = ENGINE_CASES["bigquery"] + conn = case.install(monkeypatch) + recorder = _CancelRecorder() + monkeypatch.setattr(execute_sql, "_deadline", recorder) + + case.run() + + assert conn.job_config_kwargs["job_timeout_ms"] == ( + _BUDGET_FOR_NATIVE + execute_sql._NATIVE_BOUND_SKEW_S + ) * 1000 + assert recorder.cancels == [], "BigQuery arms no watchdog — see the residual above" + + +def test_bigquery_keeps_the_default_dataset_alongside_the_job_timeout(monkeypatch): + """The job config used to be built only when a default dataset was configured. Now it is always + built, and the pre-existing setting has to survive that.""" + case = ENGINE_CASES["bigquery"] + conn = case.install(monkeypatch) + + execute_sql._run_bigquery({"project": "proj", "dataset": "shop"}, "SELECT c FROM orders") + + assert conn.job_config_kwargs["default_dataset"] == "proj.shop" + assert "job_timeout_ms" in conn.job_config_kwargs + + +@pytest.mark.parametrize("engine", sorted(set(ENGINE_CASES) - {"postgres", "snowflake", "bigquery"})) +def test_no_other_engine_grows_a_native_bound(engine, monkeypatch): + """Exactly three engines get one. The rest are bounded by the watchdog alone, and inventing a + per-engine session setting for them would be a second, unasserted timeout to keep in step.""" + case = ENGINE_CASES[engine] + conn = case.install(monkeypatch) + + case.run() + + statements = " ".join(e[1] for e in conn.log if e[0] == "execute" and isinstance(e[1], str)) + assert "timeout" not in statements.lower() + assert "session_parameters" not in conn.connect_kwargs + + +# -------------------------------------------------------------------------------------------- +# A real in-process bomb, on a real DuckDB +# -------------------------------------------------------------------------------------------- + + +def test_a_cartesian_bomb_is_bounded_on_a_real_duckdb(tmp_path, monkeypatch): + """The fakes above prove the wiring; this proves the cancel. DuckDB is in-process like SQLite, so + a genuine `interrupt()` can be driven end to end with no network and no fixture warehouse — and a + cross join of two ten-million-row ranges is 10^14 rows, far beyond what the budget allows, so a + run that returns in time can only have been stopped.""" + duckdb = pytest.importorskip("duckdb", reason="duckdb is not installed in this environment") + + path = tmp_path / "warehouse.duckdb" + con = duckdb.connect(str(path)) + con.execute("CREATE TABLE orders (id INTEGER)") + con.close() + + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + started = time.monotonic() + with pytest.raises(execute_sql._ResourceLimit): + execute_sql._run_duckdb( + {"path": str(path)}, + "SELECT count(*) AS c FROM range(10000000) a, range(10000000) b", + ) + elapsed = time.monotonic() - started + + assert elapsed < 20, f"the statement ran {elapsed:.1f}s against a {_BUDGET_S}s budget" diff --git a/tests/test_ace044_bounded_fetch.py b/tests/test_ace044_bounded_fetch.py index 2c1d8e76..e4025a8d 100644 --- a/tests/test_ace044_bounded_fetch.py +++ b/tests/test_ace044_bounded_fetch.py @@ -175,13 +175,19 @@ def __enter__(self): def __exit__(self, *a): return False - def execute(self, sql): + def execute(self, sql, params=None): + # `params` accepts the native server-side bound the engine now sets first + # (`SET LOCAL statement_timeout = %s`, on its own client-side cursor). The caller's + # statement runs last, so `seen["sql"]` still ends up holding it. seen["sql"] = sql def fetchmany(self, n): seen["fetchmany"] = n return self._rows[:n] + def close(self): + seen["closed"] = True + class FakeConn: def cursor(self, name=None): return FakeCur(name) @@ -192,6 +198,10 @@ def __enter__(self): def __exit__(self, *a): return False + def cancel(self): + # The per-statement watchdog arms `connection.cancel` on this engine; never fired here. + seen["cancelled"] = True + def close(self): pass diff --git a/tests/test_postgres_timeout_integration.py b/tests/test_postgres_timeout_integration.py new file mode 100644 index 00000000..6fe028d9 --- /dev/null +++ b/tests/test_postgres_timeout_integration.py @@ -0,0 +1,294 @@ +"""Integration guard: the per-statement timeout on a LIVE Postgres, asserted against the SERVER's +view of what happened rather than the client's. + +The in-process tests prove the refusal; they cannot prove the cancel. A client that simply stopped +waiting — abandoned the statement and closed nothing — produces the identical Envelope while the +backend keeps scanning for nobody. So the assertion here is made from a SECOND connection: the +statement is observed running in `pg_stat_activity`, and after the refusal it is GONE from +`pg_stat_activity`. That is the difference between a bound on the caller and a bound on the +database, and it is the whole security value of the feature. + +The second test proves the other half — that `SET LOCAL statement_timeout` really landed on the +session. It disarms our own watchdog and lets the statement run past the native bound; Postgres then +kills it itself, and says so in its own words ("canceling statement due to statement timeout", which +a client cancel never produces — that one reads "due to user request"). + +Opt-in: it **skips unless `AGAMI_IT_PG_PASSWORD` is set** (so normal CI, which has no Postgres, is +unaffected, and no test password lives in the source). To run it against the integration fixture: + + docker compose -f tests/integration/docker-compose.yml up -d postgres + AGAMI_IT_PG_PASSWORD= \ + uv run pytest tests/test_postgres_timeout_integration.py + +Host/port/user/db default to the fixture's values and can be overridden via the other AGAMI_IT_PG_* +vars, exactly as in `test_postgres_named_cursor_integration.py`. +""" + +from __future__ import annotations + +import contextlib +import os +import sys +import threading +import time +import urllib.parse +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +PKG_SRC = REPO_ROOT / "packages" / "agami-core" / "src" +if str(PKG_SRC) not in sys.path: + sys.path.insert(0, str(PKG_SRC)) + +import execute_sql # noqa: E402 +import guardrail # noqa: E402 + +PROFILE = "analytics" + +# Long enough that the observer below reliably catches the statement mid-flight on a loaded machine, +# short enough that the whole file stays a few seconds. The resolver deals in whole seconds, so this +# is also the smallest useful step above the 1s floor the SQLite tests use. +_BUDGET_S = 3 + +# A statement with no termination in reach on any machine: a three-way cross join of ten thousand +# rows each, so a trillion nested-loop iterations of pure CPU. Sized that way on purpose — "bounded +# rather than hanging" is only proved if the unbounded run would take far longer than the assertion +# allows. +# +# The SHAPE matters as much as the size. One wide `generate_series(1, 100000000000)` looks like the +# simpler runaway and is wrong: Postgres materializes a function scan into a tuplestore, so that +# statement burns time by filling the server's temp tablespace, and a test that leaves the disk full +# has done more damage than the bug it was checking for. Three narrow series each materialize a few +# hundred kilobytes and never grow — the cost is entirely time, which is the only axis under test. +# +# `pg_sleep` would be the obvious runaway and is deliberately not used: the read-only guard blocks it +# by name (it is a DoS function), so it would never reach the executor. This does reach it — it opens +# with SELECT, names no physical table, projects no star, and calls nothing the guard denies. +_RUNAWAY_SQL = ( + "SELECT count(*) AS n FROM generate_series(1, 10000) AS a(x), " + "generate_series(1, 10000) AS b(y), generate_series(1, 10000) AS c(z)" +) + +# What the executor's work looks like from the server side. The Postgres path runs through a +# server-side cursor, so the backend's `query` is the `DECLARE "agami_bounded" …` and then the +# `FETCH FORWARD … FROM "agami_bounded"` that actually burns the CPU — the cursor name, not the +# SELECT text, is what is on the wire for most of the run. +_CURSOR_NAME = "agami_bounded" + + +def _pg_creds() -> dict[str, str]: + # The password is env-only (no source-embedded test secret); host/port/user/db default to the + # fixture's values. Override any of them via AGAMI_IT_PG_* to point at another Postgres. + return { + "type": "postgres", + "host": os.environ.get("AGAMI_IT_PG_HOST", "127.0.0.1"), + "port": os.environ.get("AGAMI_IT_PG_PORT", "55432"), + "user": os.environ.get("AGAMI_IT_PG_USER", "agami_test"), + "password": os.environ["AGAMI_IT_PG_PASSWORD"], + "database": os.environ.get("AGAMI_IT_PG_DB", "shop"), + } + + +@pytest.fixture +def pg_observer(): + """A live Postgres, plus the SECOND connection the assertions are made from. + + `autocommit` is not a detail: `pg_stat_activity` is served from a per-transaction snapshot of the + backend-status array, so polling it inside one long transaction returns the same stale rows + forever. Each poll has to be its own transaction to see the server's current state. + """ + psycopg2 = pytest.importorskip("psycopg2") + if not os.environ.get("AGAMI_IT_PG_PASSWORD"): + pytest.skip("set AGAMI_IT_PG_PASSWORD to run the live-Postgres integration test") + creds = _pg_creds() + try: + conn = psycopg2.connect( + host=creds["host"], port=int(creds["port"]), user=creds["user"], + password=creds["password"], dbname=creds["database"], connect_timeout=3, + ) + except Exception as exc: # no DB in this environment → skip, don't fail + pytest.skip(f"no reachable Postgres for the integration test ({exc})") + conn.autocommit = True + try: + yield conn, creds + finally: + conn.close() + + +@pytest.fixture(autouse=True) +def _reset_overrides(): + # The row cap is a request-scoped ContextVar; isolate this file from whatever set it. The budget + # has no such surface at all — it travels by env var only, which is what lets a forked child and + # a thread this test starts both resolve the same number. + execute_sql._max_rows_override.set(None) + yield + execute_sql._max_rows_override.set(None) + + +def _dsn(creds: dict[str, str]) -> str: + # `safe=""` because every reserved character has to be escaped inside userinfo and the path. + # `quote`'s default leaves `/` alone, which would end the userinfo or the database name early and + # build a DSN that silently points somewhere else. + def quote(value: str) -> str: + return urllib.parse.quote(value, safe="") + + return ( + f"postgresql://{quote(creds['user'])}:{quote(creds['password'])}" + f"@{creds['host']}:{creds['port']}/{quote(creds['database'])}" + ) + + +def test_the_dsn_escapes_every_reserved_character_in_userinfo_and_path(): + """A `/` in a password must not silently truncate the DSN and point the test somewhere else. + + `urllib.parse.quote` leaves `/` alone by default, which is right for a path and wrong for every + component this builds: a `/` in the password ends the userinfo early and one in the database name + ends the path early, either way producing a URL that parses cleanly and connects elsewhere. This + one needs no database, so it runs in ordinary CI alongside the opt-in tests below. + """ + # The values spell out the characters under test rather than looking like credentials, so a + # secret scanner has nothing to flag and a reader can see what each one is for. + reserved = "slash/colon:at@" + dsn = _dsn( + {"user": "a/b", "password": reserved, "host": "127.0.0.1", "port": "5432", "database": "d/b"} + ) + parsed = urllib.parse.urlparse(dsn) + + assert urllib.parse.unquote(parsed.username or "") == "a/b" + assert urllib.parse.unquote(parsed.password or "") == reserved + assert urllib.parse.unquote(parsed.path.lstrip("/")) == "d/b" + assert parsed.hostname == "127.0.0.1" + assert parsed.port == 5432 + + +def _our_backends(conn, creds: dict[str, str]) -> list[tuple]: + """Every backend OTHER than this one whose work is the executor's — the server's own answer to + "is that statement still there?". `pg_backend_pid()` excludes the observer, whose own query text + contains the cursor name it is searching for.""" + with conn.cursor() as cur: + cur.execute( + "SELECT pid, state, left(query, 120) FROM pg_stat_activity " + "WHERE datname = %s AND pid <> pg_backend_pid() AND query LIKE %s", + (creds["database"], f"%{_CURSOR_NAME}%"), + ) + return cur.fetchall() + + +def _wait_until(predicate, limit_s: float, poll_s: float = 0.02) -> bool: + deadline = time.monotonic() + limit_s + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(poll_s) + return predicate() + + +def _guarded_on_a_thread(sql: str) -> tuple[threading.Thread, dict]: + outcome: dict = {} + + def call() -> None: + try: + outcome["env"] = execute_sql.execute_guarded( + sql, PROFILE, None, executor=execute_sql.BUILTIN_EXECUTOR, no_safety=True + ) + except BaseException as exc: # the chokepoint is total; record a breach rather than hide it + outcome["raised"] = exc + + worker = threading.Thread(target=call, name="agami-it-guarded", daemon=True) + worker.start() + return worker, outcome + + +def test_the_cancel_reaches_the_server_and_the_statement_is_gone(pg_observer, monkeypatch): + """The headline, and the one criterion the client cannot self-certify. + + Three facts in order, from a connection the executor does not own: the statement WAS running on + the server, the caller got the `resource_limit` refusal, and the statement is no longer there. + The last one is polled rather than asserted instantaneously — a cancel is a request the backend + acts on at its next check, not a synchronous return — but the window is two seconds, which a + statement counting to a hundred billion would still be inside by many minutes. + + The detail assertion is not decoration. `execute_guarded` has TWO time bounds, and the outer one + (`_execute_bounded`) is exactly the client-only abandon this test exists to rule out: it stops + waiting, leaves the worker inside the driver, and cannot say the statement stopped. Its refusal + says so in different words. Asserting the INNER wording is what proves the cancel is the thing + that ended this, rather than the caller walking away. + """ + conn, creds = pg_observer + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", _dsn(creds)) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(_BUDGET_S)) + + started = time.monotonic() + worker, outcome = _guarded_on_a_thread(_RUNAWAY_SQL) + + # 1. The server really is running it — otherwise "it is gone" proves nothing. + assert _wait_until(lambda: bool(_our_backends(conn, creds)), limit_s=_BUDGET_S), ( + "the statement never appeared in pg_stat_activity; the test would prove nothing" + ) + + # 2. The caller gets the refusal, well inside the outer bound (budget + 10s). + worker.join(timeout=_BUDGET_S + 30) + assert not worker.is_alive(), "the executor never returned" + assert "raised" not in outcome, outcome.get("raised") + elapsed = time.monotonic() - started + env = outcome["env"] + assert env.status == "refused", env + assert env.refusal.rule == guardrail.RULE_RESOURCE_LIMIT + assert env.data is None + assert f"{_BUDGET_S}s limit" in env.refusal.detail, env.refusal.detail + assert "cancelled" in env.refusal.detail, env.refusal.detail # the inner bound, not the abandon + assert elapsed < _BUDGET_S + execute_sql._OUTER_BOUND_SKEW_S, f"{elapsed:.1f}s" + + # 3. The server's view: the statement is gone, not merely disowned. A backend left `idle in + # transaction` would still be listed here and would fail this assertion, so this covers the + # half-abandon (client stopped reading, session still holding the transaction open) as well + # as the outright runaway. + assert _wait_until(lambda: not _our_backends(conn, creds), limit_s=2.0), ( + f"still on the server after the refusal: {_our_backends(conn, creds)}" + ) + + +def test_the_native_statement_timeout_is_really_set_on_the_session(pg_observer, monkeypatch): + """`SET LOCAL statement_timeout` is the backstop for the case the watchdog cannot cover — our + process dying mid-statement — so it has to be proven to have LANDED, not just to have been sent. + + It cannot be read back the ordinary way: `SHOW statement_timeout` is not a SELECT and + `current_setting()` is a denied function, so neither reaches the executor, and the setting is + transaction-scoped so it is gone by the time anything else could look. What is left is to let it + fire. With our watchdog disarmed the statement runs on to the native bound, and the proof is + Postgres's own wording: it says "statement timeout", where a client `conn.cancel()` says "user + request". Only the server can produce the former. + + The elapsed floor pins the skew as well as the fact: the native bound is the budget PLUS five + seconds, deliberately behind the watchdog so the watchdog always wins the race in normal + operation and stays the sole classification signal. + """ + conn, creds = pg_observer + budget_s = 1 + monkeypatch.setenv(f"DATASOURCE_URL__{PROFILE.upper()}", _dsn(creds)) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", str(budget_s)) + + @contextlib.contextmanager + def _disarmed(cancel, timeout_s): + # The watchdog, minus the timer: the Event never fires, so nothing client-side stops this. + yield threading.Event() + + monkeypatch.setattr(execute_sql, "_deadline", _disarmed) + + started = time.monotonic() + worker, outcome = _guarded_on_a_thread(_RUNAWAY_SQL) + worker.join(timeout=budget_s + execute_sql._OUTER_BOUND_SKEW_S + 30) + assert not worker.is_alive(), "the executor never returned" + elapsed = time.monotonic() - started + + env = outcome["env"] + # A server-side kill is not our refusal: nothing set the flag, so it is an ordinary database + # failure — which is exactly why the watchdog is set to fire first in normal operation. + assert env.status == "failed", env + assert "statement timeout" in env.failure.message.lower(), env.failure.message + native_s = budget_s + execute_sql._NATIVE_BOUND_SKEW_S + assert elapsed >= native_s - 0.5, f"killed after {elapsed:.1f}s, before the {native_s}s bound" + assert _wait_until(lambda: not _our_backends(conn, creds), limit_s=2.0), ( + f"still on the server after the native timeout: {_our_backends(conn, creds)}" + )