Proposal: Add databricks query command (Statement API) with a profile‑level safety gate; vector‑ready design

Bắt đầu bằng cách xem xét các điểm vào lệnh được đề xuất trong cmd/query, cmd/query/sql.go và cách chúng được kết nối trong cmd/cmd.go. Sau đó, kiểm tra phạm vi unit test được đề xuất cho libs/safety/sqlsafe và các yêu cầu của Statement Execution API. Công việc được xem là hoàn tất khi phạm vi v1 đã được chấp nhận hoạt động với các flags được liệt kê, tính an toàn và thứ tự ưu tiên cấu hình đã được kiểm thử, các luồng mock API đã được bao phủ và tài liệu CLI đã được thêm.

Summary

Introduce a new top‑level databricks query command for running queries from the CLI. The first subcommand, databricks query sql, would execute SQL via the Statement Execution API on a SQL Warehouse. A profile‑level safety gate enforces read‑only behavior by default so both humans and AI tools can use the command safely. The structure intentionally leaves space for databricks query vector (future) to execute vector/search queries.

---

Problem & Motivation

• Today there is no CLI command that submits a SQL statement and returns results. The databricks psql command is interactive JDBC; it’s not great for automation, scripting, or AI agents that need a one‑shot query with structured output.
• Teams often want CI jobs, notebooks, or AI dev tools to run read‑only queries safely without risking accidental DROP/ALTER/INSERT/DELETE.
• We also want a coherent home for non‑SQL queries (e.g., vector search) that return records but don’t fit into the “workspace admin” groups under workspace sql-*.

Goals

1. Provide a simple, scriptable command to run queries and get results back as table/JSON/CSV (Arrow optional later).
2. Make it safe by default for people and AI: a profile flag controls whether destructive SQL is allowed; default is read‑only only.
3. Establish a query namespace that can grow: query sql now, query vector later.

Non‑Goals (v1)

• Replace the interactive SQL CLI or JDBC clients.
• Implement every Statement API option; start with a minimal, solid core.
• Add a full SQL parser; we’ll use a conservative token/keyword approach to detect non‑read‑only statements.

---

Proposed UX

Run an inline statement

databricks query sql \
  --warehouse-id <WAREHOUSE_ID> \
  --sql "select 1 as one"

Run from a file

databricks query sql \
  --warehouse-id <WAREHOUSE_ID> \
  --file queries/top10_customers.sql

Output control

# Render as a table to TTY (default)
# Write JSON/CSV to a file

databricks query sql --warehouse-id <ID> \
  --sql "select * from system.builtin.current_queries" \
  --format json --output out/current_queries.json

Wait behavior (hybrid default)

# Wait up to 10s; if still running, fall back to async and poll.
databricks query sql --warehouse-id <ID> --sql "select count(*) from t"

# Force async return immediately
# (caller can poll with statement id; may be out of scope for v1 CLI)
# databricks query sql --wait-timeout 0s ...

Safety gate (profile/env/flag)

• Default: read‑only only. Destructive statements are blocked with a clear error.
• Opt‑in overrides:
  • CLI flag: --allow-destructive
  • Env var: DATABRICKS_CLI_ALLOW_DESTRUCTIVE_SQL=true
  • Profile key: cli.allow_destructive_sql = true

Example profile (in ~/.databrickscfg):

[dev]
host = https://dbc-xxx.cloud.databricks.com
token = dapi***********
cli.allow_destructive_sql = false  # default; safe for AI and scripts

Future: databricks query vector ... to issue vector/search queries to Databricks vector services, returning JSON/CSV similarly.

---

Why a new top‑level query command?

• Keeps execution distinct from workspace admin commands (query history, visualizations, permissions). Those are CRUD‑style operations; this proposal is about running queries.
• Avoids overloading psql (interactive) and gives a clean automation entrypoint for humans and AI.
• Provides a natural home for multiple modalities (SQL today, vector tomorrow) under one user mental model.

---

Implementation Plan (high level)

1) Command structure

• New package: cmd/query with root command query.
• Subcommand: cmd/query/sql.go exporting New() that wires into cmd/cmd.go.

2) Flags (v1)

• --warehouse-id <string> (required; no implicit default in v1)
• One of:
  • --sql "..." (inline statement)
  • --file <path.sql> (read file contents; allow multiple statements separated by ;)
• Output/format:
  • --format table|json|csv (default: table for TTY, json if --output is provided)
  • --output <path> (writes file; otherwise print)
• Wait & control:
  • --wait-timeout <duration> (default 10s hybrid: wait, then poll)
  • (Optional later) --on-wait-timeout CONTINUE|CANCEL (default CONTINUE)
• Safety:
  • --allow-destructive (override)
  • As noted: env and profile keys also supported
• (Stretch) --dry-run & --explain-safety to print parsed statements and classification results.

3) Statement Execution & polling

• Use Go SDK’s Statement Execution API client to ExecuteStatement with warehouse id, SQL text, wait timeout.
• If result returns inline: render immediately.
• If not ready: poll GetStatement until success/error; if chunked results, stream via GetStatementResultChunkN.
• Renderers: table (TTY), JSON/CSV (file or stdout).

4) Read‑only safety classifier (minimal, conservative)

• Strip -- and /* ... */ comments while respecting '" and backtick quotes.
• Split on semicolons that are not inside quotes.
• For each statement, check first meaningful token (case‑insensitive):
  • Allow: SELECT, SHOW, DESCRIBE, EXPLAIN (only if target is read‑only), WITH leading to SELECT.
  • Block: everything else (INSERT, UPDATE, DELETE, MERGE, CREATE, ALTER, DROP, TRUNCATE, REPLACE, VACUUM, OPTIMIZE, COPY, GRANT, REVOKE, RESTORE, MSCK, SET [treat as stateful; block in v1]).
• If any statement in a multi‑statement input is non‑read‑only, fail fast with a helpful error:

  "Blocked by safe‑only mode: found DDL/DML keyword CREATE at statement 1 (pos 1:1). Pass --allow-destructive or set cli.allow_destructive_sql=true to override."

• Note the classifier is intentionally conservative; we’ll document that user‑defined or engine‑side side effects cannot be perfectly detected statically.

Placement: implement classifier as libs/safety/sqlsafe with unit tests.

5) Config resolution order

1. CLI flag --allow-destructive
2. Env var DATABRICKS_CLI_ALLOW_DESTRUCTIVE_SQL
3. Profile key cli.allow_destructive_sql
4. Default: false (safe‑only)

6) Error handling & UX

• Map Statement API terminal states to clear CLI errors with exit codes.
• Spinner/progress while polling; --log-level debug should include request ids for supportability.

7) Tests & acceptance

• Unit: classifier edge cases (comments, quotes, multi‑stmt, EXPLAIN SELECT vs EXPLAIN INSERT, etc.).
• Integration (optional): behind a guard/environment, run a tiny safe query against a test warehouse.
• Acceptance: mock server for Statement API covering inline/hybrid/chunked paths.

8) Docs

• New page under CLI docs: usage, safety gate, examples.
• Mention how this differs from psql and from the Databricks SQL CLI.

---

Backwards Compatibility & Impact

• Purely additive: no changes to existing command groups.
• Defaults emphasize safety and deterministic output for automation.
• Future query vector will reuse output/UI and safety philosophy (read‑only by nature; separate rate/size limits to discuss later).

---

Open Questions for Maintainers

1. Placement & naming
   • Approve a new root query command with sql subcommand?
   • Alternate: databricks sql run under a new sql group? (My preference: query to keep modality‑agnostic.)
2. Default output
   • Table to TTY, JSON to file by default—ok?
3. Safety defaults & scope
   • Is SET considered destructive (block in v1)? Any other keywords to add to the blocklist/allowlist?
4. Warehouse resolution
   • v1 requires --warehouse-id. Should we look up a default from profile in a later PR?
5. Arrow support
   • Should we include --format arrow in v1 or leave to a follow‑up?
6. Async surface
   • Expose statement id and a query sql get --statement-id in v1, or leave until there’s demand?
7. CLA/process
   • Is a CLA required and are you open to a community PR for this feature?

---

Willing to Implement

I’m happy to implement the v1 scope outlined above (command, Statement API integration, conservative safety classifier, table/JSON/CSV output, tests, docs). I’ll incorporate your feedback on naming, defaults, and any repo‑specific conventions. If a CLA or additional steps are needed, please let me know.

---

Acceptance Criteria (proposed)

• databricks query sql executes a SQL statement on a warehouse and returns results.
• Read‑only safety gate enforced by default; destructive statements blocked with a clear error message.
• Flags: --warehouse-id, --sql/--file, --format, --output, --wait-timeout, --allow-destructive.
• Profile/env/flag precedence works as documented.
• Unit tests for classifier; basic acceptance tests for inline/hybrid/chunked flows.
• Docs added with examples and safety guidance.

---

Future Work (non‑blocking)

• databricks query vector for vector/search queries.
• --format arrow and streaming Arrow to stdout.
• Async surface: return statement id, get subcommands.
• Sane defaults for warehouse selection.
• Richer result rendering (paging, truncation, width control) if demand arises.

---

Request for feedback:

• Do you agree with the top‑level query command design and v1 scope?
• Are the safety defaults the right trade‑off (helpful for AI and automation)?
• Any concerns with implementation approach or repo structure?