Experimental postgres query (PR 2/4): provisioned + JSON/CSV streaming + types

experimental/aitools/cmd/query.go

@@ -14,10 +14,9 @@ import (
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/experimental/aitools/lib/middlewares"
 	"github.com/databricks/cli/experimental/aitools/lib/session"
+	"github.com/databricks/cli/experimental/libs/sqlcli"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
-	"github.com/databricks/cli/libs/env"
-	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/spf13/cobra"
@@ -35,16 +34,6 @@ const (
 
 	// cancelTimeout is how long to wait for server-side cancellation.
 	cancelTimeout = 10 * time.Second
-
-	// staticTableThreshold is the maximum number of rows rendered as a static table.
-	// Beyond this, an interactive scrollable table is used.
-	staticTableThreshold = 30
-
-	// outputCSV is the csv output format, supported only by the query command.
-	outputCSV = "csv"
-
-	// envOutputFormat matches the env var name in cmd/root/io.go.
-	envOutputFormat = "DATABRICKS_OUTPUT_FORMAT"
 )
 
 type queryOutputMode int
@@ -55,8 +44,13 @@ const (
 	queryOutputModeInteractiveTable
 )
 
-func selectQueryOutputMode(outputType flags.Output, stdoutInteractive, promptSupported bool, rowCount int) queryOutputMode {
-	if outputType == flags.OutputJSON {
+// selectQueryOutputMode picks the rendering mode for a single-query result.
+// JSON is the only machine-readable option; static and interactive are
+// table variants chosen by row count and TTY capabilities. Sharing only
+// the threshold with sqlcli; the three-way decision is aitools-specific
+// because the postgres command's renderers have a different shape.
+func selectQueryOutputMode(format sqlcli.Format, stdoutInteractive, promptSupported bool, rowCount int) queryOutputMode {
+	if format == sqlcli.OutputJSON {
 		return queryOutputModeJSON
 	}
 	if !stdoutInteractive {
@@ -67,7 +61,7 @@ func selectQueryOutputMode(outputType flags.Output, stdoutInteractive, promptSup
 	if !promptSupported {
 		return queryOutputModeStaticTable
 	}
-	if rowCount <= staticTableThreshold {
+	if rowCount <= sqlcli.StaticTableThreshold {
 		return queryOutputModeStaticTable
 	}
 	return queryOutputModeInteractiveTable
@@ -119,24 +113,15 @@ interactive table browser. Use --output csv to export results as CSV.`,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := cmd.Context()
 
-			// Normalize case to match root --output behavior (flags.Output.Set lowercases).
-			outputFormat = strings.ToLower(outputFormat)
-
-			// If --output wasn't explicitly passed, check the env var.
-			// Invalid env values are silently ignored, matching cmd/root/io.go.
-			if !cmd.Flag("output").Changed {
-				if v, ok := env.Lookup(ctx, envOutputFormat); ok {
-					switch flags.Output(strings.ToLower(v)) {
-					case flags.OutputText, flags.OutputJSON, outputCSV:
-						outputFormat = strings.ToLower(v)
-					}
-				}
-			}
-
-			switch flags.Output(outputFormat) {
-			case flags.OutputText, flags.OutputJSON, outputCSV:
-			default:
-				return fmt.Errorf("unsupported output format %q, accepted values: text, json, csv", outputFormat)
+			// Resolve the effective format via sqlcli so the env-var
+			// precedence and explicit-text-on-pipe handling stays in sync
+			// across commands. We pass stdoutTTY=true to keep the original
+			// aitools behavior of not auto-falling-back to JSON here; the
+			// per-result render mode further down already handles the pipe
+			// case via selectQueryOutputMode.
+			format, err := sqlcli.ResolveFormat(ctx, outputFormat, cmd.Flag("output").Changed, true)
+			if err != nil {
+				return err
 			}
 
 			sqls, err := resolveSQLs(ctx, cmd, args, filePaths)
@@ -146,7 +131,7 @@ interactive table browser. Use --output csv to export results as CSV.`,
 
 			// Reject incompatible flag combinations before any API call so the
 			// user sees the real error instead of an auth/warehouse failure.
-			if len(sqls) > 1 && flags.Output(outputFormat) != flags.OutputJSON {
+			if len(sqls) > 1 && format != sqlcli.OutputJSON {
 				return fmt.Errorf("multiple queries require --output json (got %q); pass --output json to receive a JSON array of per-statement results", outputFormat)
 			}
 
@@ -173,7 +158,7 @@ interactive table browser. Use --output csv to export results as CSV.`,
 			}
 
 			// CSV bypasses the normal output mode selection.
-			if flags.Output(outputFormat) == outputCSV {
+			if format == sqlcli.OutputCSV {
 				if len(columns) == 0 && len(rows) == 0 {
 					return nil
 				}
@@ -190,7 +175,7 @@ interactive table browser. Use --output csv to export results as CSV.`,
 			stdoutInteractive := cmdio.SupportsColor(ctx, cmd.OutOrStdout())
 			promptSupported := cmdio.IsPromptSupported(ctx)
 
-			switch selectQueryOutputMode(flags.Output(outputFormat), stdoutInteractive, promptSupported, len(rows)) {
+			switch selectQueryOutputMode(format, stdoutInteractive, promptSupported, len(rows)) {
 			case queryOutputModeJSON:
 				return renderJSON(cmd.OutOrStdout(), columns, rows)
 			case queryOutputModeStaticTable:
@@ -206,9 +191,13 @@ interactive table browser. Use --output csv to export results as CSV.`,
 	cmd.Flags().IntVar(&concurrency, "concurrency", defaultBatchConcurrency, "Maximum in-flight statements when running a batch of queries")
 	// Local --output flag shadows the root command's persistent --output flag,
 	// adding csv support for this command only.
-	cmd.Flags().StringVarP(&outputFormat, "output", "o", string(flags.OutputText), "Output format: text, json, or csv")
+	cmd.Flags().StringVarP(&outputFormat, "output", "o", string(sqlcli.OutputText), "Output format: text, json, or csv")
 	cmd.RegisterFlagCompletionFunc("output", func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) {
-		return []string{string(flags.OutputText), string(flags.OutputJSON), string(outputCSV)}, cobra.ShellCompDirectiveNoFileComp
+		out := make([]string, len(sqlcli.AllFormats))
+		for i, f := range sqlcli.AllFormats {
+			out[i] = string(f)
+		}
+		return out, cobra.ShellCompDirectiveNoFileComp
 	})
 
 	return cmd

experimental/aitools/cmd/query_test.go

@@ -10,9 +10,9 @@ import (
 	"time"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/experimental/libs/sqlcli"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/env"
-	"github.com/databricks/cli/libs/flags"
 	mocksql "github.com/databricks/databricks-sdk-go/experimental/mocks/service/sql"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/spf13/cobra"
@@ -271,57 +271,57 @@ func TestResolveWarehouseIDWithFlag(t *testing.T) {
 func TestSelectQueryOutputMode(t *testing.T) {
 	tests := []struct {
 		name              string
-		outputType        flags.Output
+		format            sqlcli.Format
 		stdoutInteractive bool
 		promptSupported   bool
 		rowCount          int
 		want              queryOutputMode
 	}{
 		{
 			name:              "json flag always returns json",
-			outputType:        flags.OutputJSON,
+			format:            sqlcli.OutputJSON,
 			stdoutInteractive: true,
 			promptSupported:   true,
 			rowCount:          999,
 			want:              queryOutputModeJSON,
 		},
 		{
 			name:              "non interactive stdout returns json",
-			outputType:        flags.OutputText,
+			format:            sqlcli.OutputText,
 			stdoutInteractive: false,
 			promptSupported:   true,
 			rowCount:          5,
 			want:              queryOutputModeJSON,
 		},
 		{
 			name:              "missing stdin interactivity falls back to static table",
-			outputType:        flags.OutputText,
+			format:            sqlcli.OutputText,
 			stdoutInteractive: true,
 			promptSupported:   false,
-			rowCount:          staticTableThreshold + 10,
+			rowCount:          sqlcli.StaticTableThreshold + 10,
 			want:              queryOutputModeStaticTable,
 		},
 		{
 			name:              "small results use static table",
-			outputType:        flags.OutputText,
+			format:            sqlcli.OutputText,
 			stdoutInteractive: true,
 			promptSupported:   true,
-			rowCount:          staticTableThreshold,
+			rowCount:          sqlcli.StaticTableThreshold,
 			want:              queryOutputModeStaticTable,
 		},
 		{
 			name:              "large results use interactive table",
-			outputType:        flags.OutputText,
+			format:            sqlcli.OutputText,
 			stdoutInteractive: true,
 			promptSupported:   true,
-			rowCount:          staticTableThreshold + 1,
+			rowCount:          sqlcli.StaticTableThreshold + 1,
 			want:              queryOutputModeInteractiveTable,
 		},
 	}
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			got := selectQueryOutputMode(tc.outputType, tc.stdoutInteractive, tc.promptSupported, tc.rowCount)
+			got := selectQueryOutputMode(tc.format, tc.stdoutInteractive, tc.promptSupported, tc.rowCount)
 			assert.Equal(t, tc.want, got)
 		})
 	}

experimental/libs/sqlcli/output.go

@@ -0,0 +1,93 @@
+// Package sqlcli holds patterns shared by experimental SQL-running commands
+// (currently `experimental aitools tools query` and `experimental postgres
+// query`). The package lives under experimental/libs/ rather than libs/ so
+// the commands depending on it inherit experimental-stability guarantees:
+// when both consumers graduate, this package can be promoted alongside
+// (or its API stabilised first).
+package sqlcli
+
+import (
+	"context"
+	"fmt"
+	"slices"
+	"strings"
+
+	"github.com/databricks/cli/libs/env"
+)
+
+// EnvOutputFormat matches the env var name in cmd/root/io.go.
+// Reading it lets pipelines set DATABRICKS_OUTPUT_FORMAT once for all
+// commands.
+const EnvOutputFormat = "DATABRICKS_OUTPUT_FORMAT"
+
+// StaticTableThreshold is the row count above which interactive callers may
+// hand off to libs/tableview's scrollable viewer. Smaller results stay in a
+// static tabwriter table so they pipe to scripts unchanged.
+const StaticTableThreshold = 30
+
+// Format is the user-selectable output shape. Using a string typedef instead
+// of an int enum keeps the help text and DATABRICKS_OUTPUT_FORMAT env var
+// values self-describing.
+type Format string
+
+const (
+	OutputText Format = "text"
+	OutputJSON Format = "json"
+	OutputCSV  Format = "csv"
+)
+
+// AllFormats is the canonical order shown in completions / help. Sharing
+// the slice avoids drift between consumers when a new format is added.
+var AllFormats = []Format{OutputText, OutputJSON, OutputCSV}
+
+// ResolveFormat picks the effective output format. Precedence:
+//
+//  1. The local --output flag if it was explicitly set.
+//  2. DATABRICKS_OUTPUT_FORMAT env var if set to a known value (invalid
+//     values are silently ignored, matching cmd/root/io.go and aitools).
+//  3. The flag default (whatever the caller passes as flagValue).
+//
+// Then the auto-selection rule applies: a *defaulted* text mode on a non-TTY
+// stdout falls back to JSON, so scripts piping the output get machine-
+// readable output by default. An *explicit* --output text (flag or env) is
+// honoured even on a pipe; per AGENTS.md we don't silently override flags
+// the user set.
+//
+// flagSet is true if the user explicitly passed --output on the CLI.
+// stdoutTTY is true if stdout is a terminal.
+func ResolveFormat(ctx context.Context, flagValue string, flagSet, stdoutTTY bool) (Format, error) {
+	chosen := Format(strings.ToLower(flagValue))
+	chosenExplicit := flagSet
+
+	if !flagSet {
+		if v, ok := env.Lookup(ctx, EnvOutputFormat); ok {
+			candidate := Format(strings.ToLower(v))
+			if IsKnown(candidate) {
+				chosen = candidate
+				chosenExplicit = true
+			}
+		}
+	}
+
+	if !IsKnown(chosen) {
+		return "", fmt.Errorf("unsupported output format %q; expected one of: %s", flagValue, joinFormats(AllFormats))
+	}
+
+	if chosen == OutputText && !stdoutTTY && !chosenExplicit {
+		return OutputJSON, nil
+	}
+	return chosen, nil
+}
+
+// IsKnown reports whether f is one of the formats in AllFormats.
+func IsKnown(f Format) bool {
+	return slices.Contains(AllFormats, f)
+}
+
+func joinFormats(formats []Format) string {
+	parts := make([]string, len(formats))
+	for i, f := range formats {
+		parts[i] = string(f)
+	}
+	return strings.Join(parts, ", ")
+}