diff --git a/experimental/aitools/cmd/query.go b/experimental/aitools/cmd/query.go index dd5b97c761..2c7cb51c21 100644 --- a/experimental/aitools/cmd/query.go +++ b/experimental/aitools/cmd/query.go @@ -69,6 +69,7 @@ func selectQueryOutputMode(outputType flags.Output, stdoutInteractive, promptSup func newQueryCmd() *cobra.Command { var warehouseID string var filePath string + var format string cmd := &cobra.Command{ Use: "query [SQL | file.sql]", @@ -83,15 +84,25 @@ The command auto-detects an available warehouse unless --warehouse is set or the DATABRICKS_WAREHOUSE_ID environment variable is configured. Output is JSON in non-interactive contexts. In interactive terminals it renders -tables, and large results open an interactive table browser.`, +tables, and large results open an interactive table browser. Use --format csv +to export results as CSV.`, Example: ` databricks experimental aitools tools query "SELECT * FROM samples.nyctaxi.trips LIMIT 5" databricks experimental aitools tools query --warehouse abc123 "SELECT 1" databricks experimental aitools tools query --file report.sql databricks experimental aitools tools query report.sql + databricks experimental aitools tools query --format csv "SELECT * FROM samples.nyctaxi.trips LIMIT 5" echo "SELECT 1" | databricks experimental aitools tools query`, Args: cobra.MaximumNArgs(1), PreRunE: root.MustWorkspaceClient, RunE: func(cmd *cobra.Command, args []string) error { + if format != "" && format != "csv" { + return fmt.Errorf("unsupported format %q, supported values: csv", format) + } + + if format != "" && cmd.Flag("output").Changed { + return errors.New("cannot use --format and --output together; use --format csv or --output json/text") + } + ctx := cmd.Context() w := cmdctx.WorkspaceClient(ctx) @@ -116,6 +127,14 @@ tables, and large results open an interactive table browser.`, return err } + // CSV format bypasses the normal output mode selection. + if format == "csv" { + if len(columns) == 0 && len(rows) == 0 { + return nil + } + return renderCSV(cmd.OutOrStdout(), columns, rows) + } + if len(columns) == 0 && len(rows) == 0 { fmt.Fprintln(cmd.OutOrStdout(), "Query executed successfully (no results)") return nil @@ -139,6 +158,7 @@ tables, and large results open an interactive table browser.`, cmd.Flags().StringVarP(&warehouseID, "warehouse", "w", "", "SQL warehouse ID to use for execution") cmd.Flags().StringVarP(&filePath, "file", "f", "", "Path to a SQL file to execute") + cmd.Flags().StringVar(&format, "format", "", "Output format (supported: csv). When omitted, uses --output flag behavior.") return cmd } diff --git a/experimental/aitools/cmd/query_test.go b/experimental/aitools/cmd/query_test.go index 22b4fbaf13..2fe18473ba 100644 --- a/experimental/aitools/cmd/query_test.go +++ b/experimental/aitools/cmd/query_test.go @@ -447,3 +447,36 @@ func TestResolveSQLMissingFileReturnsError(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "read SQL file") } + +func TestQueryCommandUnsupportedFormatReturnsError(t *testing.T) { + cmd := newQueryCmd() + cmd.PreRunE = nil + cmd.SetArgs([]string{"--format", "xml", "SELECT 1"}) + err := cmd.Execute() + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported format") +} + +func TestQueryCommandFormatAndOutputConflictReturnsError(t *testing.T) { + cmd := newQueryCmd() + cmd.PreRunE = nil + cmd.PersistentFlags().String("output", "text", "output type") + cmd.SetArgs([]string{"--format", "csv", "--output", "json", "SELECT 1"}) + err := cmd.Execute() + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot use --format and --output together") +} + +func TestRenderCSVOutput(t *testing.T) { + var buf strings.Builder + err := renderCSV(&buf, []string{"id", "name"}, [][]string{{"1", "alice"}, {"2", "bob"}}) + require.NoError(t, err) + assert.Equal(t, "id,name\r\n1,alice\r\n2,bob\r\n", buf.String()) +} + +func TestRenderCSVHeadersOnlyWhenNoRows(t *testing.T) { + var buf strings.Builder + err := renderCSV(&buf, []string{"id", "name"}, nil) + require.NoError(t, err) + assert.Equal(t, "id,name\r\n", buf.String()) +} diff --git a/experimental/aitools/cmd/render.go b/experimental/aitools/cmd/render.go index b7eadb401c..09cffb3049 100644 --- a/experimental/aitools/cmd/render.go +++ b/experimental/aitools/cmd/render.go @@ -1,6 +1,7 @@ package mcp import ( + "encoding/csv" "encoding/json" "fmt" "io" @@ -51,6 +52,28 @@ func renderJSON(w io.Writer, columns []string, rows [][]string) error { return nil } +// renderCSV writes query results as CSV with column headers as the first row. +func renderCSV(w io.Writer, columns []string, rows [][]string) error { + cw := csv.NewWriter(w) + cw.UseCRLF = true + if err := cw.Write(columns); err != nil { + return fmt.Errorf("write CSV header: %w", err) + } + for _, row := range rows { + record := make([]string, len(columns)) + for i := range columns { + if i < len(row) { + record[i] = row[i] + } + } + if err := cw.Write(record); err != nil { + return fmt.Errorf("write CSV row: %w", err) + } + } + cw.Flush() + return cw.Error() +} + // renderStaticTable writes query results as a formatted text table. func renderStaticTable(w io.Writer, columns []string, rows [][]string) error { tw := tabwriter.NewWriter(w, 0, 4, 2, ' ', 0) diff --git a/experimental/aitools/cmd/render_test.go b/experimental/aitools/cmd/render_test.go index f559609749..25b8b2386c 100644 --- a/experimental/aitools/cmd/render_test.go +++ b/experimental/aitools/cmd/render_test.go @@ -93,3 +93,52 @@ func TestRenderStaticTableEmpty(t *testing.T) { assert.Contains(t, output, "id") assert.Contains(t, output, "0 rows") } + +func TestRenderCSVBasic(t *testing.T) { + var buf bytes.Buffer + columns := []string{"id", "name", "city"} + rows := [][]string{ + {"1", "Alice", "New York"}, + {"2", "Bob", "London"}, + } + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "id,name,city\r\n1,Alice,New York\r\n2,Bob,London\r\n", buf.String()) +} + +func TestRenderCSVSpecialCharacters(t *testing.T) { + var buf bytes.Buffer + columns := []string{"name", "description"} + rows := [][]string{ + {"Alice", "has a comma, here"}, + {"Bob", `has "quotes" here`}, + {"Carol", "has a\nnewline"}, + } + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "name,description\r\nAlice,\"has a comma, here\"\r\nBob,\"has \"\"quotes\"\" here\"\r\nCarol,\"has a\r\nnewline\"\r\n", buf.String()) +} + +func TestRenderCSVEmptyResultSet(t *testing.T) { + var buf bytes.Buffer + columns := []string{"id", "name"} + var rows [][]string + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "id,name\r\n", buf.String()) +} + +func TestRenderCSVShortRows(t *testing.T) { + var buf bytes.Buffer + columns := []string{"a", "b", "c"} + rows := [][]string{ + {"1"}, + } + + err := renderCSV(&buf, columns, rows) + require.NoError(t, err) + assert.Equal(t, "a,b,c\r\n1,,\r\n", buf.String()) +}