diff --git a/.claude/skills/otai/SKILL.md b/.claude/skills/otai/SKILL.md index fddd25d..e82af5f 100644 --- a/.claude/skills/otai/SKILL.md +++ b/.claude/skills/otai/SKILL.md @@ -42,15 +42,31 @@ Positional `` (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 ""` +### `otai run-sql "" [--timeout SECONDS]` Positional ``, 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 ` 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 + ` 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. diff --git a/README.md b/README.md index 1c72d48..f8f326f 100644 --- a/README.md +++ b/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 [--release X]` — a dataset's columns, types, and relationships. -- `run-sql ""` — 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 "" [--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 ` 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 diff --git a/src/otai/cli.py b/src/otai/cli.py index 0f06afe..2064c5a 100644 --- a/src/otai/cli.py +++ b/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, diff --git a/tests/test_cli.py b/tests/test_cli.py index bb842a8..4dc3a48 100644 --- a/tests/test_cli.py +++ b/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"]