Add --timeout flag to run-sql, document env var configuration

.claude/skills/otai/SKILL.md

@@ -42,15 +42,31 @@ Positional `<name>` (the dataset to describe) plus `--release` (default
 column names, types, descriptions, and cross-dataset relationships/nested
 subfields, parsed from the release's croissant schema.
 
-### `otai run-sql "<query>"`
+### `otai run-sql "<query>" [--timeout SECONDS]`
 Positional `<query>`, a read-only SQL string. No `--release` flag:
 - Unqualified table names resolve against `latest`.
 - Schema-qualify a table to target a specific (possibly non-latest)
   release, e.g. `"26.03".target`; this also enables cross-release joins in
   a single query, e.g. `"26.06".target JOIN "26.03".target ...`.
 - The CLI enforces read-only SQL, a ~1000-row cap (response says whether
-  results were truncated), and a ~45s timeout — do not attempt to
-  replicate or second-guess these checks yourself.
+  results were truncated), and a timeout — do not attempt to replicate or
+  second-guess these checks yourself.
+- `--timeout SECONDS` overrides the default ~45s timeout for this call
+  only. Use it when a query is legitimately slow but useful (e.g. a full
+  aggregate over a very large dataset) rather than a mistake to fix — see
+  rule 5 below for when to reach for it instead of narrowing the query.
+
+## Environment variables
+
+Configuration is via env vars, not CLI flags (except `--timeout`, which is
+per-call). You normally won't need to set any of these — the defaults are
+correct for regular use — but they're worth knowing about:
+
+| Variable            | Default                          | Purpose |
+|----------------------|-----------------------------------|---------|
+| `OTAI_CACHE_DIR`     | `~/.cache/otai`                   | Where the DuckDB catalog, the "latest release" cache, and cached `croissant.json` files live. |
+| `OTAI_BASE_URI`      | the public Open Targets S3 bucket | Root the CLI reads parquet/`croissant.json` from. Only relevant for testing against local fixtures — never point this anywhere else in normal use. |
+| `OTAI_LOG_LEVEL`     | `INFO`                            | Verbosity of the CLI's stderr logging (progress/cache/retry messages). Set to `DEBUG` if you need more detail while diagnosing an issue; logging never touches stdout, so it's always safe to leave at the default. |
 
 ## JSON envelope
 
@@ -69,7 +85,7 @@ Every command emits one of:
 |------------------------|------------------------------------------------------|------------|
 | `guardrail_violation`  | Query isn't a single read-only SELECT/WITH           | Fix the SQL (e.g. remove the mutating/DDL statement) and retry |
 | `sql_error`            | SQL failed to parse, or failed at execution           | Fix the SQL syntax/logic and retry |
-| `timeout`              | Query ran past the execution time limit               | Narrow the query (add filters/LIMIT, reduce scope) and retry |
+| `timeout`              | Query ran past the execution time limit               | If the query is doing more work than the question needs, narrow it (add filters/LIMIT, reduce scope) and retry. If it's already minimal and legitimately slow (e.g. a full aggregate over a huge dataset), retry the *same* query with `--timeout <seconds>` instead |
 | `release_not_found`    | A schema-qualified release in the query is unknown    | Run `list-releases` to see valid release identifiers, then retry with a correct qualifier |
 | `dataset_not_found`    | `describe-dataset` name doesn't exist in that release | Run `list-datasets` for that release to find the correct name |
 | `s3_error`             | Couldn't list/reach the S3 bucket                     | Report the failure to the user; retrying immediately is unlikely to help |
@@ -91,9 +107,10 @@ Every command emits one of:
    specific past release or spans more than one release; leave table names
    unqualified when the question is about the latest release.
 5. **On a `run-sql` error, branch on `error.type`** per the table above —
-   in short: `timeout` → narrow and retry; `sql_error` /
-   `guardrail_violation` → fix the SQL; `release_not_found` → check
-   `list-releases` before retrying.
+   in short: `timeout` → narrow and retry, or retry with `--timeout
+   <seconds>` if the query is already minimal and legitimately slow;
+   `sql_error` / `guardrail_violation` → fix the SQL; `release_not_found`
+   → check `list-releases` before retrying.
 6. **Cite your sources in the final answer**: state which release(s) were
    queried and show the actual SQL you executed, so the user can verify or
    rerun it.

README.md

@@ -35,14 +35,17 @@ The four commands:
 - `list-releases` — what releases exist on S3, which is `latest`, which are cached locally.
 - `list-datasets [--release X]` — the datasets (tables) available in a release.
 - `describe-dataset <name> [--release X]` — a dataset's columns, types, and relationships.
-- `run-sql "<query>"` — read-only SQL against the views, guarded by `sqlglot`-based
-  validation: rejects anything but a single `SELECT`/`WITH` (including mutations
-  nested in a CTE or subquery) and rejects table-valued functions like
-  `read_csv_auto`/`read_parquet` as a data source (only plain, optionally
-  schema-qualified table/view names are allowed — `run-sql` can only query the
-  release catalog, never arbitrary local/remote files), plus a ~1000-row cap
-  and a ~45s timeout. A proactive `EXPLAIN`-based complexity check is scoped
-  but not yet implemented (see [issues/07](issues/07-query-complexity-guard.md)).
+- `run-sql "<query>" [--timeout SECONDS]` — read-only SQL against the views,
+  guarded by `sqlglot`-based validation: rejects anything but a single
+  `SELECT`/`WITH` (including mutations nested in a CTE or subquery) and
+  rejects table-valued functions like `read_csv_auto`/`read_parquet` as a
+  data source (only plain, optionally schema-qualified table/view names are
+  allowed — `run-sql` can only query the release catalog, never arbitrary
+  local/remote files), plus a ~1000-row cap and a timeout (default ~45s,
+  overridable per call with `--timeout` for a query that's legitimately
+  slow rather than a mistake to fix). A proactive `EXPLAIN`-based complexity
+  check is scoped but not yet implemented (see
+  [issues/07](issues/07-query-complexity-guard.md)).
 
 Every command emits a JSON envelope (`{"ok": true, "data": {...}}` /
 `{"ok": false, "error": {"type": "...", "message": "..."}}`) by default, or
@@ -51,8 +54,21 @@ a human-readable table with `--format table`.
 Building a release's schema for the first time can take a while (each
 dataset resolves a glob against real S3) — a progress bar and log messages
 report on that, always on stderr so they never interfere with the JSON on
-stdout. Set `OTAI_LOG_LEVEL` (default `INFO`) to `DEBUG` for more detail or
-`WARNING` to quiet it down.
+stdout.
+
+## Configuration
+
+Everything below is optional — the defaults are correct for regular use.
+
+| Variable         | Default                           | Purpose |
+| ---------------- | ---------------------------------- | ------- |
+| `OTAI_CACHE_DIR` | `~/.cache/otai`                    | Where the shared DuckDB catalog, the "latest release" cache, and cached `croissant.json` files live. |
+| `OTAI_BASE_URI`  | the public Open Targets S3 bucket  | Root the CLI reads parquet/`croissant.json` from. Tests point this at local fixtures; there's no reason to change it otherwise. |
+| `OTAI_LOG_LEVEL` | `INFO`                             | Verbosity of stderr logging (progress bars, cache hits/misses, catalog lock retries). `DEBUG` for more detail, `WARNING` to quiet it down — logging never touches stdout, so it's always safe to change. |
+
+`run-sql` also takes a `--timeout <seconds>` CLI flag (see above) to
+override the default timeout for one call, rather than an env var, since
+it's a per-query decision rather than a standing configuration choice.
 
 ## Requirements
 
@@ -93,6 +109,7 @@ uvx --from . otai list-releases
 uvx --from . otai list-datasets [--release 26.03]
 uvx --from . otai describe-dataset target [--release 26.03]
 uvx --from . otai run-sql "SELECT count(*) FROM target"
+uvx --from . otai run-sql "SELECT count(*) FROM colocalisation" --timeout 90
 ```
 
 Add `--format table` to any command for human-readable output; the default

src/otai/cli.py

@@ -10,7 +10,7 @@
 
 import typer
 
-from otai import commands, config, croissant, formatting
+from otai import commands, config, croissant, formatting, sql_guard
 from otai import releases as releases_mod
 from otai.logging_setup import configure_logging
 
@@ -44,6 +44,12 @@ def _validate_format(value: str) -> str:
     return value
 
 
+def _validate_timeout(value: float | None) -> float | None:
+    if value is not None and value <= 0:
+        raise typer.BadParameter("--timeout must be a positive number of seconds.")
+    return value
+
+
 def _emit_error(result: dict, output_format: str) -> None:
     error = result["error"]
     if output_format == "table":
@@ -146,6 +152,14 @@ def describe_dataset_cmd(
 @app.command("run-sql")
 def run_sql_cmd(
     query: str = typer.Argument(..., help="Read-only SQL query to execute."),
+    timeout: float | None = typer.Option(
+        None,
+        "--timeout",
+        callback=_validate_timeout,
+        help="Override the query timeout, in seconds "
+        f"(default: {sql_guard.DEFAULT_TIMEOUT_SECONDS:g}). Use for a call "
+        "known to need more time than the default allows.",
+    ),
     format: str = typer.Option(
         "json",
         "--format",
@@ -166,6 +180,9 @@ def run_sql_cmd(
         fetch_xml=releases_mod.default_fetch_listing_xml,
         fetch_croissant=croissant.default_fetch_croissant,
         base_uri=base_uri,
+        timeout_seconds=(
+            timeout if timeout is not None else sql_guard.DEFAULT_TIMEOUT_SECONDS
+        ),
     )
     _emit(
         result,

tests/test_cli.py

@@ -4,6 +4,7 @@
 from test_croissant import CROISSANT_FIXTURE
 from typer.testing import CliRunner
 
+from otai import catalog, sql_guard
 from otai.cli import app
 
 runner = CliRunner()
@@ -306,3 +307,80 @@ def test_run_sql_malformed_query_returns_sql_error(tmp_path, fixture_release_lay
     payload = json.loads(result.stdout)
     assert payload["ok"] is False
     assert payload["error"]["type"] == "sql_error"
+
+
+def test_run_sql_rejects_non_positive_timeout(tmp_path, fixture_release_layout):
+    base_uri, _release, _dataset_rows = fixture_release_layout
+
+    result = _invoke_with_fixtures(
+        ["run-sql", "SELECT 1", "--timeout", "0"], tmp_path / "cache", base_uri
+    )
+
+    # Match this file's convention for invalid-option tests (e.g.
+    # test_unknown_format_is_rejected): check the exit code only, not the
+    # styled error text - Typer/Click's box-drawn error panel wraps
+    # differently depending on the terminal width Rich detects, so exact
+    # substring matches against result.output are environment-fragile.
+    assert result.exit_code == 2
+
+
+def test_run_sql_passes_custom_timeout_through_to_commands(
+    tmp_path, fixture_release_layout
+):
+    base_uri, _release, _dataset_rows = fixture_release_layout
+
+    with patch("otai.cli.commands.run_sql") as mock_run_sql:
+        mock_run_sql.return_value = {"ok": True, "data": {"columns": [], "rows": []}}
+        _invoke_with_fixtures(
+            ["run-sql", "SELECT 1", "--timeout", "120"], tmp_path / "cache", base_uri
+        )
+
+    assert mock_run_sql.call_args.kwargs["timeout_seconds"] == 120.0
+
+
+def test_run_sql_uses_default_timeout_when_flag_omitted(
+    tmp_path, fixture_release_layout
+):
+    base_uri, _release, _dataset_rows = fixture_release_layout
+
+    with patch("otai.cli.commands.run_sql") as mock_run_sql:
+        mock_run_sql.return_value = {"ok": True, "data": {"columns": [], "rows": []}}
+        _invoke_with_fixtures(["run-sql", "SELECT 1"], tmp_path / "cache", base_uri)
+
+    assert (
+        mock_run_sql.call_args.kwargs["timeout_seconds"]
+        == sql_guard.DEFAULT_TIMEOUT_SECONDS
+    )
+
+
+def test_run_sql_timeout_flag_actually_shortens_execution(
+    tmp_path, fixture_release_layout
+):
+    # Real end-to-end confirmation (not mocked): a view wrapping range()
+    # (created outside the guard, same pattern used in test_commands.py/
+    # test_sql_guard.py) is cheap to set up but genuinely slow to query -
+    # a short --timeout must cause it to time out for real.
+    base_uri, release, _dataset_rows = fixture_release_layout
+    cache_dir = tmp_path / "cache"
+
+    _invoke_with_fixtures(["list-datasets"], cache_dir, base_uri)
+    conn = catalog.connect_catalog(cache_dir)
+    try:
+        conn.execute(
+            f'CREATE VIEW "{release}".slow_a AS SELECT * FROM range(100000000)'
+        )
+        conn.execute(f'CREATE VIEW "{release}".slow_b AS SELECT * FROM range(100000)')
+    finally:
+        conn.close()
+
+    result = _invoke_with_fixtures(
+        ["run-sql", "SELECT count(*) FROM slow_a a, slow_b b", "--timeout", "0.2"],
+        cache_dir,
+        base_uri,
+    )
+
+    assert result.exit_code != 0
+    payload = json.loads(result.stdout)
+    assert payload["ok"] is False
+    assert payload["error"]["type"] == "timeout"
+    assert "0.2" in payload["error"]["message"]