Add CSV output format for SQL query results

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
 }

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())
+}

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)

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())
+}