aitools: run multiple SQL queries in parallel from one invocation

[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) — this PR
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)
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

Today databricks experimental aitools tools query runs one SQL at a time. ai-dev-kit's data-exploration phase fires 5-10 probes per dashboard (cardinality, top values, distributions, trend viability) and they all run in series because each is a separate CLI invocation that blocks. End-to-end exploration takes about a minute when it could take seconds.

Quentin already wired up a bash workaround that fans out via the raw /api/2.0/sql/statements endpoint with wait_timeout=0s and harvests results separately. This PR exposes that pattern natively so the skill can drop the hack and other CLI users get the same speed-up.

Changes

Before: query accepted at most one positional SQL or a single --file. Mixing the two errored. JSON output was an array of row objects.

Now: query accepts any number of positional SQLs and/or repeated --file paths. With one input, behavior is unchanged (back-compat). With two or more, the queries run in parallel against the warehouse and the result is a JSON array of one object per input in input order:

[
  {
    "sql": "SELECT count(*) FROM t",
    "statement_id": "01ef...",
    "state": "SUCCEEDED",
    "elapsed_ms": 412,
    "columns": ["count"],
    "rows": [["12345"]]
  },
  {
    "sql": "SELECT bad_syntax",
    "statement_id": "01ef...",
    "state": "FAILED",
    "elapsed_ms": 87,
    "error": {
      "message": "near 'bad_syntax': syntax error",
      "error_code": "SYNTAX_ERROR"
    }
  }
]

Implementation:

• New experimental/aitools/cmd/batch.go with executeBatch (errgroup with bounded parallelism) and runOneBatchQuery. Each goroutine submits with OnWaitTimeout: CONTINUE, polls via the helper from aitools: extract pollStatement helper and pin OnWaitTimeout #5092, and encodes its outcome into a batchResult struct. Failures don't abort siblings.
• New --concurrency flag (default 8). Same value used by cmd/fs/cp.go for similar fan-out. Validated > 0 in PreRunE (a 0 value would deadlock errgroup.SetLimit).
• --file is now a repeatable string slice. Previous --file + positional conflict error is removed; both compose.
• resolveSQL is replaced by resolveSQLs returning []string. Result order is --file inputs first (in flag order), then positional SQLs (in arg order).
• Multi-query output is JSON-only. --output text and --output csv are rejected with an actionable error before any API call.
• On Ctrl+C, in-flight statements are cancelled server-side via CancelExecution after g.Wait() returns. Statements that finished normally before the cancel are left alone.
• Exit code is non-zero (root.ErrAlreadyPrinted) when any statement failed; the JSON already contains the error detail, no extra stderr noise.

Test plan

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

• make checks clean.

• make fmt no drift.

• make lint 0 issues.

• New unit tests cover:

  • all-succeed batch with input-order preservation
  • server-reported failure on one of N (others still complete)
  • submission-time transport error encoded into per-result error
  • explicit OnWaitTimeout: CONTINUE on every ExecuteStatement
  • staggered completion (1 slow + 2 fast) preserves input order in results
  • context cancellation triggers CancelExecution for each in-flight statement
  • cobra-level rejection of --output text and --output csv with multiple positionals
  • cobra-level rejection of --concurrency 0 and --concurrency -1
  • resolveSQLs covering mixed sources, multiple files, multiple positionals, indexed-error message

• Manual smoke against a real warehouse:

  databricks experimental aitools tools query \
    --warehouse <wh> --output json \
    "SELECT 1" "SELECT 2" "SELECT current_timestamp()"

[arsenyinfo]

Cancel-RPC contexts derived from an already-cancelled parent, causing server-side cancel to be a no-op

• Priority: P2
• Location: experimental/aitools/cmd/batch.go:201-202, experimental/aitools/cmd/query.go:347-355
• Scenario: When the parent ctx passed to executeBatch (or executeAndPoll) is itself cancelled — as exercised by TestExecuteBatchContextCancellationCancelsInFlight, which cancels the parent context before invocation — context.WithTimeout(ctx, cancelTimeout) in cancelInFlight and cancelStatement inherits the already-done parent, so the derived context is immediately cancelled. api.CancelExecution short-circuits with ctx.Err() and no cancel RPC reaches the warehouse, leaving in-flight statements running server-side. The test only passes today because the mock accepts mock.Anything for the context argument and never inspects liveness.
• Potential solution: Detach the cancel timeout from any inbound cancellation by using context.WithTimeout(context.WithoutCancel(ctx), cancelTimeout) (Go 1.21+) in both cancelInFlight (batch.go:201) and cancelStatement (query.go:347); fall back to context.Background() if the minimum Go version is below 1.21.

---

🔍 Reviewed by nitpicker