aitools: add 'tools statement' lifecycle commands

[simonfaltum]

Stack

This PR is part of a 4-PR stack making aitools data exploration faster for ai-dev-kit. Each PR is independently reviewable; merge in order.

1. aitools: extract pollStatement helper and pin OnWaitTimeout #5092 — aitools: extract pollStatement helper and pin OnWaitTimeout (base: main)
2. aitools: run multiple SQL queries in parallel from one invocation #5093 — aitools: run multiple SQL queries in parallel from one query invocation (base: aitools: extract pollStatement helper and pin OnWaitTimeout #5092)
3. aitools: add 'tools statement' lifecycle commands #5095 — aitools: add 'tools statement' lifecycle commands (base: aitools: run multiple SQL queries in parallel from one invocation #5093) — this PR
4. aitools: parallelize discover-schema across tables and probes #5097 — aitools: parallelize discover-schema across tables and probes (base: aitools: add 'tools statement' lifecycle commands #5095)

Use git diff <base>...HEAD or set the comparison base in the GitHub UI to see only this PR's changes; the default "Files changed" diff against main includes ancestor PRs.

---

Why

Quentin's ai-dev-kit skill works against synchronous tools query. That covers most cases, but there are workflows where the agent wants a server-side handle it can poll separately: long-running maintenance queries, parallel exploration where the agent does other work in between, and any "submit-now-harvest-later" pattern.

tools query with a single SQL is for "I want results now." This PR adds a low-level command tree, tools statement, for "I want a handle." Cleaner separation than overloading query with --async/--cancel flags (which would be semantically forced — a query shouldn't manage someone else's statement_id).

Changes

Four new subcommands under databricks experimental aitools tools statement:

# Fire and exit with a handle.
databricks experimental aitools tools statement submit \
  --warehouse <wh> "SELECT pg_sleep(60)"

# Output:
# { "statement_id": "01ef...", "state": "PENDING", "warehouse_id": "..." }

# Block until terminal and emit rows.
databricks experimental aitools tools statement get <statement_id>

# Peek at current state without polling.
databricks experimental aitools tools statement status <statement_id>

# Request cancellation.
databricks experimental aitools tools statement cancel <statement_id>

Implementation notes:

• All four subcommands emit a uniform statementInfo JSON shape: {statement_id, state, warehouse_id, columns, rows, error} with omitempty on every field except statement_id. So submit doesn't include columns/rows, cancel doesn't include warehouse_id, etc. Consumer parsing is uniform.
• submit uses WaitTimeout: "0s" and OnWaitTimeout: CONTINUE (matching the helper from aitools: extract pollStatement helper and pin OnWaitTimeout #5092).
• get uses pollStatement (from aitools: extract pollStatement helper and pin OnWaitTimeout #5092) and inherits its "ctx cancellation does NOT cancel server-side" semantics. This is the important UX difference from tools query: hitting Ctrl+C on get stops polling but leaves the statement running on the warehouse. Use cancel for explicit termination. That asymmetry is intentional, since get is poll-only by design — the user already submitted async.
• status does a single GetStatementByStatementId with no polling.
• cancel calls CancelExecution and optimistically reports state=CANCELED. The Statements API returns no body on cancel; the actual server-side state transitions asynchronously. The Long help points users at status if they need certainty.
• A shared helper statementErrorFromStatus populates the error field for every non-success terminal state (FAILED, CANCELED, CLOSED), even when the server returns no Status.Error payload. So skill consumers can branch on error == null alone instead of inspecting state.
• Each subcommand has a small testable helper (submitStatement, getStatementResult, getStatementStatus, cancelStatementExecution) extracted from the cobra RunE. Tests target the helpers directly with a mock StatementExecutionInterface.
• Parent statement.go registers the four subcommands and is wired into tools.go next to query, discover-schema, and get-default-warehouse.
• submit validates input (rejects mixed --file + positional) BEFORE accessing WorkspaceClient, so the error surfaces cleanly without an auth or warehouse roundtrip.

Test plan

• go test ./experimental/aitools/... passes.

• make checks clean.

• make fmt no drift.

• make lint 0 issues.

• New tests cover:

  • submit returns the statement_id and pins OnWaitTimeout: CONTINUE
  • submit wraps transport errors with execute statement: ...
  • get polls until terminal and assembles rows
  • get reports server-side errors in the JSON without raising a Go error
  • get ctx cancellation propagates without calling CancelExecution (the deliberate UX difference from query)
  • get synthesizes error for terminal CLOSED / FAILED with no backend payload
  • status does a single GET, no polling
  • status reports server-side errors in the JSON; running/pending stay error-free
  • status synthesizes error for FAILED with no backend payload
  • cancel calls CancelExecution and reports state=CANCELED
  • cancel wraps API errors
  • statementErrorFromStatus table-driven across nil, succeeded, running, failed-with-error, failed/canceled/closed-without-error
  • renderStatementInfo JSON shape (full and minimal)
  • cobra-level: submit rejects mixed --file + positional, submit enforces MaximumNArgs(1), get and cancel require a positional statement_id

• Manual smoke against a real warehouse:

  SID=$(databricks experimental aitools tools statement submit \
    --warehouse <wh> "SELECT pg_sleep(5)" | jq -r '.statement_id')
  databricks experimental aitools tools statement status "$SID"
  databricks experimental aitools tools statement get "$SID"

[arsenyinfo]

getStatementResult discards partial JSON when fetchAllRows fails after a SUCCEEDED state

• Priority: P2
• Location: experimental/aitools/cmd/statement_get.go:78-85 and statement_get.go:32-35
• Scenario: Statement reaches SUCCEEDED and getStatementResult populates info.StatementID, info.State, and info.Columns, but a subsequent GetStatementResultChunkN call inside fetchAllRows fails (network blip, throttling, transient 5xx). Line 82 returns (info, err); RunE at lines 33-34 discards info and surfaces only the raw Go error. The user sees fetch result chunk N: ... with no JSON, no statement_id, and no machine-readable error field — inconsistent with the terminal-failure path at lines 41-44 which renders JSON and returns root.ErrAlreadyPrinted.
• Potential solution: Render the partially-populated statementInfo with a synthesized Error describing the chunk-fetch failure, then return root.ErrAlreadyPrinted, mirroring the terminal-state failure path just above in the same function.

submit input validation runs after WorkspaceClient setup, contradicting the PR description

• Priority: P3
• Location: experimental/aitools/cmd/statement_submit.go:29-44
• Scenario: User runs statement submit --file foo.sql "select 2" (or statement submit with no source) without configured auth. PreRunE: root.MustWorkspaceClient at line 30 runs before RunE, so auth/workspace setup fails first; resolveSQLs (line 38) and the len(sqls) != 1 check (lines 42-43) never execute. The PR description's claim that submit validates input before accessing WorkspaceClient is not delivered.
• Potential solution: Move resolveSQLs and the length check into a custom PreRunE that runs before root.MustWorkspaceClient, mirroring the pattern already used in query.go:113-118, or correct the PR description to remove the validation-order claim.

---

🔍 Reviewed by nitpicker