From fbba15fe4afa602bf8588c9d21ec7eb2c8538aa5 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Sun, 12 Jul 2026 15:26:00 -0700 Subject: [PATCH 1/2] ACE-038: per-statement timeout (the availability guarantee) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A generated query can hang or exhaust the DB (recursive CTE, cartesian bomb, unbounded scan). AGAMI_SQL_TIMEOUT_S (default 30) bounds every statement: each engine sets its NATIVE server-side timeout where it has one (Postgres SET LOCAL statement_timeout, MySQL/MariaDB max_execution_time / max_statement_time, Snowflake STATEMENT_TIMEOUT_IN_SECONDS, BigQuery job_timeout_ms, Oracle call_timeout, SQL Server / Trino / Databricks native caps) — the genuine DB-side cancel — layered under a universal wall-clock watchdog (_run_bounded + _deadline) that cancels the in-flight query client-side for any engine without one (SQLite / DuckDB conn.interrupt). A fired deadline kills the query and returns NO result. Rebased onto the reconciled ACE-035 seam: a timeout raises _ResourceLimit up through the engine's transaction handling (so `with conn` rolls the cancelled txn back), and execute_guarded maps it — at the single chokepoint both surfaces funnel through — to the typed resource_limit Refusal. So the subprocess wire and the in-process handler both surface one refused Envelope with no partial data. _run_bounded now returns the bounded ExecResult; the CLI adapter (_emit_or_err) surfaces the same refusal + exit code. Spec: ACE-038 Co-Authored-By: Claude Opus 4.8 (1M context) --- deploy/agami.env.example | 8 + packages/agami-core/src/execute_sql.py | 307 ++++++++++++--- plugins/agami/lib/execute_sql.py | 307 ++++++++++++--- .../agami-deploy/bundle/agami.env.example | 8 + tests/test_ace044_bounded_fetch.py | 3 + tests/test_execute_sql_envelope.py | 22 ++ tests/test_resource_limits.py | 362 ++++++++++++++++++ 7 files changed, 901 insertions(+), 116 deletions(-) create mode 100644 tests/test_resource_limits.py diff --git a/deploy/agami.env.example b/deploy/agami.env.example index 6d7dc71c..deba2c5f 100644 --- a/deploy/agami.env.example +++ b/deploy/agami.env.example @@ -62,3 +62,11 @@ AGAMI_ADMIN_PASSWORD= # LATERAL source. This is the safe default. `warn` logs and allows instead; use it only as a # temporary staged-rollout escape hatch, never as the shipped setting. # AGAMI_SQL_UNSCOPABLE_POSTURE=enforce + +# Optional: SQL resource limits — bound a runaway query (a recursive CTE, a cartesian product, an +# unbounded scan). Both are enforced at the executor, so every SQL surface inherits them. +# AGAMI_SQL_TIMEOUT_S=30 # per-statement timeout in seconds; a query that exceeds it is cancelled +# # (server-side where the engine supports it) and refused. Default 30; +# # keep it below ~240 (the executor's own process wall). +# AGAMI_SQL_MAX_ROWS=1000 # max rows returned per query; a larger result is truncated and flagged, +# # never silently cut. Raise it for bigger pulls. Default 1000. diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index 3211fbe3..d60b99b5 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -45,8 +45,11 @@ import os import stat import sys +import threading +import time import urllib.parse from collections.abc import Callable +from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass from pathlib import Path @@ -424,15 +427,30 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"Postgres connect failed: {e}", code=4) try: with conn: + # Genuine server-side cancel: the backend aborts the statement itself after the deadline. + # Set per-transaction via SET LOCAL — NOT the libpq `options` startup parameter, which a + # transaction-mode pooler (Supabase Supavisor, PgBouncer) can reject at connect. SET LOCAL + # is transaction-scoped, which those poolers pass through cleanly. The watchdog's + # conn.cancel() (a pg_cancel request) is the client-initiated backstop. + with conn.cursor() as setup: + setup.execute(f"SET LOCAL statement_timeout = {_timeout_ms()}") # 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) # 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") # not `with`: a cancelled txn makes CLOSE raise + cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window + try: + result = _run_bounded(lambda: _exec(cur, sql), conn.cancel) + finally: + try: + cur.close() + except Exception: + pass # a cancelled query leaves the txn aborted; closing the server-side cursor + # would raise and mask _ResourceLimit — swallow so the timeout refusal survives + except _ResourceLimit: + raise # timed out; `with conn` rolled back — execute_guarded maps it to a resource_limit refusal except Exception as e: raise ExecutorError(f"Postgres execution error: {e}", code=5) finally: @@ -446,6 +464,7 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: except ImportError: raise ExecutorError("pymysql not installed. Run: pip install pymysql", code=3) _require(creds, "host", "port", "user", "password", "database") + secs = _resolve_timeout_s() try: conn = pymysql.connect( host=creds["host"], @@ -455,18 +474,37 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: database=creds["database"], charset="utf8mb4", connect_timeout=10, + # The reliable client-side bound: a socket read/write blocked past the deadline raises on + # its own. Closing the connection from the watchdog thread does NOT unblock a recv() on + # Linux (that needs shutdown(), which pymysql's close() skips), so this — not the close — + # is what guarantees MySQL/MariaDB can't hang past the deadline. + read_timeout=secs, + write_timeout=secs, autocommit=True, ) except Exception as e: raise ExecutorError(f"MySQL connect failed: {e}", code=4) + # Native server-side timeout. MySQL 5.7.8+ spells it max_execution_time (ms); MariaDB spells it + # max_statement_time (seconds). Try both — each is a no-op on the other dialect (a swallowed + # unknown-variable error) — so whichever the server is gets a genuine server-side abort. + for stmt in ( + f"SET SESSION max_execution_time={secs * 1000}", + f"SET SESSION max_statement_time={secs}", + ): + try: + with conn.cursor() as _c: + _c.execute(stmt) + except Exception: + pass # not this dialect; read_timeout still bounds the client try: - with conn.cursor() as cur: - cur.execute(sql) - result = _collect_cursor(cur) + cur = conn.cursor() # not `with` — the watchdog may close conn, so avoid a close-on-close + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: - conn.close() + _safe_close(conn) return result @@ -492,6 +530,8 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: "user": creds["user"], "client_session_keep_alive": False, "login_timeout": 15, + # Native server-side statement timeout (seconds) — Snowflake aborts the query itself. + "session_parameters": {"STATEMENT_TIMEOUT_IN_SECONDS": _resolve_timeout_s()}, } for k in ("password", "warehouse", "database", "schema", "role", "authenticator"): if creds.get(k): @@ -502,15 +542,13 @@ 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) + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Snowflake execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -573,7 +611,8 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: # 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] = {} + # `job_timeout_ms` is the native server-side job timeout — BigQuery ends the job itself. + job_config_kwargs: dict[str, Any] = {"job_timeout_ms": _timeout_ms()} if creds.get("dataset"): try: job_config_kwargs["default_dataset"] = f"{project}.{creds['dataset']}" @@ -581,16 +620,30 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: pass cap = _resolve_row_cap() + job_box: dict[str, Any] = {} + started = time.monotonic() + timeout_s = _resolve_timeout_s() + + def _bq_cancel() -> None: + job = job_box.get("job") + if job is not None: + job.cancel() # genuine server-side job cancel (the client-initiated backstop to job_timeout_ms) + 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) - # 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 + with _deadline(_bq_cancel) as fired: + try: + job = client.query(sql, job_config=bigquery.QueryJobConfig(**job_config_kwargs)) + job_box["job"] = job + # 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 Exception: + if _deadline_hit(fired, started, timeout_s): + raise _ResourceLimit from None + raise + except _ResourceLimit: + raise # execute_guarded maps a fired BigQuery deadline to a resource_limit refusal except Exception as e: raise ExecutorError(f"BigQuery execution error: {e}", code=5) @@ -619,8 +672,9 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"SQLite connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + result = _run_bounded(lambda: _exec(cur, sql), conn.interrupt) + except _ResourceLimit: + raise # timed out — execute_guarded maps it to a resource_limit refusal except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: @@ -643,20 +697,19 @@ def _run_sqlserver(creds: dict[str, str], sql: str) -> ExecResult: password=creds["password"], database=creds.get("database", ""), login_timeout=15, + timeout=_resolve_timeout_s(), # native per-query timeout (seconds) ) except Exception as e: raise ExecutorError(f"SQL Server connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"SQL Server execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -677,17 +730,19 @@ 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) + try: + conn.call_timeout = _timeout_ms() # native round-trip timeout (ms) + except Exception: + pass # older driver / mode; the wall-clock watchdog still bounds the query try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Oracle execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -709,17 +764,21 @@ def _run_databricks(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Databricks connect failed: {e}", code=4) + try: + with conn.cursor() as _c: # best-effort native server-side bound (seconds; 0 = no timeout) + _c.execute(f"SET STATEMENT_TIMEOUT = {_resolve_timeout_s()}") + except Exception: + pass # older runtime without the config; cur.cancel() below is the cross-thread cancel try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # cur.cancel() is a genuine server-side cancel of the running statement. + result = _run_bounded(lambda: _exec(cur, sql), cur.cancel) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Databricks execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -742,20 +801,21 @@ def _run_trino(creds: dict[str, str], sql: str) -> ExecResult: schema=creds.get("schema"), http_scheme="https" if creds.get("password") else "http", auth=auth, + # Native server-side cap on total query runtime. + session_properties={"query_max_run_time": f"{_resolve_timeout_s()}s"}, ) except Exception as e: raise ExecutorError(f"Trino connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # cur.cancel() cancels the running Trino query server-side. + result = _run_bounded(lambda: _exec(cur, sql), cur.cancel) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Trino execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -771,15 +831,15 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"DuckDB open failed: {e}", code=4) try: - cur = conn.execute(sql) - result = _collect_cursor(cur) + # DuckDB runs the query in conn.execute() and returns the streaming source; conn.interrupt() + # aborts an in-flight query (in-process, thread-safe) — same shape as SQLite. + result = _run_bounded(lambda: conn.execute(sql), conn.interrupt) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"DuckDB execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -859,6 +919,123 @@ def _write_cursor_csv(cur: Any) -> None: _emit_result_csv(_collect_cursor(cur)) +# --------------------------------------------------------------------------- +# Per-statement timeout — the availability guarantee +# +# A generated query can hang or exhaust the DB (recursive CTE, cartesian bomb, unbounded scan). The +# read-only role is not mandated to carry a role-level statement_timeout, so this app-layer bound is +# the SOLE availability guarantee: a wall-clock watchdog cancels the in-flight query after +# AGAMI_SQL_TIMEOUT_S (default 30) — the universal backstop for every engine — layered under each +# engine's native statement timeout where it has one (the primary, genuine DB-side cancel). A fired +# deadline emits a `resource_limit` refusal; the query is killed, no result is returned. +# --------------------------------------------------------------------------- + +_DEFAULT_TIMEOUT_S = 30 # per-statement wall-clock timeout, overridable by AGAMI_SQL_TIMEOUT_S + + +def _resolve_timeout_s() -> int: + """Per-statement timeout in seconds — `AGAMI_SQL_TIMEOUT_S` (default 30). A missing / invalid / + non-positive value falls back to the default; there is no "0 = unlimited" (availability is a + safety obligation — a query is always bounded). Read once at the point of use, like + `_resolve_row_cap`.""" + raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() + return int(raw) if raw.isdigit() and int(raw) > 0 else _DEFAULT_TIMEOUT_S + + +@contextmanager +def _deadline(cancel: Callable[[], None]): + """Bound the enclosed statement by a wall-clock watchdog: after `_resolve_timeout_s()` seconds, + `cancel()` interrupts the in-flight query (per driver — `conn.cancel()` / `conn.interrupt()` / + `cur.cancel()`), which makes the blocked execute/fetch raise. The universal availability backstop + so no engine hangs past the deadline, even one with no native statement timeout. Yields an Event + that is set iff the deadline fired (so the caller can tell a timeout from a real error).""" + fired = threading.Event() + + def _fire() -> None: + fired.set() + try: + cancel() + except Exception: + pass # best-effort cancel; the deadline having fired is what the caller keys on + + timer = threading.Timer(_resolve_timeout_s(), _fire) + timer.daemon = True + timer.start() + try: + yield fired + finally: + timer.cancel() + + +class _ResourceLimit(Exception): + """The per-statement deadline fired and the query was cancelled. Raised out through the engine's + `with conn` / transaction handling so it ROLLS BACK (a cancelled query leaves the txn aborted) + rather than committing; the engine re-raises it and `execute_guarded` maps it to the typed + `resource_limit` refusal (one refusal, both surfaces).""" + + +def _timeout_ms() -> int: + """The statement timeout in milliseconds — for the engines whose native param takes ms.""" + return _resolve_timeout_s() * 1000 + + +def _safe_close(conn: Any) -> None: + """Close a connection, swallowing errors — safe to call twice (the watchdog may close it to + interrupt an in-flight query, then the engine's `finally` closes it again).""" + try: + conn.close() + except Exception: + pass + + +def _exec(cur: Any, sql: str) -> Any: + """Run `sql` on a DB-API cursor and return it as the streaming source (for `_run_bounded`).""" + cur.execute(sql) + return cur + + +def _resource_limit_refusal() -> Refusal: + """Build the typed `resource_limit` refusal for a query cancelled at the statement deadline. The + single source of the timeout refusal — `execute_guarded` raises it as a `GuardRefused` so both the + subprocess wire and the in-process path surface the identical refused Envelope.""" + secs = _resolve_timeout_s() + return Refusal( + kind="resource_limit", + reason=f"the query exceeded the {secs}s statement timeout and was cancelled", + remediation=f"Narrow the query, or raise AGAMI_SQL_TIMEOUT_S (currently {secs}s).", + ) + + +def _deadline_hit(fired: threading.Event, started: float, timeout_s: int) -> bool: + """True if the query was ended by the statement deadline. Primary signal: the watchdog `fired` + (it sets the flag before it cancels, so a watchdog-driven kill is always flagged). Secondary + signal: wall-clock elapsed reached the timeout — an engine's NATIVE server timeout (set to the + same duration) can win the race and raise before the watchdog's callback runs, and without the + elapsed check that genuine timeout would be mislabeled a generic error. Either signal → the query + ran to the deadline, so the raise is a `resource_limit`, not a real error.""" + return fired.is_set() or (time.monotonic() - started) >= timeout_s + + +def _run_bounded(execute: Callable[[], Any], cancel: Callable[[], None]) -> ExecResult: + """Run a statement under the deadline and return its bounded ``ExecResult``. `execute()` runs the + SQL and returns the cursor/result to fetch from (a thunk so both the `cur.execute(sql)` engines and + DuckDB's `conn.execute(sql)` fit). On a deadline hit — the client watchdog fired OR wall-clock + reached the timeout (whichever mechanism, client cancel or native server timeout, killed the + query) — it raises `_ResourceLimit`; the engine re-raises it and `execute_guarded` maps it to a + single `resource_limit` refusal, so no partial result is ever returned. A non-timeout error + (raised before the deadline) propagates unchanged to the engine's handler.""" + started = time.monotonic() + timeout_s = _resolve_timeout_s() + with _deadline(cancel) as fired: + try: + cur = execute() + return _collect_cursor(cur) + except Exception: + if _deadline_hit(fired, started, timeout_s): + raise _ResourceLimit from None + raise + + def _hosted() -> bool: """The served (hosted) path is signalled by a configured database — the same signal `tools._load_org` / `Store.from_env` use. On it, a missing model is a safety failure (fail @@ -1045,9 +1222,16 @@ 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). A per-statement + deadline (``_ResourceLimit``) surfaces as the shared ``resource_limit`` refusal on stderr + exit + code 1 — the same JSON ``execute_guarded``/``main`` emit — and ``run()`` raises it BEFORE returning + a result, so no partial CSV is written.""" try: _emit_result_csv(run()) + except _ResourceLimit: + json.dump({"refusal": _resource_limit_refusal().as_dict()}, sys.stderr) + sys.stderr.write("\n") + return 1 except ExecutorError as e: return _err(e.msg, code=e.code) return 0 @@ -1177,7 +1361,14 @@ def execute_guarded( if refusal is not None: raise GuardRefused(refusal, code=1) creds = _load_credentials(profile) - return executor.execute(sql, creds, profile=profile) + try: + return executor.execute(sql, creds, profile=profile) + except _ResourceLimit: + # The per-statement deadline fired inside the engine (client watchdog or native server + # timeout). Map it to the shared refusal HERE — the one place both surfaces funnel through — + # so the timeout is a `resource_limit` refused Envelope with no partial data, identically for + # the subprocess wire and the in-process handler. + raise GuardRefused(_resource_limit_refusal(), code=1) from None def main() -> int: diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index 3211fbe3..d60b99b5 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -45,8 +45,11 @@ import os import stat import sys +import threading +import time import urllib.parse from collections.abc import Callable +from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass from pathlib import Path @@ -424,15 +427,30 @@ def _run_postgres(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"Postgres connect failed: {e}", code=4) try: with conn: + # Genuine server-side cancel: the backend aborts the statement itself after the deadline. + # Set per-transaction via SET LOCAL — NOT the libpq `options` startup parameter, which a + # transaction-mode pooler (Supabase Supavisor, PgBouncer) can reject at connect. SET LOCAL + # is transaction-scoped, which those poolers pass through cleanly. The watchdog's + # conn.cancel() (a pg_cancel request) is the client-initiated backstop. + with conn.cursor() as setup: + setup.execute(f"SET LOCAL statement_timeout = {_timeout_ms()}") # 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) # 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") # not `with`: a cancelled txn makes CLOSE raise + cur.itersize = _resolve_row_cap() + 1 # server fetch batch = the bounded window + try: + result = _run_bounded(lambda: _exec(cur, sql), conn.cancel) + finally: + try: + cur.close() + except Exception: + pass # a cancelled query leaves the txn aborted; closing the server-side cursor + # would raise and mask _ResourceLimit — swallow so the timeout refusal survives + except _ResourceLimit: + raise # timed out; `with conn` rolled back — execute_guarded maps it to a resource_limit refusal except Exception as e: raise ExecutorError(f"Postgres execution error: {e}", code=5) finally: @@ -446,6 +464,7 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: except ImportError: raise ExecutorError("pymysql not installed. Run: pip install pymysql", code=3) _require(creds, "host", "port", "user", "password", "database") + secs = _resolve_timeout_s() try: conn = pymysql.connect( host=creds["host"], @@ -455,18 +474,37 @@ def _run_mysql(creds: dict[str, str], sql: str) -> ExecResult: database=creds["database"], charset="utf8mb4", connect_timeout=10, + # The reliable client-side bound: a socket read/write blocked past the deadline raises on + # its own. Closing the connection from the watchdog thread does NOT unblock a recv() on + # Linux (that needs shutdown(), which pymysql's close() skips), so this — not the close — + # is what guarantees MySQL/MariaDB can't hang past the deadline. + read_timeout=secs, + write_timeout=secs, autocommit=True, ) except Exception as e: raise ExecutorError(f"MySQL connect failed: {e}", code=4) + # Native server-side timeout. MySQL 5.7.8+ spells it max_execution_time (ms); MariaDB spells it + # max_statement_time (seconds). Try both — each is a no-op on the other dialect (a swallowed + # unknown-variable error) — so whichever the server is gets a genuine server-side abort. + for stmt in ( + f"SET SESSION max_execution_time={secs * 1000}", + f"SET SESSION max_statement_time={secs}", + ): + try: + with conn.cursor() as _c: + _c.execute(stmt) + except Exception: + pass # not this dialect; read_timeout still bounds the client try: - with conn.cursor() as cur: - cur.execute(sql) - result = _collect_cursor(cur) + cur = conn.cursor() # not `with` — the watchdog may close conn, so avoid a close-on-close + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"MySQL execution error: {e}", code=5) finally: - conn.close() + _safe_close(conn) return result @@ -492,6 +530,8 @@ def _run_snowflake(creds: dict[str, str], sql: str) -> ExecResult: "user": creds["user"], "client_session_keep_alive": False, "login_timeout": 15, + # Native server-side statement timeout (seconds) — Snowflake aborts the query itself. + "session_parameters": {"STATEMENT_TIMEOUT_IN_SECONDS": _resolve_timeout_s()}, } for k in ("password", "warehouse", "database", "schema", "role", "authenticator"): if creds.get(k): @@ -502,15 +542,13 @@ 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) + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Snowflake execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -573,7 +611,8 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: # 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] = {} + # `job_timeout_ms` is the native server-side job timeout — BigQuery ends the job itself. + job_config_kwargs: dict[str, Any] = {"job_timeout_ms": _timeout_ms()} if creds.get("dataset"): try: job_config_kwargs["default_dataset"] = f"{project}.{creds['dataset']}" @@ -581,16 +620,30 @@ def _run_bigquery(creds: dict[str, str], sql: str) -> ExecResult: pass cap = _resolve_row_cap() + job_box: dict[str, Any] = {} + started = time.monotonic() + timeout_s = _resolve_timeout_s() + + def _bq_cancel() -> None: + job = job_box.get("job") + if job is not None: + job.cancel() # genuine server-side job cancel (the client-initiated backstop to job_timeout_ms) + 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) - # 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 + with _deadline(_bq_cancel) as fired: + try: + job = client.query(sql, job_config=bigquery.QueryJobConfig(**job_config_kwargs)) + job_box["job"] = job + # 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 Exception: + if _deadline_hit(fired, started, timeout_s): + raise _ResourceLimit from None + raise + except _ResourceLimit: + raise # execute_guarded maps a fired BigQuery deadline to a resource_limit refusal except Exception as e: raise ExecutorError(f"BigQuery execution error: {e}", code=5) @@ -619,8 +672,9 @@ def _run_sqlite(creds: dict[str, str], sql: str) -> ExecResult: raise ExecutorError(f"SQLite connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + result = _run_bounded(lambda: _exec(cur, sql), conn.interrupt) + except _ResourceLimit: + raise # timed out — execute_guarded maps it to a resource_limit refusal except Exception as e: raise ExecutorError(f"SQLite execution error: {e}", code=5) finally: @@ -643,20 +697,19 @@ def _run_sqlserver(creds: dict[str, str], sql: str) -> ExecResult: password=creds["password"], database=creds.get("database", ""), login_timeout=15, + timeout=_resolve_timeout_s(), # native per-query timeout (seconds) ) except Exception as e: raise ExecutorError(f"SQL Server connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"SQL Server execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -677,17 +730,19 @@ 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) + try: + conn.call_timeout = _timeout_ms() # native round-trip timeout (ms) + except Exception: + pass # older driver / mode; the wall-clock watchdog still bounds the query try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + result = _run_bounded(lambda: _exec(cur, sql), lambda: _safe_close(conn)) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Oracle execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -709,17 +764,21 @@ def _run_databricks(creds: dict[str, str], sql: str) -> ExecResult: ) except Exception as e: raise ExecutorError(f"Databricks connect failed: {e}", code=4) + try: + with conn.cursor() as _c: # best-effort native server-side bound (seconds; 0 = no timeout) + _c.execute(f"SET STATEMENT_TIMEOUT = {_resolve_timeout_s()}") + except Exception: + pass # older runtime without the config; cur.cancel() below is the cross-thread cancel try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # cur.cancel() is a genuine server-side cancel of the running statement. + result = _run_bounded(lambda: _exec(cur, sql), cur.cancel) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Databricks execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -742,20 +801,21 @@ def _run_trino(creds: dict[str, str], sql: str) -> ExecResult: schema=creds.get("schema"), http_scheme="https" if creds.get("password") else "http", auth=auth, + # Native server-side cap on total query runtime. + session_properties={"query_max_run_time": f"{_resolve_timeout_s()}s"}, ) except Exception as e: raise ExecutorError(f"Trino connect failed: {e}", code=4) try: cur = conn.cursor() - cur.execute(sql) - result = _collect_cursor(cur) + # cur.cancel() cancels the running Trino query server-side. + result = _run_bounded(lambda: _exec(cur, sql), cur.cancel) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"Trino execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -771,15 +831,15 @@ def _run_duckdb(creds: dict[str, str], sql: str) -> ExecResult: except Exception as e: raise ExecutorError(f"DuckDB open failed: {e}", code=4) try: - cur = conn.execute(sql) - result = _collect_cursor(cur) + # DuckDB runs the query in conn.execute() and returns the streaming source; conn.interrupt() + # aborts an in-flight query (in-process, thread-safe) — same shape as SQLite. + result = _run_bounded(lambda: conn.execute(sql), conn.interrupt) + except _ResourceLimit: + raise except Exception as e: raise ExecutorError(f"DuckDB execution error: {e}", code=5) finally: - try: - conn.close() - except Exception: - pass + _safe_close(conn) return result @@ -859,6 +919,123 @@ def _write_cursor_csv(cur: Any) -> None: _emit_result_csv(_collect_cursor(cur)) +# --------------------------------------------------------------------------- +# Per-statement timeout — the availability guarantee +# +# A generated query can hang or exhaust the DB (recursive CTE, cartesian bomb, unbounded scan). The +# read-only role is not mandated to carry a role-level statement_timeout, so this app-layer bound is +# the SOLE availability guarantee: a wall-clock watchdog cancels the in-flight query after +# AGAMI_SQL_TIMEOUT_S (default 30) — the universal backstop for every engine — layered under each +# engine's native statement timeout where it has one (the primary, genuine DB-side cancel). A fired +# deadline emits a `resource_limit` refusal; the query is killed, no result is returned. +# --------------------------------------------------------------------------- + +_DEFAULT_TIMEOUT_S = 30 # per-statement wall-clock timeout, overridable by AGAMI_SQL_TIMEOUT_S + + +def _resolve_timeout_s() -> int: + """Per-statement timeout in seconds — `AGAMI_SQL_TIMEOUT_S` (default 30). A missing / invalid / + non-positive value falls back to the default; there is no "0 = unlimited" (availability is a + safety obligation — a query is always bounded). Read once at the point of use, like + `_resolve_row_cap`.""" + raw = os.environ.get("AGAMI_SQL_TIMEOUT_S", "").strip() + return int(raw) if raw.isdigit() and int(raw) > 0 else _DEFAULT_TIMEOUT_S + + +@contextmanager +def _deadline(cancel: Callable[[], None]): + """Bound the enclosed statement by a wall-clock watchdog: after `_resolve_timeout_s()` seconds, + `cancel()` interrupts the in-flight query (per driver — `conn.cancel()` / `conn.interrupt()` / + `cur.cancel()`), which makes the blocked execute/fetch raise. The universal availability backstop + so no engine hangs past the deadline, even one with no native statement timeout. Yields an Event + that is set iff the deadline fired (so the caller can tell a timeout from a real error).""" + fired = threading.Event() + + def _fire() -> None: + fired.set() + try: + cancel() + except Exception: + pass # best-effort cancel; the deadline having fired is what the caller keys on + + timer = threading.Timer(_resolve_timeout_s(), _fire) + timer.daemon = True + timer.start() + try: + yield fired + finally: + timer.cancel() + + +class _ResourceLimit(Exception): + """The per-statement deadline fired and the query was cancelled. Raised out through the engine's + `with conn` / transaction handling so it ROLLS BACK (a cancelled query leaves the txn aborted) + rather than committing; the engine re-raises it and `execute_guarded` maps it to the typed + `resource_limit` refusal (one refusal, both surfaces).""" + + +def _timeout_ms() -> int: + """The statement timeout in milliseconds — for the engines whose native param takes ms.""" + return _resolve_timeout_s() * 1000 + + +def _safe_close(conn: Any) -> None: + """Close a connection, swallowing errors — safe to call twice (the watchdog may close it to + interrupt an in-flight query, then the engine's `finally` closes it again).""" + try: + conn.close() + except Exception: + pass + + +def _exec(cur: Any, sql: str) -> Any: + """Run `sql` on a DB-API cursor and return it as the streaming source (for `_run_bounded`).""" + cur.execute(sql) + return cur + + +def _resource_limit_refusal() -> Refusal: + """Build the typed `resource_limit` refusal for a query cancelled at the statement deadline. The + single source of the timeout refusal — `execute_guarded` raises it as a `GuardRefused` so both the + subprocess wire and the in-process path surface the identical refused Envelope.""" + secs = _resolve_timeout_s() + return Refusal( + kind="resource_limit", + reason=f"the query exceeded the {secs}s statement timeout and was cancelled", + remediation=f"Narrow the query, or raise AGAMI_SQL_TIMEOUT_S (currently {secs}s).", + ) + + +def _deadline_hit(fired: threading.Event, started: float, timeout_s: int) -> bool: + """True if the query was ended by the statement deadline. Primary signal: the watchdog `fired` + (it sets the flag before it cancels, so a watchdog-driven kill is always flagged). Secondary + signal: wall-clock elapsed reached the timeout — an engine's NATIVE server timeout (set to the + same duration) can win the race and raise before the watchdog's callback runs, and without the + elapsed check that genuine timeout would be mislabeled a generic error. Either signal → the query + ran to the deadline, so the raise is a `resource_limit`, not a real error.""" + return fired.is_set() or (time.monotonic() - started) >= timeout_s + + +def _run_bounded(execute: Callable[[], Any], cancel: Callable[[], None]) -> ExecResult: + """Run a statement under the deadline and return its bounded ``ExecResult``. `execute()` runs the + SQL and returns the cursor/result to fetch from (a thunk so both the `cur.execute(sql)` engines and + DuckDB's `conn.execute(sql)` fit). On a deadline hit — the client watchdog fired OR wall-clock + reached the timeout (whichever mechanism, client cancel or native server timeout, killed the + query) — it raises `_ResourceLimit`; the engine re-raises it and `execute_guarded` maps it to a + single `resource_limit` refusal, so no partial result is ever returned. A non-timeout error + (raised before the deadline) propagates unchanged to the engine's handler.""" + started = time.monotonic() + timeout_s = _resolve_timeout_s() + with _deadline(cancel) as fired: + try: + cur = execute() + return _collect_cursor(cur) + except Exception: + if _deadline_hit(fired, started, timeout_s): + raise _ResourceLimit from None + raise + + def _hosted() -> bool: """The served (hosted) path is signalled by a configured database — the same signal `tools._load_org` / `Store.from_env` use. On it, a missing model is a safety failure (fail @@ -1045,9 +1222,16 @@ 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). A per-statement + deadline (``_ResourceLimit``) surfaces as the shared ``resource_limit`` refusal on stderr + exit + code 1 — the same JSON ``execute_guarded``/``main`` emit — and ``run()`` raises it BEFORE returning + a result, so no partial CSV is written.""" try: _emit_result_csv(run()) + except _ResourceLimit: + json.dump({"refusal": _resource_limit_refusal().as_dict()}, sys.stderr) + sys.stderr.write("\n") + return 1 except ExecutorError as e: return _err(e.msg, code=e.code) return 0 @@ -1177,7 +1361,14 @@ def execute_guarded( if refusal is not None: raise GuardRefused(refusal, code=1) creds = _load_credentials(profile) - return executor.execute(sql, creds, profile=profile) + try: + return executor.execute(sql, creds, profile=profile) + except _ResourceLimit: + # The per-statement deadline fired inside the engine (client watchdog or native server + # timeout). Map it to the shared refusal HERE — the one place both surfaces funnel through — + # so the timeout is a `resource_limit` refused Envelope with no partial data, identically for + # the subprocess wire and the in-process handler. + raise GuardRefused(_resource_limit_refusal(), code=1) from None def main() -> int: diff --git a/plugins/agami/skills/agami-deploy/bundle/agami.env.example b/plugins/agami/skills/agami-deploy/bundle/agami.env.example index 0c034ddd..20a496bf 100644 --- a/plugins/agami/skills/agami-deploy/bundle/agami.env.example +++ b/plugins/agami/skills/agami-deploy/bundle/agami.env.example @@ -54,3 +54,11 @@ AGAMI_ADMIN_PASSWORD= # LATERAL source. This is the safe default. `warn` logs and allows instead; use it only as a # temporary staged-rollout escape hatch, never as the shipped setting. # AGAMI_SQL_UNSCOPABLE_POSTURE=enforce + +# Optional: SQL resource limits — bound a runaway query (a recursive CTE, a cartesian product, an +# unbounded scan). Both are enforced at the executor, so every SQL surface inherits them. +# AGAMI_SQL_TIMEOUT_S=30 # per-statement timeout in seconds; a query that exceeds it is cancelled +# # (server-side where the engine supports it) and refused. Default 30; +# # keep it below ~240 (the executor's own process wall). +# AGAMI_SQL_MAX_ROWS=1000 # max rows returned per query; a larger result is truncated and flagged, +# # never silently cut. Raise it for bigger pulls. Default 1000. diff --git a/tests/test_ace044_bounded_fetch.py b/tests/test_ace044_bounded_fetch.py index 038c6e01..52fad597 100644 --- a/tests/test_ace044_bounded_fetch.py +++ b/tests/test_ace044_bounded_fetch.py @@ -192,6 +192,9 @@ def __enter__(self): def __exit__(self, *a): return False + def cancel(self): # the per-statement watchdog's cancel primitive (no-op here — never fires) + pass + def close(self): pass diff --git a/tests/test_execute_sql_envelope.py b/tests/test_execute_sql_envelope.py index a1cce809..e0c62df4 100644 --- a/tests/test_execute_sql_envelope.py +++ b/tests/test_execute_sql_envelope.py @@ -73,6 +73,28 @@ def test_refusal_line_is_found_among_interleaved_notices(monkeypatch): assert env["status"] == "refused" and env["refusal"]["kind"] == "sensitive_columns" +def test_resource_limit_refusal_discards_partial_stdout(monkeypatch): + # The per-statement timeout emits a {"refusal": {"kind": "resource_limit"}} line from the + # executor. Both transports (stdio + HTTP) funnel through tool_execute_sql, so this single relay + # is the both-surfaces guarantee. Critically: a streaming engine may have flushed a partial CSV + # (e.g. the header row) to stdout before the cancel — the relay must DISCARD it on the non-zero + # exit, so a killed query is a refused Envelope with no data. Feed partial stdout to pin that. + stderr = json.dumps( + { + "refusal": { + "kind": "resource_limit", + "reason": "the query exceeded the 30s statement timeout and was cancelled", + "remediation": "Narrow the query, or raise AGAMI_SQL_TIMEOUT_S (currently 30s).", + } + } + ) + env = _run(monkeypatch, _Proc(1, stdout="n\n1\n2\n", stderr=stderr)) + assert env["status"] == "refused" + assert env["refusal"]["kind"] == "resource_limit" + assert "data" not in env # the partial CSV on stdout is discarded — no partial data leaks + assert env["audit_id"] + + def test_bare_operational_failure_is_classified_by_exit_code(monkeypatch): # Exit 5 = SQL execution error with no {"refusal"} line → classified as 'syntax' by exit code. env = _run(monkeypatch, _Proc(5, stderr='relation "x" does not exist')) diff --git a/tests/test_resource_limits.py b/tests/test_resource_limits.py new file mode 100644 index 00000000..bcc08915 --- /dev/null +++ b/tests/test_resource_limits.py @@ -0,0 +1,362 @@ +"""ACE-038 — per-statement timeout (the availability guarantee). + +The row-cap half shipped via ACE-044; this covers the timeout. The genuine-cancel proof runs +in-process against SQLite (a real recursive-CTE bomb is interrupted by the wall-clock watchdog and +returns a `resource_limit` refusal); the cloud engines' native mechanisms are asserted-by-config in +unit tests (CI has no live warehouses — same constraint as ACE-037's SC5). +""" + +from __future__ import annotations + +import json +import sqlite3 +import sys +import threading +import types +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT / "plugins" / "agami" / "scripts")) +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 + +# A recursive CTE that counts to 2 billion — seconds of pure DB-side work, no memory growth (it +# aggregates, storing no rows), so a sub-second timeout must interrupt it mid-flight. +_BOMB = ( + "WITH RECURSIVE r(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM r WHERE n < 2000000000) " + "SELECT count(*) FROM r" +) + + +# ── _resolve_timeout_s ─────────────────────────────────────────────────────── + + +def test_timeout_default_is_30(monkeypatch): + monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) + assert execute_sql._resolve_timeout_s() == 30 + + +@pytest.mark.parametrize("val,expected", [("5", 5), ("", 30), ("nope", 30), ("0", 30), ("-3", 30)]) +def test_timeout_env_override_and_invalid_fall_back(monkeypatch, val, expected): + # A missing / invalid / non-positive value falls back to the default — never "0 = unlimited". + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", val) + assert execute_sql._resolve_timeout_s() == expected + + +# ── the genuine cancel (SQLite, in-process) ────────────────────────────────── + + +def test_sqlite_runaway_is_killed_with_resource_limit(tmp_path, monkeypatch, capsys): + # The watchdog genuinely cancels an in-flight query: a recursive-CTE bomb under a 1 s timeout is + # interrupted and returns a resource_limit refusal — not hung, not a generic execution error. + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "1") + db = tmp_path / "t.db" + sqlite3.connect(str(db)).close() # create the file + + code = execute_sql._execute_sqlite({"path": str(db)}, _BOMB) + + assert code == 1 # killed (refused), not 0 (success) or 5 (generic error) + out = capsys.readouterr() + refusal = json.loads(out.err.strip().splitlines()[-1])["refusal"] + assert refusal["kind"] == "resource_limit" + assert "timeout" in refusal["reason"] + assert out.out == "" # the cancelled aggregate wrote nothing to stdout — no partial data + + +def test_duckdb_runaway_is_killed_with_resource_limit(tmp_path, monkeypatch, capsys): + # The other in-process engine: conn.interrupt() genuinely aborts a DuckDB runaway the same way. + duckdb = pytest.importorskip("duckdb") + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "1") + db = tmp_path / "t.duckdb" + duckdb.connect(str(db)).close() # create the file so _execute_duckdb can open it read-only + + code = execute_sql._execute_duckdb({"path": str(db)}, _BOMB) + + assert code == 1 + out = capsys.readouterr() + refusal = json.loads(out.err.strip().splitlines()[-1])["refusal"] + assert refusal["kind"] == "resource_limit" + + +def test_sqlite_fast_query_still_returns(tmp_path, monkeypatch, capsys): + # A quick query well under the timeout is unaffected — the watchdog is a bound, not a tax. + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + db = tmp_path / "t.db" + con = sqlite3.connect(str(db)) + con.execute("CREATE TABLE t (id INTEGER)") + con.execute("INSERT INTO t VALUES (1), (2), (3)") + con.commit() + con.close() + + code = execute_sql._execute_sqlite({"path": str(db)}, "SELECT count(*) AS n FROM t") + + assert code == 0 + out = capsys.readouterr() + assert out.out.splitlines() == ["n", "3"] # header + the real result + assert '"resource_limit"' not in out.err + + +# ── native server-side timeouts (asserted-by-config) ───────────────────────── +# +# CI has no live warehouse, so each engine's native timeout is asserted by construction: the right +# param, in the right UNITS, is handed to the driver. This guards the ms-vs-s mix directly — a +# `_timeout_ms()`/`_resolve_timeout_s()` swap would make an engine's timeout 1000x wrong and silently +# disable the availability guarantee. (Mirrors the ACE-037 SC5 call — the runnable path is proven +# in-process, the cloud path asserted by config.) + + +class _RecCursor: + """A DB-API cursor that records every executed statement and yields one row.""" + + description = [("n",)] + itersize = 0 + + def __init__(self, log): + self._log = log + + def execute(self, sql, *_a): + self._log.append(sql) + + def fetchmany(self, _n): + return [(1,)] + + def cancel(self): + pass + + def close(self): + pass + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + +class _RecConn: + """A DB-API connection over a shared statement log; `cursor(...)` ignores kwargs like `name`.""" + + def __init__(self, log): + self._log = log + + def cursor(self, *_a, **_k): + return _RecCursor(self._log) + + def cancel(self): + pass + + def close(self): + pass + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + +def _set_timeout_env(monkeypatch, timeout_s): + if timeout_s: + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", timeout_s) + else: + monkeypatch.delenv("AGAMI_SQL_TIMEOUT_S", raising=False) + + +@pytest.mark.parametrize("timeout_s,expected_ms", [("30", 30000), ("5", 5000), ("", 30000)]) +def test_postgres_sets_statement_timeout_via_set_local(monkeypatch, timeout_s, expected_ms): + log: list = [] + connect_kw: dict = {} + fake = types.ModuleType("psycopg2") + fake.connect = lambda **kw: (connect_kw.update(kw), _RecConn(log))[1] + monkeypatch.setitem(sys.modules, "psycopg2", fake) + _set_timeout_env(monkeypatch, timeout_s) + + creds = {"host": "h", "port": "5432", "user": "u", "password": "p", "database": "d"} + code = execute_sql._execute_postgres(creds, "SELECT 1 AS n") + + assert code == 0 + # Genuine server-side abort, delivered per-transaction via SET LOCAL (ms)... + assert f"SET LOCAL statement_timeout = {expected_ms}" in log + # ...NOT the libpq `options` startup param a transaction-mode pooler (Supabase/PgBouncer) rejects. + assert "options" not in connect_kw + + +def test_mysql_native_timeouts_and_client_read_timeout(monkeypatch, capsys): + log: list = [] + connect_kw: dict = {} + fake = types.ModuleType("pymysql") + fake.connect = lambda **kw: (connect_kw.update(kw), _RecConn(log))[1] + monkeypatch.setitem(sys.modules, "pymysql", fake) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + + creds = {"host": "h", "port": "3306", "user": "u", "password": "p", "database": "d"} + code = execute_sql._execute_mysql(creds, "SELECT 1 AS n") + + assert code == 0 + # The reliable client bound (a close from the watchdog thread can't unblock a recv on Linux). + assert connect_kw["read_timeout"] == 30 and connect_kw["write_timeout"] == 30 + assert "SET SESSION max_execution_time=30000" in log # MySQL 5.7.8+ (ms) + assert "SET SESSION max_statement_time=30" in log # MariaDB (seconds) + + +def test_snowflake_native_statement_timeout(monkeypatch): + connect_kw: dict = {} + connector = types.ModuleType("snowflake.connector") + connector.connect = lambda **kw: (connect_kw.update(kw), _RecConn([]))[1] + pkg = types.ModuleType("snowflake") + pkg.connector = connector + monkeypatch.setitem(sys.modules, "snowflake", pkg) + monkeypatch.setitem(sys.modules, "snowflake.connector", connector) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + + code = execute_sql._execute_snowflake( + {"account": "a", "user": "u", "password": "p"}, "SELECT 1 AS n" + ) + + assert code == 0 + assert connect_kw["session_parameters"] == {"STATEMENT_TIMEOUT_IN_SECONDS": 30} # seconds + + +def test_bigquery_native_job_timeout(monkeypatch): + cfg: dict = {} + + class _Job: + def result(self, max_results=None): + return types.SimpleNamespace(schema=[]) + + def cancel(self): + pass + + class _Client: + def __init__(self, **_kw): + pass + + def query(self, _sql, **_kw): + return _Job() + + gcloud = types.ModuleType("google.cloud") + gcloud.bigquery = types.SimpleNamespace( + Client=_Client, QueryJobConfig=lambda **k: (cfg.update(k), object())[1] + ) + goauth = types.ModuleType("google.oauth2") + goauth.service_account = types.SimpleNamespace() + monkeypatch.setitem(sys.modules, "google.cloud", gcloud) + monkeypatch.setitem(sys.modules, "google.oauth2", goauth) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + + code = execute_sql._execute_bigquery({"project": "p"}, "SELECT 1 AS n") + + assert code == 0 + assert cfg["job_timeout_ms"] == 30000 # milliseconds + + +def test_sqlserver_native_query_timeout(monkeypatch): + connect_kw: dict = {} + fake = types.ModuleType("pymssql") + fake.connect = lambda **kw: (connect_kw.update(kw), _RecConn([]))[1] + monkeypatch.setitem(sys.modules, "pymssql", fake) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + + code = execute_sql._execute_sqlserver( + {"host": "h", "user": "u", "password": "p"}, "SELECT 1 AS n" + ) + + assert code == 0 + assert connect_kw["timeout"] == 30 # pymssql per-query timeout (seconds) + + +def test_oracle_native_call_timeout(monkeypatch): + conn = _RecConn([]) + fake = types.ModuleType("oracledb") + fake.connect = lambda **_kw: conn + fake.makedsn = lambda *_a, **_k: "dsn" + monkeypatch.setitem(sys.modules, "oracledb", fake) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + + code = execute_sql._execute_oracle({"user": "u", "password": "p", "dsn": "d"}, "SELECT 1 AS n") + + assert code == 0 + assert conn.call_timeout == 30000 # oracledb round-trip timeout (ms) + + +def test_trino_native_query_max_run_time(monkeypatch): + connect_kw: dict = {} + fake = types.ModuleType("trino") + fake.dbapi = types.SimpleNamespace( + connect=lambda **kw: (connect_kw.update(kw), _RecConn([]))[1] + ) + fake.auth = types.SimpleNamespace(BasicAuthentication=lambda *_a: object()) + monkeypatch.setitem(sys.modules, "trino", fake) + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "30") + + code = execute_sql._execute_trino({"host": "h", "user": "u"}, "SELECT 1 AS n") + + assert code == 0 + assert connect_kw["session_properties"] == {"query_max_run_time": "30s"} # seconds-string + + +def test_postgres_timeout_returns_refusal_exit_even_when_cursor_close_raises(monkeypatch, capsys): + # A cancelled Postgres query aborts the txn, so closing the server-side cursor raises. That close + # must NOT mask _ResourceLimit — the executor must still return exit 1 with the resource_limit + # refusal (not exit 5 / a generic error). Deterministic: the watchdog cancel unblocks execute(). + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "1") + cancelled = threading.Event() + + class _NamedCursor: + description = [("n",)] + itersize = 0 + + def execute(self, _sql): + if cancelled.wait( + timeout=5 + ): # blocks like a real slow query until the watchdog cancels + raise RuntimeError("canceling statement due to statement timeout") + + def fetchmany(self, _n): + return [] + + def close(self): + if cancelled.is_set(): + raise RuntimeError("current transaction is aborted") # the masking error + + class _SetupCursor: + def execute(self, _sql): + pass + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + class _Conn: + def cursor(self, name=None): + return _NamedCursor() if name else _SetupCursor() + + def cancel(self): + cancelled.set() # the watchdog's pg_cancel — unblocks the blocked execute() + + def close(self): + pass + + def __enter__(self): + return self + + def __exit__(self, *_a): + return False + + fake = types.ModuleType("psycopg2") + fake.connect = lambda **_kw: _Conn() + monkeypatch.setitem(sys.modules, "psycopg2", fake) + + creds = {"host": "h", "port": "5432", "user": "u", "password": "p", "database": "d"} + code = execute_sql._execute_postgres(creds, "SELECT 1 AS n") + + assert code == 1 # refused, NOT 5 — the cursor-close error did not mask _ResourceLimit + refusal = json.loads(capsys.readouterr().err.strip().splitlines()[-1])["refusal"] + assert refusal["kind"] == "resource_limit" From 17fe5854aadf97f85b9fcb60c79f4f637a494ec4 Mon Sep 17 00:00:00 2001 From: Sandeep Date: Sun, 12 Jul 2026 19:09:40 -0700 Subject: [PATCH 2/2] =?UTF-8?q?ACE-038:=20address=20review=20=E2=80=94=20?= =?UTF-8?q?=5Fdeadline=20takes=20the=20resolved=20timeout=20(no=20env=20re?= =?UTF-8?q?-read)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Copilot review: _run_bounded (and the BigQuery path) resolve `timeout_s` once, feed it to `_deadline_hit` for the timeout-vs-real-error classification, but _deadline independently RE-READ `AGAMI_SQL_TIMEOUT_S` when constructing the watchdog Timer. A mid-flight change to the env var could make the watchdog duration and the classification diverge. Pass the already-resolved `timeout_s` into _deadline so both key off the SAME value; the env is read once per call, at the call site. Test pins that _deadline schedules the Timer with the passed value, not a fresh env read. Spec: ACE-038 Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/agami-core/src/execute_sql.py | 22 +++++++++++++--------- plugins/agami/lib/execute_sql.py | 22 +++++++++++++--------- tests/test_resource_limits.py | 24 ++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 18 deletions(-) diff --git a/packages/agami-core/src/execute_sql.py b/packages/agami-core/src/execute_sql.py index d60b99b5..b95e61dd 100644 --- a/packages/agami-core/src/execute_sql.py +++ b/packages/agami-core/src/execute_sql.py @@ -630,7 +630,7 @@ def _bq_cancel() -> None: job.cancel() # genuine server-side job cancel (the client-initiated backstop to job_timeout_ms) try: - with _deadline(_bq_cancel) as fired: + with _deadline(_bq_cancel, timeout_s) as fired: try: job = client.query(sql, job_config=bigquery.QueryJobConfig(**job_config_kwargs)) job_box["job"] = job @@ -943,12 +943,16 @@ def _resolve_timeout_s() -> int: @contextmanager -def _deadline(cancel: Callable[[], None]): - """Bound the enclosed statement by a wall-clock watchdog: after `_resolve_timeout_s()` seconds, - `cancel()` interrupts the in-flight query (per driver — `conn.cancel()` / `conn.interrupt()` / - `cur.cancel()`), which makes the blocked execute/fetch raise. The universal availability backstop - so no engine hangs past the deadline, even one with no native statement timeout. Yields an Event - that is set iff the deadline fired (so the caller can tell a timeout from a real error).""" +def _deadline(cancel: Callable[[], None], timeout_s: int): + """Bound the enclosed statement by a wall-clock watchdog: after `timeout_s` seconds, `cancel()` + interrupts the in-flight query (per driver — `conn.cancel()` / `conn.interrupt()` / `cur.cancel()`), + which makes the blocked execute/fetch raise. The universal availability backstop so no engine hangs + past the deadline, even one with no native statement timeout. Yields an Event that is set iff the + deadline fired (so the caller can tell a timeout from a real error). + + `timeout_s` is passed in (not re-read from the env here) so the watchdog duration and the caller's + `_deadline_hit` elapsed-vs-timeout classification key off the SAME resolved value — a mid-flight + change to `AGAMI_SQL_TIMEOUT_S` can't make the two diverge (Copilot review).""" fired = threading.Event() def _fire() -> None: @@ -958,7 +962,7 @@ def _fire() -> None: except Exception: pass # best-effort cancel; the deadline having fired is what the caller keys on - timer = threading.Timer(_resolve_timeout_s(), _fire) + timer = threading.Timer(timeout_s, _fire) timer.daemon = True timer.start() try: @@ -1026,7 +1030,7 @@ def _run_bounded(execute: Callable[[], Any], cancel: Callable[[], None]) -> Exec (raised before the deadline) propagates unchanged to the engine's handler.""" started = time.monotonic() timeout_s = _resolve_timeout_s() - with _deadline(cancel) as fired: + with _deadline(cancel, timeout_s) as fired: try: cur = execute() return _collect_cursor(cur) diff --git a/plugins/agami/lib/execute_sql.py b/plugins/agami/lib/execute_sql.py index d60b99b5..b95e61dd 100644 --- a/plugins/agami/lib/execute_sql.py +++ b/plugins/agami/lib/execute_sql.py @@ -630,7 +630,7 @@ def _bq_cancel() -> None: job.cancel() # genuine server-side job cancel (the client-initiated backstop to job_timeout_ms) try: - with _deadline(_bq_cancel) as fired: + with _deadline(_bq_cancel, timeout_s) as fired: try: job = client.query(sql, job_config=bigquery.QueryJobConfig(**job_config_kwargs)) job_box["job"] = job @@ -943,12 +943,16 @@ def _resolve_timeout_s() -> int: @contextmanager -def _deadline(cancel: Callable[[], None]): - """Bound the enclosed statement by a wall-clock watchdog: after `_resolve_timeout_s()` seconds, - `cancel()` interrupts the in-flight query (per driver — `conn.cancel()` / `conn.interrupt()` / - `cur.cancel()`), which makes the blocked execute/fetch raise. The universal availability backstop - so no engine hangs past the deadline, even one with no native statement timeout. Yields an Event - that is set iff the deadline fired (so the caller can tell a timeout from a real error).""" +def _deadline(cancel: Callable[[], None], timeout_s: int): + """Bound the enclosed statement by a wall-clock watchdog: after `timeout_s` seconds, `cancel()` + interrupts the in-flight query (per driver — `conn.cancel()` / `conn.interrupt()` / `cur.cancel()`), + which makes the blocked execute/fetch raise. The universal availability backstop so no engine hangs + past the deadline, even one with no native statement timeout. Yields an Event that is set iff the + deadline fired (so the caller can tell a timeout from a real error). + + `timeout_s` is passed in (not re-read from the env here) so the watchdog duration and the caller's + `_deadline_hit` elapsed-vs-timeout classification key off the SAME resolved value — a mid-flight + change to `AGAMI_SQL_TIMEOUT_S` can't make the two diverge (Copilot review).""" fired = threading.Event() def _fire() -> None: @@ -958,7 +962,7 @@ def _fire() -> None: except Exception: pass # best-effort cancel; the deadline having fired is what the caller keys on - timer = threading.Timer(_resolve_timeout_s(), _fire) + timer = threading.Timer(timeout_s, _fire) timer.daemon = True timer.start() try: @@ -1026,7 +1030,7 @@ def _run_bounded(execute: Callable[[], Any], cancel: Callable[[], None]) -> Exec (raised before the deadline) propagates unchanged to the engine's handler.""" started = time.monotonic() timeout_s = _resolve_timeout_s() - with _deadline(cancel) as fired: + with _deadline(cancel, timeout_s) as fired: try: cur = execute() return _collect_cursor(cur) diff --git a/tests/test_resource_limits.py b/tests/test_resource_limits.py index bcc08915..10466f59 100644 --- a/tests/test_resource_limits.py +++ b/tests/test_resource_limits.py @@ -48,6 +48,30 @@ def test_timeout_env_override_and_invalid_fall_back(monkeypatch, val, expected): assert execute_sql._resolve_timeout_s() == expected +def test_deadline_uses_the_passed_timeout_not_a_reread_env(monkeypatch): + # Copilot review: the watchdog Timer must key off the timeout the caller already resolved (and + # feeds to _deadline_hit), NOT a fresh env read — else a mid-flight change to AGAMI_SQL_TIMEOUT_S + # makes the watchdog duration and the timeout classification diverge. Pin that _deadline schedules + # the Timer with the PASSED value, not the (different) env value a re-read would pick up. + monkeypatch.setenv("AGAMI_SQL_TIMEOUT_S", "99") # what a re-read inside _deadline would have used + captured = {} + + class _FakeTimer: + def __init__(self, interval, fn): + captured["interval"] = interval + + def start(self): + pass + + def cancel(self): + pass + + monkeypatch.setattr(execute_sql.threading, "Timer", _FakeTimer) + with execute_sql._deadline(lambda: None, timeout_s=7): + pass + assert captured["interval"] == 7 # the caller's resolved value, not 99 from the env + + # ── the genuine cancel (SQLite, in-process) ──────────────────────────────────