diff --git a/internal/command.go b/internal/command.go index 66472cd363..daf984bc33 100644 --- a/internal/command.go +++ b/internal/command.go @@ -40,6 +40,7 @@ import ( "github.com/confluentinc/cli/v4/internal/plugin" "github.com/confluentinc/cli/v4/internal/prompt" providerintegration "github.com/confluentinc/cli/v4/internal/provider-integration" + "github.com/confluentinc/cli/v4/internal/query" "github.com/confluentinc/cli/v4/internal/rtce" schemaregistry "github.com/confluentinc/cli/v4/internal/schema-registry" "github.com/confluentinc/cli/v4/internal/secret" @@ -132,6 +133,7 @@ func NewConfluentCommand(cfg *config.Config) *cobra.Command { cmd.AddCommand(plugin.New(cfg, prerunner)) cmd.AddCommand(prompt.New(cfg)) cmd.AddCommand(providerintegration.New(prerunner)) + cmd.AddCommand(query.New(cfg, prerunner)) cmd.AddCommand(rtce.New(cfg, prerunner)) cmd.AddCommand(schemaregistry.New(cfg, prerunner)) cmd.AddCommand(secret.New(prerunner, secrets.NewPasswordProtectionPlugin())) diff --git a/internal/query/command.go b/internal/query/command.go new file mode 100644 index 0000000000..5215b5f2dc --- /dev/null +++ b/internal/query/command.go @@ -0,0 +1,683 @@ +package query + +import ( + "context" + goerrors "errors" + "fmt" + "net/http" + "os" + "os/signal" + "syscall" + "time" + + "github.com/olekukonko/tablewriter" + "github.com/spf13/cobra" + + flinkgatewayv1 "github.com/confluentinc/ccloud-sdk-go-v2/flink-gateway/v1" + + "github.com/confluentinc/cli/v4/pkg/auth" + "github.com/confluentinc/cli/v4/pkg/ccloudv2" + pcmd "github.com/confluentinc/cli/v4/pkg/cmd" + cliconfig "github.com/confluentinc/cli/v4/pkg/config" + "github.com/confluentinc/cli/v4/pkg/errors" + flinkerror "github.com/confluentinc/cli/v4/pkg/errors/flink" + "github.com/confluentinc/cli/v4/pkg/examples" + "github.com/confluentinc/cli/v4/pkg/featureflags" + "github.com/confluentinc/cli/v4/pkg/flink/config" + "github.com/confluentinc/cli/v4/pkg/flink/query" + "github.com/confluentinc/cli/v4/pkg/flink/types" + "github.com/confluentinc/cli/v4/pkg/jwt" + "github.com/confluentinc/cli/v4/pkg/output" + "github.com/confluentinc/cli/v4/pkg/properties" +) + +const ( + // snapshotModeProperty makes the statement a bounded, point-in-time read. Not + // overridable via --property: the append_only envelope field's guarantee only + // holds for a snapshot read. + snapshotModeProperty = "sql.snapshot.mode" + snapshotModeNow = "now" + + // engineSnapshot is the only value Engine ever takes today; see queryOut.Engine. + engineSnapshot = "snapshot" + + // stopTimeout bounds how long we wait for a statement to stop after interrupt. + stopTimeout = 5 * time.Second + + // queryFeatureFlag gates the command's visibility. + queryFeatureFlag = "cli.query" +) + +type command struct { + *pcmd.AuthenticatedCLICommand +} + +// New mounts `confluent query` at the top level, not under `flink`: the same +// one-shot ergonomics are meant to cover other backends (e.g. Lightning Tables) +// later without a rename. +func New(cfg *cliconfig.Config, prerunner pcmd.PreRunner) *cobra.Command { + cmd := &cobra.Command{ + Use: "query [sql]", + Short: "Run a bounded Flink SQL query and print its results.", + Long: "Run a bounded (snapshot) Flink SQL query, block until it finishes, and print the complete result set.\n\n" + + "The SQL can be given as \"--sql\" or as a positional argument, but not both.\n\n" + + "Unlike statement creation, which submits a statement and returns immediately, this command waits for every " + + "result page and exits with a non-zero status if the statement fails. It is intended for scripting and " + + "one-shot queries against a bounded (point-in-time) result set.\n\n" + + "With \"-o json\" or \"-o yaml\", output defaults to an envelope carrying the column schema alongside the rows, " + + "since the rows on their own carry no type information. Pass \"--raw\" for a bare array of row objects instead.", + Args: cobra.MaximumNArgs(1), + // Hidden until the flag targets an org; cfg.IsTest keeps it visible to the + // integration suite regardless of the (unreachable in tests) LD evaluation. + Hidden: !(cfg.IsTest || featureflags.Manager.BoolVariation(queryFeatureFlag, cfg.Context(), cliconfig.CliLaunchDarklyClient, true, false)), + Annotations: map[string]string{ + pcmd.RunRequirement: pcmd.RequireNonAPIKeyCloudLogin, + }, + Example: examples.BuildExampleString( + examples.Example{ + Text: "Run a bounded query in the current compute pool and print the rows as a table.", + Code: `confluent query --sql "SELECT * FROM orders LIMIT 10;"`, + }, + examples.Example{ + Text: "Run a bounded query against Kafka cluster \"my-cluster\" and emit JSON for a script to consume.", + Code: `confluent query --sql "SELECT status, COUNT(*) FROM orders GROUP BY status;" --compute-pool lfcp-123456 --database my-cluster --output json`, + }, + examples.Example{ + Text: "Emit a bare JSON array of rows, with no envelope, for a script that only wants the data.", + Code: `confluent query --sql "SELECT * FROM orders LIMIT 10;" --output json --raw`, + }, + examples.Example{ + Text: "Pass the SQL as a positional argument instead of \"--sql\".", + Code: `confluent query "SELECT * FROM orders LIMIT 10;"`, + }, + ), + } + + c := &command{pcmd.NewAuthenticatedCLICommand(cmd, prerunner)} + cmd.RunE = c.runQuery + + cmd.Flags().String("sql", "", `The Flink SQL statement. Alternatively, pass it as a positional argument or with "-f".`) + cmd.Flags().StringP("file", "f", "", `Path to a file containing the Flink SQL statement. Alternatively, pass the SQL with "--sql" or as a positional argument.`) + c.addComputePoolFlag(cmd) + pcmd.AddServiceAccountFlag(cmd, c.AuthenticatedCLICommand) + c.addDatabaseFlag(cmd) + c.addClusterAlias(cmd) + cmd.Flags().StringSlice("property", []string{}, "A mechanism to pass properties in the form key=value when creating a Flink statement.") + cmd.Flags().Duration("wait-timeout", config.DefaultTimeoutDuration, "Maximum time to wait for the query to finish before giving up.") + cmd.Flags().Int("max-rows", 0, "Stop fetching and discard the rest after this many rows, or 0 to fetch every row. Client-side only: rows past the limit are still produced by the query.") + cmd.Flags().Bool("raw", false, `Emit the rows as a bare array with no envelope. Requires "-o json" or "-o yaml".`) + pcmd.AddEnvironmentFlag(cmd, c.AuthenticatedCLICommand) + c.addCatalogAlias(cmd) + pcmd.AddContextFlag(cmd, c.CLICommand) + pcmd.AddOutputFlag(cmd) + pcmd.AddCloudFlag(cmd) + pcmd.AddRegionFlagFlink(cmd, c.AuthenticatedCLICommand) + + return cmd +} + +// addComputePoolFlag and addDatabaseFlag mirror internal/flink's helpers, duplicated +// since this command lives outside the `flink` package boundary. +func (c *command) addComputePoolFlag(cmd *cobra.Command) { + cmd.Flags().String("compute-pool", "", "Flink compute pool ID.") + pcmd.RegisterFlagCompletionFunc(cmd, "compute-pool", c.autocompleteComputePools) +} + +func (c *command) autocompleteComputePools(cmd *cobra.Command, args []string) []string { + if err := c.PersistentPreRunE(cmd, args); err != nil { + return nil + } + + environmentId, err := c.Context.EnvironmentId() + if err != nil { + return nil + } + + computePools, err := c.V2Client.ListFlinkComputePools("", environmentId, "") + if err != nil { + return nil + } + + suggestions := make([]string, len(computePools)) + for i, computePool := range computePools { + suggestions[i] = fmt.Sprintf("%s\t%s", computePool.GetId(), computePool.Spec.GetDisplayName()) + } + return suggestions +} + +func (c *command) addDatabaseFlag(cmd *cobra.Command) { + cmd.Flags().String("database", "", "The database which will be used as the default database. When using Kafka, this is the cluster ID.") + pcmd.RegisterFlagCompletionFunc(cmd, "database", c.autocompleteDatabases) +} + +// addClusterAlias keeps "--cluster" as an independent flag rather than sharing +// storage with "--database": ParseFlagsIntoContext persists "cluster" as the active +// Kafka context but never "database", so sharing storage would leak that side effect +// onto every "--database" call. See resolveDatabase for how the two reconcile. +func (c *command) addClusterAlias(cmd *cobra.Command) { + cmd.Flags().String("cluster", "", `Alias for "--database". Unlike "--database", this also sets the CLI's active Kafka cluster context, the same as it does on every other command.`) + pcmd.RegisterFlagCompletionFunc(cmd, "cluster", c.autocompleteDatabases) +} + +// addCatalogAlias shares the same pflag.Value as "--environment" rather than +// copying it in RunE: ParseFlagsIntoContext reads "--environment" before RunE runs, +// so a later copy would be too late. Safe to share, unlike --database/--cluster, +// because "--environment" already persists to context — aliasing adds no new +// side effect. +func (c *command) addCatalogAlias(cmd *cobra.Command) { + environmentFlag := cmd.Flags().Lookup("environment") + cmd.Flags().Var(environmentFlag.Value, "catalog", `Alias for "--environment".`) +} + +func (c *command) autocompleteDatabases(cmd *cobra.Command, args []string) []string { + if err := c.PersistentPreRunE(cmd, args); err != nil { + return nil + } + + environmentId, err := c.Context.EnvironmentId() + if err != nil { + return nil + } + + clusters, err := c.V2Client.ListKafkaClusters(environmentId) + if err != nil { + return nil + } + + suggestions := make([]string, len(clusters)) + for i, cluster := range clusters { + suggestions[i] = fmt.Sprintf("%s\t%s", cluster.GetId(), cluster.Spec.GetDisplayName()) + } + return suggestions +} + +type queryColumnOut struct { + Name string `json:"name" yaml:"name"` + Type string `json:"type" yaml:"type"` +} + +type queryOut struct { + StatementName string `json:"statement_name" yaml:"statement_name"` + // Engine is always "snapshot" until M2 (Lightning routing) ships; emitting it + // now means a script switching on this field today won't need to change later. + Engine string `json:"engine" yaml:"engine"` + Phase string `json:"phase" yaml:"phase"` + Columns []queryColumnOut `json:"columns" yaml:"columns"` + // Values carry their SQL type (number, null, etc); see + // types.StatementResultField.ToSerializedValue. + Rows []map[string]any `json:"rows" yaml:"rows"` + RowCount int `json:"row_count" yaml:"row_count"` + Truncated bool `json:"truncated" yaml:"truncated"` + // Incomplete mirrors Result.Incomplete: rows may be missing. Repeated here since + // the stderr warning is invisible to a script reading only stdout. + Incomplete bool `json:"incomplete" yaml:"incomplete"` + // AppendOnly is nil until traits are known, true for insert-only statements, + // false when Rows is a changelog rather than a materialized table — the only + // such signal a script gets, since rows carry no per-row operation marker. + AppendOnly *bool `json:"append_only,omitempty" yaml:"append_only,omitempty"` +} + +func (c *command) runQuery(cmd *cobra.Command, args []string) error { + if err := resolveEnvironmentAlias(cmd); err != nil { + return err + } + + environmentId, err := c.Context.EnvironmentId() + if err != nil { + return err + } + + environment, _, err := c.V2Client.GetOrgEnvironment(environmentId) + if err != nil { + return errors.NewErrorWithSuggestions(err.Error(), "List available environments with `confluent environment list`.") + } + + // computePool falls back to context (`flink compute-pool use`); + // GetFlinkGatewayClient below decides whether pool or cloud/region is enough. + computePool := c.Context.GetCurrentFlinkComputePool() + + name := types.GenerateStatementName() + + sql, err := resolveSQL(cmd, args) + if err != nil { + return err + } + + database, err := c.resolveDatabase(cmd) + if err != nil { + return err + } + + timeout, err := cmd.Flags().GetDuration("wait-timeout") + if err != nil { + return err + } + + maxRows, err := cmd.Flags().GetInt("max-rows") + if err != nil { + return err + } + if maxRows < 0 { + return errors.New("the `--max-rows` flag must not be negative") + } + + raw, err := cmd.Flags().GetBool("raw") + if err != nil { + return err + } + if raw && !output.GetFormat(cmd).IsSerialized() { + return errors.New("the `--raw` flag requires `-o json` or `-o yaml`") + } + + statementProperties, err := c.buildQueryProperties(cmd, environment.GetDisplayName(), database) + if err != nil { + return err + } + + statement := flinkgatewayv1.SqlV1Statement{ + Name: flinkgatewayv1.PtrString(name), + Spec: &flinkgatewayv1.SqlV1StatementSpec{ + Statement: flinkgatewayv1.PtrString(sql), + Properties: &statementProperties, + }, + } + + var client *ccloudv2.FlinkGatewayClient + if computePool != "" { + statement.Spec.ComputePoolId = flinkgatewayv1.PtrString(computePool) + client, err = c.GetFlinkGatewayClient(true) + } else { + client, err = c.GetFlinkGatewayClient(false) + } + if err != nil { + return err + } + + jwtValidator := jwt.NewValidator() + + serviceAccount, err := cmd.Flags().GetString("service-account") + if err != nil { + return err + } + + principal := serviceAccount + if serviceAccount == "" { + principal = c.Context.GetUser().GetResourceId() + } + + if _, err := client.CreateStatement(statement, principal, environmentId, c.Context.LastOrgId); err != nil { + return err + } + + // From here on the statement exists server-side and is consuming pool capacity. A + // job left running with nothing draining its collect-sink buffer stalls + // indefinitely and keeps burning the compute it reserved — that was the truncation + // bug: one of several exit paths simply forgot to stop it. Making that structurally + // impossible, rather than remembering it at every exit, is the point of this defer: + // it fires unless settled is true, and settled is only set once we know the + // statement's fate — either it reached a terminal phase on its own, or something + // already made one stop attempt on our behalf. + settled := false + defer func() { + if !settled { + c.stopStatement(client, environmentId, name) + } + }() + + ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer cancel() + ctx, cancelTimeout := context.WithTimeout(ctx, timeout) + defer cancelTimeout() + + options := query.Options{ + Client: client, + EnvironmentId: environmentId, + OrganizationId: c.Context.LastOrgId, + MaxRows: maxRows, + RequireBounded: true, + RefreshToken: c.refreshGatewayToken(client, jwtValidator), + } + + result, err := query.Run(ctx, options, name) + if err != nil { + return c.handleQueryError(client, environmentId, name, err, &settled) + } + + // result.Statement's phase predates the drain loop's decision to truncate or + // give up — it says nothing about whether the job kept running after. Truncated + // and Incomplete both mean "we chose to stop reading," so both warrant a stop. + settled = !result.Truncated && !result.Incomplete && query.IsTerminal(result.Phase()) + + if result.Phase() == types.FAILED { + return errors.NewErrorWithSuggestions( + fmt.Sprintf(`statement "%s" failed: %s`, name, result.Statement.Status.GetDetail()), + fmt.Sprintf("Inspect the failure with `confluent flink statement exception list %s`.", name), + ) + } + + if result.Incomplete { + output.ErrPrintf(false, "Warning: the gateway stopped returning result pages while statement \"%s\" was still in phase %s. The result set below may be incomplete.\n", name, result.Phase()) + } + + if result.Truncated { + output.ErrPrintf(false, "Warning: stopped after %d rows because of the \"--max-rows\" flag. The result set below is truncated.\n", maxRows) + } + + traits := result.Statement.Status.GetTraits() + statementTraits := types.StatementTraits{FlinkGatewayV1StatementTraits: &traits} + isAppendOnly, appendOnlyKnown := statementTraits.GetIsAppendOnly() + if appendOnlyKnown && !isAppendOnly { + output.ErrPrintln(false, "Warning: this statement emits updates and deletions. The rows below are the raw changelog, not a materialized table.") + } + + return c.printQueryResult(cmd, name, result, isAppendOnly, appendOnlyKnown, raw) +} + +// resolveEnvironmentAlias errors if both "--environment" and "--catalog" were +// given. They share the same flag storage (see addCatalogAlias), so this only +// guards against an ambiguous invocation silently picking whichever came last. +func resolveEnvironmentAlias(cmd *cobra.Command) error { + if cmd.Flags().Changed("environment") && cmd.Flags().Changed("catalog") { + return errors.New("the environment must not be given both with the `--environment` flag and with the `--catalog` flag") + } + return nil +} + +// resolveDatabase returns whichever of "--database" or "--cluster" was given (see +// addClusterAlias), falling back to the active Kafka cluster context — the same +// "flag, then CLI context" chain environment and compute pool follow. Giving both +// flags is a usage error, mirroring resolveSQL. +func (c *command) resolveDatabase(cmd *cobra.Command) (string, error) { + database, err := cmd.Flags().GetString("database") + if err != nil { + return "", err + } + cluster, err := cmd.Flags().GetString("cluster") + if err != nil { + return "", err + } + if database != "" && cluster != "" { + return "", errors.New("the database must not be given both with the `--database` flag and with the `--cluster` flag") + } + if cluster != "" { + return cluster, nil + } + if database != "" { + return database, nil + } + return c.Context.KafkaClusterContext.GetActiveKafkaClusterId(), nil +} + +// resolveSQL returns the SQL text from whichever of "--sql", "--file", or the +// positional argument was given. Exactly one of the three is required; giving none or +// more than one is a usage error. +func resolveSQL(cmd *cobra.Command, args []string) (string, error) { + sql, err := cmd.Flags().GetString("sql") + if err != nil { + return "", err + } + + file, err := cmd.Flags().GetString("file") + if err != nil { + return "", err + } + + positional := len(args) == 1 + + switch { + case positional && sql != "", positional && file != "", sql != "" && file != "": + return "", errors.New("the SQL statement must be given exactly one way: as a positional argument, with the `--sql` flag, or with the `--file` flag") + case positional: + return args[0], nil + case sql != "": + return sql, nil + case file != "": + contents, err := os.ReadFile(file) + if err != nil { + return "", fmt.Errorf(`failed to read the SQL statement from "%s": %v`, file, err) + } + return string(contents), nil + default: + return "", errors.New("the SQL statement is required: pass it as a positional argument, with the `--sql` flag, or with the `--file` flag") + } +} + +// buildQueryProperties seeds the statement with the catalog and snapshot mode, then +// lets --property override anything else. snapshotModeProperty is the one exception: +// see its doc comment for why it must not be overridable. +func (c *command) buildQueryProperties(cmd *cobra.Command, catalog, database string) (map[string]string, error) { + statementProperties := map[string]string{ + config.KeyCatalog: catalog, + snapshotModeProperty: snapshotModeNow, + } + if database != "" { + statementProperties[config.KeyDatabase] = database + } + + configs, err := cmd.Flags().GetStringSlice("property") + if err != nil { + return nil, err + } + + if len(configs) > 0 { + configMap, err := properties.ConfigSliceToMap(configs) + if err != nil { + return nil, err + } + if mode, ok := configMap[snapshotModeProperty]; ok && mode != snapshotModeNow { + return nil, errors.NewErrorWithSuggestions( + fmt.Sprintf("the `--property` flag must not set `%s`", snapshotModeProperty), + "This command only supports a snapshot (point-in-time, append-only) read, so this property is fixed and cannot be overridden.", + ) + } + for key, value := range configMap { + statementProperties[key] = value + } + } + + return statementProperties, nil +} + +// handleQueryError turns a failed or interrupted run into a message that always +// names the statement. settled is runQuery's deferred-cleanup flag: a branch that +// stops the statement itself sets it true; a branch with no better information +// leaves it false so cleanup does the stop instead. +func (c *command) handleQueryError(client *ccloudv2.FlinkGatewayClient, environmentId, name string, err error, settled *bool) error { + var unbounded *query.UnboundedError + if goerrors.As(err, &unbounded) { + // Only claim the statement was stopped when it actually was. + fate := fmt.Sprintf("Statement \"%s\" was stopped.", name) + if !c.stopStatement(client, environmentId, name) { + fate = fmt.Sprintf("Stop it with `confluent flink statement stop %s`.", name) + } + *settled = true + return errors.NewErrorWithSuggestions( + err.Error(), + fmt.Sprintf("Bound the query with a LIMIT clause or a time predicate and run `confluent query` again. %s", fate), + ) + } + + if goerrors.Is(err, context.Canceled) || goerrors.Is(err, context.DeadlineExceeded) { + c.stopStatement(client, environmentId, name) + *settled = true + reason := "interrupted" + if goerrors.Is(err, context.DeadlineExceeded) { + reason = "timed out" + } + return errors.NewErrorWithSuggestions( + fmt.Sprintf(`query %s before statement "%s" finished`, reason, name), + fmt.Sprintf("Check the statement with `confluent flink statement describe %s`, or raise the limit with the `--wait-timeout` flag.", name), + ) + } + + // Every other error, including ResultsFetchError below, leaves settled false: + // the caller's deferred cleanup makes the stop attempt these branches skip. + var resultsFetchErr *query.ResultsFetchError + if goerrors.As(err, &resultsFetchErr) { + var coder flinkerror.Coder + if goerrors.As(err, &coder) { + switch coder.StatusCode() { + case http.StatusNotFound: + // A 404 here means the statement itself is gone or was mistyped, not + // an expired result window — the gateway signals that separately (see + // the 408 case below). Confirmed from the gateway's + // GetStatementResultEndpoint, not inferred from client behavior. + return errors.NewErrorWithSuggestions( + resultsFetchErr.Error(), + fmt.Sprintf("Statement \"%s\" no longer exists — it may have been deleted, or the name is mistyped. Check `confluent flink statement describe %s`.", name, name), + ) + case http.StatusRequestTimeout: + // Snapshot query results are retained for exactly one hour after the + // statement is created; past that the gateway returns 408 with its own + // explicit message (already carried in resultsFetchErr.Error()) instead + // of continuing to page. + return errors.NewErrorWithSuggestions( + resultsFetchErr.Error(), + "Re-run the query — the result window for this statement has closed.", + ) + } + } + return errors.NewErrorWithSuggestions( + resultsFetchErr.Error(), + fmt.Sprintf("Check the statement with `confluent flink statement describe %s`.", name), + ) + } + + return errors.NewErrorWithSuggestions( + err.Error(), + fmt.Sprintf("Check the statement with `confluent flink statement describe %s`.", name), + ) +} + +// refreshGatewayToken mirrors the shell's pre-call check: without it, a query +// outliving the short-lived dataplane token dies on a 401 before --wait-timeout. +func (c *command) refreshGatewayToken(client *ccloudv2.FlinkGatewayClient, jwtValidator jwt.Validator) func() error { + return func() error { + jwtCtx := &cliconfig.Context{State: &cliconfig.ContextState{AuthToken: client.AuthToken}} + if jwtValidator.Validate(jwtCtx) == nil { + return nil + } + + dataplaneToken, err := auth.GetDataplaneToken(c.Context) + if err != nil { + return err + } + client.AuthToken = dataplaneToken + return nil + } +} + +// stopStatement makes a best-effort, bounded attempt to stop an abandoned +// statement and reports the outcome either way. +func (c *command) stopStatement(client *ccloudv2.FlinkGatewayClient, environmentId, name string) bool { + done := make(chan error, 1) + go func() { + // The gateway rejects a body carrying only spec.stopped as malformed; read + // the statement back and flip the flag on what it returns. + statement, err := client.GetStatement(environmentId, name, c.Context.LastOrgId) + if err != nil { + done <- err + return + } + if statement.Spec == nil { + done <- fmt.Errorf(`statement "%s" has no spec`, name) + return + } + statement.Spec.Stopped = flinkgatewayv1.PtrBool(true) + done <- client.UpdateStatement(environmentId, name, c.Context.LastOrgId, statement) + }() + + select { + case err := <-done: + if err != nil { + output.ErrPrintf(false, "Warning: could not stop statement \"%s\": %v. It may still be running.\n", name, err) + return false + } + output.ErrPrintf(false, "Stopped statement \"%s\".\n", name) + return true + case <-time.After(stopTimeout): + output.ErrPrintf(false, "Warning: timed out trying to stop statement \"%s\". It may still be running.\n", name) + return false + } +} + +func (c *command) printQueryResult(cmd *cobra.Command, name string, result *query.Result, isAppendOnly, appendOnlyKnown, raw bool) error { + columns := make([]queryColumnOut, len(result.Columns)) + headers := make([]string, len(result.Columns)) + for i, column := range result.Columns { + columnType := column.GetType() + columns[i] = queryColumnOut{Name: column.GetName(), Type: columnType.GetType()} + headers[i] = column.GetName() + } + + showOperation := appendOnlyKnown && !isAppendOnly + + if output.GetFormat(cmd).IsSerialized() { + rows := make([]map[string]any, len(result.Rows)) + for i, row := range result.Rows { + // Every row is guaranteed len(headers) fields: Run() hard-errors on a + // row/schema mismatch before this function ever sees a result. + fields := make(map[string]any, len(headers)) + for j, field := range row.GetFields() { + fields[headers[j]] = field.ToSerializedValue() + } + rows[i] = fields + } + + // A bare array has nowhere to put the schema, so the envelope is the + // default and --raw opts into the bare array. + if raw { + return output.SerializedOutput(cmd, rows) + } + + var appendOnly *bool + if appendOnlyKnown { + appendOnly = &isAppendOnly + } + + return output.SerializedOutput(cmd, &queryOut{ + StatementName: name, + Engine: engineSnapshot, + Phase: string(result.Phase()), + Columns: columns, + Rows: rows, + RowCount: len(rows), + Truncated: result.Truncated, + Incomplete: result.Incomplete, + AppendOnly: appendOnly, + }) + } + + if len(headers) == 0 || len(result.Rows) == 0 { + output.ErrPrintf(false, "The query returned no rows. Statement \"%s\" is in phase %s.\n", name, result.Phase()) + return nil + } + + if showOperation { + headers = append([]string{"Operation"}, headers...) + } + + rows := make([][]string, len(result.Rows)) + for i, row := range result.Rows { + fields := make([]string, 0, len(headers)) + if showOperation { + fields = append(fields, row.Operation.String()) + } + for _, field := range row.GetFields() { + fields = append(fields, field.ToString()) + } + rows[i] = fields + } + + // No column truncation: this command is expected to be piped, and shortening + // a value would corrupt whatever reads it. + table := tablewriter.NewWriter(os.Stdout) + table.SetAutoFormatHeaders(false) + table.SetAutoWrapText(false) + table.SetAlignment(tablewriter.ALIGN_LEFT) + table.SetHeader(headers) + table.AppendBulk(rows) + table.Render() + + return nil +} diff --git a/internal/query/command_test.go b/internal/query/command_test.go new file mode 100644 index 0000000000..01ecee89e9 --- /dev/null +++ b/internal/query/command_test.go @@ -0,0 +1,637 @@ +package query + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/require" + + flinkgatewayv1 "github.com/confluentinc/ccloud-sdk-go-v2/flink-gateway/v1" + + climock "github.com/confluentinc/cli/v4/mock" + "github.com/confluentinc/cli/v4/pkg/ccloudv2" + pcmd "github.com/confluentinc/cli/v4/pkg/cmd" + cliconfig "github.com/confluentinc/cli/v4/pkg/config" + "github.com/confluentinc/cli/v4/pkg/errors" + flinkerror "github.com/confluentinc/cli/v4/pkg/errors/flink" + "github.com/confluentinc/cli/v4/pkg/flink/query" + "github.com/confluentinc/cli/v4/pkg/flink/types" + testserver "github.com/confluentinc/cli/v4/test/test-server" +) + +func TestNew(t *testing.T) { + cfg := cliconfig.AuthenticatedCloudConfigMock() + prerunner := climock.NewPreRunnerMock(nil, nil, nil, nil, cfg) + + cmd := New(cfg, prerunner) + + require.Equal(t, "query [sql]", cmd.Use) + require.False(t, cmd.Hidden, "cfg.IsTest should keep the command visible in tests") + + for _, name := range []string{"sql", "file", "compute-pool", "service-account", "database", "cluster", "property", "wait-timeout", "max-rows", "raw", "environment", "catalog", "context", "output", "cloud", "region"} { + require.NotNil(t, cmd.Flags().Lookup(name), "expected --%s to be registered", name) + } + + // "sql" is deliberately not cobra-required: it can come from the positional + // argument instead, so requiredness is enforced by resolveSQL, not by cobra. + sqlFlag := cmd.Flags().Lookup("sql") + require.Empty(t, sqlFlag.Annotations[cobra.BashCompOneRequiredFlag]) + + // --catalog shares storage with --environment: setting one is setting the other. + require.NoError(t, cmd.Flags().Set("catalog", "env-999")) + environment, err := cmd.Flags().GetString("environment") + require.NoError(t, err) + require.Equal(t, "env-999", environment) +} + +func TestResolveEnvironmentAlias(t *testing.T) { + newEnvCmd := func() *cobra.Command { + cmd := &cobra.Command{} + v := new(stringValueForTest) + cmd.Flags().Var(v, "environment", "") + cmd.Flags().Var(v, "catalog", "") + return cmd + } + + // Neither given, or only one given: fine. + require.NoError(t, resolveEnvironmentAlias(newEnvCmd())) + + cmd := newEnvCmd() + require.NoError(t, cmd.Flags().Set("environment", "env-123")) + require.NoError(t, resolveEnvironmentAlias(cmd)) + + // Both explicitly given: a usage error, even though they'd resolve to the same + // underlying value. + cmd = newEnvCmd() + require.NoError(t, cmd.Flags().Set("environment", "env-123")) + require.NoError(t, cmd.Flags().Set("catalog", "env-123")) + require.ErrorContains(t, resolveEnvironmentAlias(cmd), "must not be given both") +} + +// stringValueForTest is a minimal pflag.Value letting two flag names share one +// value, the way addCatalogAlias does for "environment"/"catalog". +type stringValueForTest string + +func (s *stringValueForTest) String() string { return string(*s) } +func (s *stringValueForTest) Set(v string) error { *s = stringValueForTest(v); return nil } +func (s *stringValueForTest) Type() string { return "string" } + +func TestResolveDatabase(t *testing.T) { + newDBCmd := func(database, cluster string) *cobra.Command { + cmd := &cobra.Command{} + cmd.Flags().String("database", "", "") + cmd.Flags().String("cluster", "", "") + if database != "" { + require.NoError(t, cmd.Flags().Set("database", database)) + } + if cluster != "" { + require.NoError(t, cmd.Flags().Set("cluster", cluster)) + } + return cmd + } + + commandWithActiveCluster := func(activeCluster string) *command { + return newTestCommand(&cliconfig.Context{KafkaClusterContext: &cliconfig.KafkaClusterContext{ActiveKafkaCluster: activeCluster}}) + } + + c := commandWithActiveCluster("") + + got, err := c.resolveDatabase(newDBCmd("", "")) + require.NoError(t, err) + require.Empty(t, got) + + got, err = c.resolveDatabase(newDBCmd("lkc-database", "")) + require.NoError(t, err) + require.Equal(t, "lkc-database", got) + + got, err = c.resolveDatabase(newDBCmd("", "lkc-cluster")) + require.NoError(t, err) + require.Equal(t, "lkc-cluster", got) + + _, err = c.resolveDatabase(newDBCmd("lkc-database", "lkc-cluster")) + require.ErrorContains(t, err, "must not be given both") + + // Neither flag given: falls back to the active Kafka cluster context, same + // "flag, then context" chain environment and compute pool follow. + c = commandWithActiveCluster("lkc-context-default") + got, err = c.resolveDatabase(newDBCmd("", "")) + require.NoError(t, err) + require.Equal(t, "lkc-context-default", got) + + // An explicit --database still wins over the context default. + got, err = c.resolveDatabase(newDBCmd("lkc-explicit", "")) + require.NoError(t, err) + require.Equal(t, "lkc-explicit", got) +} + +func TestResolveSQL(t *testing.T) { + newSQLCmd := func(sqlFlagValue, fileFlagValue string) *cobra.Command { + cmd := &cobra.Command{} + cmd.Flags().String("sql", "", "") + cmd.Flags().String("file", "", "") + if sqlFlagValue != "" { + require.NoError(t, cmd.Flags().Set("sql", sqlFlagValue)) + } + if fileFlagValue != "" { + require.NoError(t, cmd.Flags().Set("file", fileFlagValue)) + } + return cmd + } + + sql, err := resolveSQL(newSQLCmd("SELECT 1", ""), nil) + require.NoError(t, err) + require.Equal(t, "SELECT 1", sql) + + sql, err = resolveSQL(newSQLCmd("", ""), []string{"SELECT 2"}) + require.NoError(t, err) + require.Equal(t, "SELECT 2", sql) + + sqlFile := filepath.Join(t.TempDir(), "query.sql") + require.NoError(t, os.WriteFile(sqlFile, []byte("SELECT 3"), 0o600)) + sql, err = resolveSQL(newSQLCmd("", sqlFile), nil) + require.NoError(t, err) + require.Equal(t, "SELECT 3", sql) + + _, err = resolveSQL(newSQLCmd("", "/nonexistent/query.sql"), nil) + require.ErrorContains(t, err, "failed to read the SQL statement") + + _, err = resolveSQL(newSQLCmd("SELECT 1", ""), []string{"SELECT 2"}) + require.ErrorContains(t, err, "must be given exactly one way") + + _, err = resolveSQL(newSQLCmd("SELECT 1", sqlFile), nil) + require.ErrorContains(t, err, "must be given exactly one way") + + _, err = resolveSQL(newSQLCmd("", sqlFile), []string{"SELECT 2"}) + require.ErrorContains(t, err, "must be given exactly one way") + + _, err = resolveSQL(newSQLCmd("", ""), nil) + require.ErrorContains(t, err, "is required") +} + +// captureStdout redirects the package-level os.Stdout (which output.Print and +// tablewriter both write to directly) for the duration of fn and returns what +// was written. +func captureStdout(t *testing.T, fn func()) string { + t.Helper() + + orig := os.Stdout + r, w, err := os.Pipe() + require.NoError(t, err) + os.Stdout = w + defer func() { os.Stdout = orig }() + + fn() + + require.NoError(t, w.Close()) + out, err := io.ReadAll(r) + require.NoError(t, err) + return string(out) +} + +func newTestCommand(ctx *cliconfig.Context) *command { + return &command{AuthenticatedCLICommand: &pcmd.AuthenticatedCLICommand{Context: ctx}} +} + +func newTestContext(platformServer, authToken string) *cliconfig.Context { + return &cliconfig.Context{ + Platform: &cliconfig.Platform{Server: platformServer}, + State: &cliconfig.ContextState{AuthToken: authToken}, + LastOrgId: "org-1", + } +} + +func TestBuildQueryProperties(t *testing.T) { + tests := []struct { + name string + catalog string + database string + flags []string + expected map[string]string + wantErr bool + }{ + { + name: "catalog and default snapshot mode, no database", + catalog: "env-123", + expected: map[string]string{ + "sql.current-catalog": "env-123", + "sql.snapshot.mode": "now", + }, + }, + { + name: "database is included when set", + catalog: "env-123", + database: "my-cluster", + expected: map[string]string{ + "sql.current-catalog": "env-123", + "sql.snapshot.mode": "now", + "sql.current-database": "my-cluster", + }, + }, + { + name: "property flag cannot override the snapshot mode", + catalog: "env-123", + flags: []string{"sql.snapshot.mode=earliest"}, + wantErr: true, + }, + { + name: "property flag redundantly setting the snapshot mode to its default is allowed", + catalog: "env-123", + flags: []string{"sql.snapshot.mode=now"}, + expected: map[string]string{ + "sql.current-catalog": "env-123", + "sql.snapshot.mode": "now", + }, + }, + { + name: "malformed property flag is rejected", + catalog: "env-123", + flags: []string{"not-a-key-value-pair"}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cmd := &cobra.Command{} + cmd.Flags().StringSlice("property", []string{}, "") + for _, f := range tt.flags { + require.NoError(t, cmd.Flags().Set("property", f)) + } + + c := newTestCommand(nil) + got, err := c.buildQueryProperties(cmd, tt.catalog, tt.database) + if tt.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Equal(t, tt.expected, got) + }) + } +} + +func newOutputCmd(t *testing.T, format string) *cobra.Command { + t.Helper() + cmd := &cobra.Command{} + pcmd.AddOutputFlag(cmd) + if format != "" { + require.NoError(t, cmd.Flags().Set("output", format)) + } + return cmd +} + +func testColumns() []flinkgatewayv1.ColumnDetails { + return []flinkgatewayv1.ColumnDetails{ + {Name: "id", Type: flinkgatewayv1.DataType{Type: "INTEGER"}}, + {Name: "status", Type: flinkgatewayv1.DataType{Type: "VARCHAR"}}, + } +} + +func testRow() types.StatementResultRow { + return types.StatementResultRow{ + Operation: types.Insert, + Fields: []types.StatementResultField{ + types.AtomicStatementResultField{Type: types.Integer, Value: "1021"}, + types.AtomicStatementResultField{Type: types.Varchar, Value: "SHIPPED"}, + }, + } +} + +func TestPrintQueryResult(t *testing.T) { + t.Run("human table output with rows", func(t *testing.T) { + c := newTestCommand(nil) + cmd := newOutputCmd(t, "") + result := &query.Result{ + Statement: flinkgatewayv1.SqlV1Statement{Status: &flinkgatewayv1.SqlV1StatementStatus{Phase: "COMPLETED"}}, + Columns: testColumns(), + Rows: []types.StatementResultRow{testRow()}, + } + + out := captureStdout(t, func() { + require.NoError(t, c.printQueryResult(cmd, "stmt", result, false, false, false)) + }) + require.Contains(t, out, "1021") + require.Contains(t, out, "SHIPPED") + require.NotContains(t, out, "Operation") + }) + + t.Run("human table output shows Operation column when requested", func(t *testing.T) { + c := newTestCommand(nil) + cmd := newOutputCmd(t, "") + result := &query.Result{ + Statement: flinkgatewayv1.SqlV1Statement{Status: &flinkgatewayv1.SqlV1StatementStatus{Phase: "COMPLETED"}}, + Columns: testColumns(), + Rows: []types.StatementResultRow{testRow()}, + } + + out := captureStdout(t, func() { + require.NoError(t, c.printQueryResult(cmd, "stmt", result, false, true, false)) + }) + require.Contains(t, out, "Operation") + }) + + t.Run("human output with no rows prints a message to stderr, not stdout", func(t *testing.T) { + c := newTestCommand(nil) + cmd := newOutputCmd(t, "") + result := &query.Result{ + Statement: flinkgatewayv1.SqlV1Statement{Status: &flinkgatewayv1.SqlV1StatementStatus{Phase: "COMPLETED"}}, + } + + out := captureStdout(t, func() { + require.NoError(t, c.printQueryResult(cmd, "stmt", result, false, false, false)) + }) + require.Empty(t, out) + }) + + t.Run("json envelope carries schema, truncated and incomplete", func(t *testing.T) { + c := newTestCommand(nil) + cmd := newOutputCmd(t, "json") + result := &query.Result{ + Statement: flinkgatewayv1.SqlV1Statement{Status: &flinkgatewayv1.SqlV1StatementStatus{Phase: "RUNNING"}}, + Columns: testColumns(), + Rows: []types.StatementResultRow{testRow()}, + Truncated: true, + Incomplete: true, + } + + out := captureStdout(t, func() { + require.NoError(t, c.printQueryResult(cmd, "stmt", result, false, false, false)) + }) + require.Contains(t, out, `"statement_name": "stmt"`) + require.Contains(t, out, `"engine": "snapshot"`) + require.Contains(t, out, `"phase": "RUNNING"`) + require.Contains(t, out, `"truncated": true`) + require.Contains(t, out, `"incomplete": true`) + require.Contains(t, out, `"id": 1021`) + require.NotContains(t, out, "append_only") + }) + + t.Run("json envelope carries append_only when known", func(t *testing.T) { + c := newTestCommand(nil) + cmd := newOutputCmd(t, "json") + result := &query.Result{ + Statement: flinkgatewayv1.SqlV1Statement{Status: &flinkgatewayv1.SqlV1StatementStatus{Phase: "COMPLETED"}}, + Columns: testColumns(), + Rows: []types.StatementResultRow{testRow()}, + } + + out := captureStdout(t, func() { + require.NoError(t, c.printQueryResult(cmd, "stmt", result, false, true, false)) + }) + require.Contains(t, out, `"append_only": false`) + }) + + t.Run("raw serialized output is a bare array with no envelope", func(t *testing.T) { + c := newTestCommand(nil) + cmd := newOutputCmd(t, "json") + result := &query.Result{ + Statement: flinkgatewayv1.SqlV1Statement{Status: &flinkgatewayv1.SqlV1StatementStatus{Phase: "COMPLETED"}}, + Columns: testColumns(), + Rows: []types.StatementResultRow{testRow()}, + } + + out := captureStdout(t, func() { + require.NoError(t, c.printQueryResult(cmd, "stmt", result, false, false, true)) + }) + require.NotContains(t, out, "statement_name") + require.NotContains(t, out, "engine") + require.Contains(t, out, `"id": 1021`) + }) + + t.Run("yaml serialized output", func(t *testing.T) { + c := newTestCommand(nil) + cmd := newOutputCmd(t, "yaml") + result := &query.Result{ + Statement: flinkgatewayv1.SqlV1Statement{Status: &flinkgatewayv1.SqlV1StatementStatus{Phase: "COMPLETED"}}, + Columns: testColumns(), + Rows: []types.StatementResultRow{testRow()}, + } + + out := captureStdout(t, func() { + require.NoError(t, c.printQueryResult(cmd, "stmt", result, false, false, false)) + }) + require.Contains(t, out, "statement_name: stmt") + require.Contains(t, out, "engine: snapshot") + }) +} + +// fakeJwtValidator lets tests control whether refreshGatewayToken thinks the +// current token is still valid without needing a real signed JWT. +type fakeJwtValidator struct { + err error +} + +func (f fakeJwtValidator) Validate(*cliconfig.Context) error { + return f.err +} + +func TestRefreshGatewayToken(t *testing.T) { + t.Run("valid token is left alone", func(t *testing.T) { + client := ccloudv2.NewFlinkGatewayClient("http://unused.invalid", "test", false, "still-valid") + c := newTestCommand(newTestContext("http://unused.invalid", "still-valid")) + + refresh := c.refreshGatewayToken(client, fakeJwtValidator{err: nil}) + require.NoError(t, refresh()) + require.Equal(t, "still-valid", client.AuthToken) + }) + + t.Run("expired token is refreshed from the platform", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/api/access_tokens", r.URL.Path) + require.Equal(t, "Bearer old-cloud-token", r.Header.Get("Authorization")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"token":"new-dataplane-token"}`)) + })) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient("http://unused.invalid", "test", false, "expired") + c := newTestCommand(newTestContext(server.URL, "old-cloud-token")) + + refresh := c.refreshGatewayToken(client, fakeJwtValidator{err: errors.New("expired")}) + require.NoError(t, refresh()) + require.Equal(t, "new-dataplane-token", client.AuthToken) + }) + + t.Run("refresh failure surfaces the platform's error", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"error":"could not mint a dataplane token"}`)) + })) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient("http://unused.invalid", "test", false, "expired") + c := newTestCommand(newTestContext(server.URL, "old-cloud-token")) + + refresh := c.refreshGatewayToken(client, fakeJwtValidator{err: errors.New("expired")}) + require.ErrorContains(t, refresh(), "could not mint a dataplane token") + require.Equal(t, "expired", client.AuthToken) + }) +} + +func TestStopStatement(t *testing.T) { + t.Run("successful stop", func(t *testing.T) { + server := httptest.NewServer(testserver.NewFlinkGatewayRouter(t)) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient(server.URL, "test", false, "token") + c := newTestCommand(newTestContext(server.URL, "token")) + + out := captureStderr(t, func() { + require.True(t, c.stopStatement(client, "env-1", "stmt")) + }) + require.Contains(t, out, `Stopped statement "stmt"`) + }) + + t.Run("failed stop reports a warning and returns false", func(t *testing.T) { + // A 4xx, not 5xx: the retryable HTTP client retries 5xx/429 responses, which + // would blow past stopTimeout and hit the timeout branch instead of this one. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + })) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient(server.URL, "test", false, "token") + c := newTestCommand(newTestContext(server.URL, "token")) + + out := captureStderr(t, func() { + require.False(t, c.stopStatement(client, "env-1", "stmt")) + }) + require.Contains(t, out, `could not stop statement "stmt"`) + }) +} + +func captureStderr(t *testing.T, fn func()) string { + t.Helper() + + orig := os.Stderr + r, w, err := os.Pipe() + require.NoError(t, err) + os.Stderr = w + defer func() { os.Stderr = orig }() + + fn() + + require.NoError(t, w.Close()) + out, err := io.ReadAll(r) + require.NoError(t, err) + return string(out) +} + +func TestHandleQueryError(t *testing.T) { + t.Run("unbounded error stops the statement and names it in the suggestion", func(t *testing.T) { + server := httptest.NewServer(testserver.NewFlinkGatewayRouter(t)) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient(server.URL, "test", false, "token") + c := newTestCommand(newTestContext(server.URL, "token")) + + settled := false + var out string + var err error + out = captureStderr(t, func() { + err = c.handleQueryError(client, "env-1", "stmt", &query.UnboundedError{StatementName: "stmt"}, &settled) + }) + require.Error(t, err) + require.Contains(t, err.Error(), "unbounded result") + var withSuggestions errors.ErrorWithSuggestions + require.ErrorAs(t, err, &withSuggestions) + require.Contains(t, withSuggestions.GetSuggestionsMsg(), `Statement "stmt" was stopped.`) + require.True(t, settled) + require.Contains(t, out, `Stopped statement "stmt"`) + }) + + t.Run("unbounded error names the manual stop command when the stop attempt fails", func(t *testing.T) { + // A 4xx, not 5xx: see the equivalent comment in TestStopStatement. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + })) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient(server.URL, "test", false, "token") + c := newTestCommand(newTestContext(server.URL, "token")) + + settled := false + err := c.handleQueryError(client, "env-1", "stmt", &query.UnboundedError{StatementName: "stmt"}, &settled) + require.Error(t, err) + var withSuggestions errors.ErrorWithSuggestions + require.ErrorAs(t, err, &withSuggestions) + require.Contains(t, withSuggestions.GetSuggestionsMsg(), "confluent flink statement stop stmt") + require.True(t, settled) + }) + + t.Run("context canceled is reported as interrupted", func(t *testing.T) { + server := httptest.NewServer(testserver.NewFlinkGatewayRouter(t)) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient(server.URL, "test", false, "token") + c := newTestCommand(newTestContext(server.URL, "token")) + + settled := false + err := c.handleQueryError(client, "env-1", "stmt", context.Canceled, &settled) + require.Error(t, err) + require.Contains(t, err.Error(), "interrupted") + require.True(t, settled) + }) + + t.Run("context deadline exceeded is reported as timed out", func(t *testing.T) { + server := httptest.NewServer(testserver.NewFlinkGatewayRouter(t)) + defer server.Close() + + client := ccloudv2.NewFlinkGatewayClient(server.URL, "test", false, "token") + c := newTestCommand(newTestContext(server.URL, "token")) + + settled := false + err := c.handleQueryError(client, "env-1", "stmt", context.DeadlineExceeded, &settled) + require.Error(t, err) + require.Contains(t, err.Error(), "timed out") + require.True(t, settled) + }) + + t.Run("a 404 results-fetch error suggests the statement is gone or mistyped", func(t *testing.T) { + c := newTestCommand(nil) + settled := false + err := c.handleQueryError(nil, "env-1", "stmt", &query.ResultsFetchError{Err: flinkerror.NewError("not found", "", http.StatusNotFound)}, &settled) + require.Error(t, err) + var withSuggestions errors.ErrorWithSuggestions + require.ErrorAs(t, err, &withSuggestions) + require.Contains(t, withSuggestions.GetSuggestionsMsg(), "no longer exists") + require.False(t, settled) + }) + + t.Run("a 408 results-fetch error tells the user to re-run the query", func(t *testing.T) { + c := newTestCommand(nil) + settled := false + err := c.handleQueryError(nil, "env-1", "stmt", &query.ResultsFetchError{Err: flinkerror.NewError("Snapshot statement results are only available for 1 hour after the statement is created.", "", http.StatusRequestTimeout)}, &settled) + require.Error(t, err) + var withSuggestions errors.ErrorWithSuggestions + require.ErrorAs(t, err, &withSuggestions) + require.Contains(t, withSuggestions.GetSuggestionsMsg(), "Re-run the query") + require.False(t, settled) + }) + + t.Run("a non-404 results-fetch error gets the generic suggestion", func(t *testing.T) { + c := newTestCommand(nil) + settled := false + err := c.handleQueryError(nil, "env-1", "stmt", &query.ResultsFetchError{Err: flinkerror.NewError("boom", "", http.StatusInternalServerError)}, &settled) + require.Error(t, err) + var withSuggestions errors.ErrorWithSuggestions + require.ErrorAs(t, err, &withSuggestions) + require.Contains(t, withSuggestions.GetSuggestionsMsg(), "confluent flink statement describe stmt") + require.False(t, settled) + }) + + t.Run("any other error falls back to the generic suggestion", func(t *testing.T) { + c := newTestCommand(nil) + settled := false + err := c.handleQueryError(nil, "env-1", "stmt", errors.New("some other failure"), &settled) + require.Error(t, err) + require.Contains(t, err.Error(), "some other failure") + require.False(t, settled) + }) +} diff --git a/pkg/ccloudv2/utils.go b/pkg/ccloudv2/utils.go index 3a5794f1b8..63335fc1bc 100644 --- a/pkg/ccloudv2/utils.go +++ b/pkg/ccloudv2/utils.go @@ -121,7 +121,11 @@ func getServerUrl(baseURL string) string { return u.String() } -func extractPageToken(nextPageUrlString string) (string, error) { +// ExtractPageToken pulls page_token out of an absolute "next page" URL. Exported +// for callers outside this package with a similar next-URL shape. Errors on a +// non-empty URL with no page_token — callers treating an empty URL itself as "no +// more pages" must check for that first, as extractNextPageToken below does. +func ExtractPageToken(nextPageUrlString string) (string, error) { nextPageUrl, err := url.Parse(nextPageUrlString) if err != nil { plog.CliLogger.Errorf("Could not parse %s into URL, %v", nextPageUrlString, err) @@ -142,6 +146,6 @@ func extractNextPageToken(nextPageUrl NullableString) (string, bool, error) { if nextPageUrlString == "" { return "", true, nil } - pageToken, err := extractPageToken(nextPageUrlString) + pageToken, err := ExtractPageToken(nextPageUrlString) return pageToken, false, err } diff --git a/pkg/ccloudv2/utils_test.go b/pkg/ccloudv2/utils_test.go index 72668bb63f..521bcf9884 100644 --- a/pkg/ccloudv2/utils_test.go +++ b/pkg/ccloudv2/utils_test.go @@ -51,3 +51,14 @@ func TestToLower(t *testing.T) { func TestToUpper(t *testing.T) { require.Equal(t, "SASL_SSL", ToUpper("sasl-ssl")) } + +func TestExtractPageToken(t *testing.T) { + token, err := ExtractPageToken("https://example.com/results?page_token=20") + require.NoError(t, err) + require.Equal(t, "20", token) +} + +func TestExtractPageToken_MissingToken(t *testing.T) { + _, err := ExtractPageToken("https://example.com/results") + require.ErrorContains(t, err, `could not parse the value for query parameter "page_token"`) +} diff --git a/pkg/flink/query/README.md b/pkg/flink/query/README.md new file mode 100644 index 0000000000..f47eb916bd --- /dev/null +++ b/pkg/flink/query/README.md @@ -0,0 +1,125 @@ +### query + +Runs a bounded ("snapshot") Flink SQL statement to completion and returns the whole +result set, synchronously, from the client. Backs `confluent query` +(`internal/query/command.go`). + +The verb, the flags and the result shape are all expected to move. + +#### Why this exists + +The Flink gateway has no synchronous execute endpoint. A statement is submitted, polled +until it leaves `PENDING`, and then its result pages are pulled one at a time following +`metadata.next`. `Run` performs that handshake behind a single call so a non-interactive +command can behave like an ordinary database query. + +```go +result, err := query.Run(ctx, query.Options{ + Client: gatewayClient, + EnvironmentId: environmentId, + OrganizationId: organizationId, + RequireBounded: true, +}, statementName) +``` + +The statement must already exist — submitting it is the caller's job, so the caller keeps +ownership of naming, properties and cleanup. + +#### Why not reuse the shell + +`Store` + `ResultFetcher` + `MaterializedStatementResults` already does submit, poll and +page. It is not reused here, because it was written for a scrolling viewer where being +wrong degrades to "the user presses refresh again". Three of those degradations become +silent data corruption once a script is reading stdout: + +| Shell behavior | Consequence for a synchronous command | +| --- | --- | +| `MaterializedStatementResults.cleanup()` evicts from the front past `MaxResultsCapacity` (10,000) | a 50k-row `SELECT` prints the **last** 10k and exits 0 | +| `Append` skips rows whose field count ≠ header count, returning a bool that `fetchNextPageAndUpdateState` discards | short result set, no signal | +| `updateState` sets `Completed` on `PageToken == ""` without checking the phase | exit 0 with a partial result set | + +This package handles each explicitly: no cap unless `Options.MaxRows` is set (and then +`Result.Truncated` says so), a hard error from `ConvertToInternalResults` on a schema +mismatch, and termination only when the page token is gone **and** the statement has +reached a terminal phase. + +It also skips the shell's table-mode materialization. `Result.Rows` is the raw changelog +as the gateway delivered it. For a bounded append-only snapshot the changelog and the +materialized table are identical; for anything else the caller decides. + +#### The drain loop + +`page_token` is a positional offset into the collect-sink protocol, not an opaque cursor. +There is therefore no token that advances past a page which did not supply one. + +The gateway's own foreground/streaming endpoint (`GetStatementResultEndpoint` in +`cc-flink-gateway-service-v2`, `internal/service/sql/v1/service.go`) only leaves `next` +empty when the JobManager itself reports `IsFinished == true` — every other case, even an +empty page, gets a fresh `next` token. So an empty `next` is not ambiguous at the +protocol level; it means the JobManager is genuinely done producing rows. + +What can still go wrong is a race between two *separately updated* signals: the +JobManager's own `IsFinished` (which drives whether `next` is populated) versus the +statement's `Status.Phase`, read here via a separate `GetStatement` call reconciled by a +different subsystem. If that call lands just before the JobManager flips to finished, +`terminalBeforeFetch` comes back `false` even though the page fetched right after is +already the last one. `drain()` handles this by re-reading the phase once more instead of +conceding immediately: + +- the page was **empty** — treated as "nothing yet". Re-requesting the same offset is + harmless, so the loop backs off and retries. +- the page carried **rows** and the phase read before the fetch wasn't terminal — the + loop re-reads the phase once more before giving up. Only if that second read is still + non-terminal does it set `Result.Incomplete`, which the command surfaces as a warning. + +This is confirmed from gateway source for the one execution path this package exercises +(foreground, JobManager-backed snapshot statements) — not verified for other paths (e.g. +any legacy/batch branch this package never hits), so the phase re-check stays as +defense-in-depth rather than being narrowed or removed. + +#### Known limitations + +- **Cloud only.** `Options.Client` is a `ccloudv2.GatewayClientInterface`. On-prem goes + through CMF (`store_onprem.go`, itself a near-copy of `store.go`), so parity means a + second implementation and a second test surface. +- **Token refresh is best-effort, not retry-aware.** `Options.RefreshToken` is invoked + before each gateway call (see the command's `refreshGatewayToken`), unlike the shell's + `synchronizedTokenRefresh`, which wraps every call including mid-flight retries. In + practice this rarely matters: the command's default 10-minute `--timeout` is on the + same order as the dataplane token's own lifetime, so a run is unlikely to still be + going when a refresh would be needed. It only bites if `--timeout` is raised well past + the default. +- **Expired-result handling lives in the command, not here.** This package just returns + `ResultsFetchError` on any failed page fetch. `internal/query/command.go`'s + `handleQueryError` is what distinguishes a 404 (statement deleted or mistyped) from a + 408 (the snapshot result window has closed) and gives each a targeted suggestion. + Confirmed from gateway source: for `sql.snapshot.mode=now` statements, results are + retained for exactly one hour after statement creation + (`cc-flink-gateway-service-v2` `internal/service/sql/v1/service.go`, + `GetStatementResultEndpoint`), after which the gateway returns 408 with an explicit + message instead of continuing to page. +- **The whole result set is held in memory** as `[]types.StatementResultRow`. There is no + streaming-to-stdout path, so the peak footprint scales with the result. +- **Ops are dropped from serialized output.** The command's `-o json` / `-o yaml` rows are + keyed by column name and carry no `op`, so a non-append-only statement loses its + update/delete markers. Human output grows an `Operation` column instead. Real gap if + non-append-only ever needs supporting. +- **No statement cleanup on success.** Each query leaves a terminal statement behind + against the 50K-per-environment pool. +- **`--unsafe-trace` dumps customer rows**, and this is the surface most likely to run in + CI with retained logs. + +#### Working on this + +```bash +go test ./pkg/flink/query/ # unit tests, mocked gateway +``` + +The unit tests drive `pkg/flink/test/mock.MockGatewayClientInterface` and inject +`Options.sleep`, so backoff costs no wall time. + +Two unrelated failures reproduce on a clean `main` and are not caused by changes here: +`pkg/flink/internal/controller` and `TestFlinkShell`/`TestFlinkShellOnPrem` panic without +a TTY. If `make lint-go` reports `unsupported version of the configuration`, a +golangci-lint v2 on `PATH` is shadowing the v1.64.8 the Makefile pins — run +`$(go env GOPATH)/bin/golangci-lint run` directly. diff --git a/pkg/flink/query/query.go b/pkg/flink/query/query.go new file mode 100644 index 0000000000..f109ebcccc --- /dev/null +++ b/pkg/flink/query/query.go @@ -0,0 +1,279 @@ +// Package query runs a bounded ("snapshot") Flink SQL statement to completion, +// synchronously, and returns the whole result set. +// +// It does not reuse the interactive shell's Store/ResultFetcher stack: that +// pipeline silently evicts rows past a 10,000-row cap, drops schema-mismatched +// rows, and treats a missing page token as "done" without checking phase — all +// silent data loss once a script reads stdout. This package's drain loop reports +// each condition instead. +package query + +import ( + "context" + "fmt" + "time" + + flinkgatewayv1 "github.com/confluentinc/ccloud-sdk-go-v2/flink-gateway/v1" + + "github.com/confluentinc/cli/v4/pkg/ccloudv2" + "github.com/confluentinc/cli/v4/pkg/flink/internal/results" + "github.com/confluentinc/cli/v4/pkg/flink/types" + "github.com/confluentinc/cli/v4/pkg/log" +) + +const ( + initialBackoff = 300 * time.Millisecond + maxBackoff = 2 * time.Second +) + +// Options configures a single run. Only Client, EnvironmentId and OrganizationId +// are required. +type Options struct { + Client ccloudv2.GatewayClientInterface + EnvironmentId string + OrganizationId string + + // MaxRows caps how many rows are collected; 0 means no cap. Hitting it sets + // Result.Truncated rather than silently dropping rows. + MaxRows int + + // RequireBounded rejects a statement whose traits say it is unbounded, rather + // than draining a stream that never ends. + RequireBounded bool + + // RefreshToken runs before every gateway call, since the dataplane token is + // short-lived. Failures are logged, not returned — the next gateway call + // surfaces its own error if the token is truly bad. Nil means no refresh. + RefreshToken func() error + + // sleep is swapped out in tests so they do not wait in real time. + sleep func(context.Context, time.Duration) error +} + +// authenticatedClient refreshes the token if configured, mirroring +// Store.authenticatedGatewayClient. +func (opts Options) authenticatedClient() ccloudv2.GatewayClientInterface { + if opts.RefreshToken != nil { + if err := opts.RefreshToken(); err != nil { + log.CliLogger.Warnf("Failed to refresh Flink gateway token: %v", err) + } + } + return opts.Client +} + +// Result is the outcome of a completed run. +type Result struct { + // Statement as the gateway last reported it, including status and traits. + Statement flinkgatewayv1.SqlV1Statement + // Columns is the result schema, in order. + Columns []flinkgatewayv1.ColumnDetails + // Rows is the raw changelog as delivered. Every row is an insert for a bounded + // append-only snapshot; otherwise the caller decides how to materialize it. + Rows []types.StatementResultRow + // Truncated reports that MaxRows stopped the drain before the result set ended. + Truncated bool + // Incomplete reports the gateway stopped giving page tokens while still + // running, so rows may be missing. See drain. + Incomplete bool +} + +// Phase is the statement phase at the end of the run. +func (r *Result) Phase() types.PHASE { + return types.PHASE(r.Statement.Status.GetPhase()) +} + +// UnboundedError is returned when RequireBounded is set and the gateway reports the +// statement produces an unbounded result. +type UnboundedError struct { + StatementName string +} + +func (e *UnboundedError) Error() string { + return fmt.Sprintf(`statement "%s" produces an unbounded result and cannot be run as a snapshot query`, e.StatementName) +} + +// ResultsFetchError distinguishes a failed page fetch from a failed +// statement-status read, so the caller can give a more specific suggestion for +// the gateway's page-retention window. +type ResultsFetchError struct { + Err error +} + +func (e *ResultsFetchError) Error() string { + return e.Err.Error() +} + +func (e *ResultsFetchError) Unwrap() error { + return e.Err +} + +// Run waits for an already-submitted statement to start, then drains every result +// page. The statement must already exist — submitting it is the caller's job, so the +// caller keeps ownership of naming, properties and cleanup. +func Run(ctx context.Context, opts Options, statementName string) (*Result, error) { + if opts.sleep == nil { + opts.sleep = sleepContext + } + + statement, err := await(ctx, opts, statementName) + if err != nil { + return nil, err + } + + result := &Result{Statement: statement} + + traits := statement.Status.GetTraits() + statementTraits := types.StatementTraits{FlinkGatewayV1StatementTraits: &traits} + if isBounded, known := statementTraits.GetIsBounded(); opts.RequireBounded && known && !isBounded { + return result, &UnboundedError{StatementName: statementName} + } + + schema := traits.GetSchema() + result.Columns = schema.GetColumns() + + // DDL/INSERT INTO statements have no schema and nothing to poll; return early + // rather than reporting an empty table. + if len(result.Columns) == 0 { + return result, nil + } + + if err := drain(ctx, opts, statementName, schema, result); err != nil { + return result, err + } + + return result, nil +} + +// await polls until the statement leaves PENDING, so that its traits — the result +// schema and the boundedness flag — are populated. +func await(ctx context.Context, opts Options, statementName string) (flinkgatewayv1.SqlV1Statement, error) { + backoff := initialBackoff + for { + if err := ctx.Err(); err != nil { + return flinkgatewayv1.SqlV1Statement{}, err + } + + statement, err := opts.authenticatedClient().GetStatement(opts.EnvironmentId, statementName, opts.OrganizationId) + if err != nil { + return flinkgatewayv1.SqlV1Statement{}, err + } + + if types.PHASE(statement.Status.GetPhase()) != types.PENDING { + return statement, nil + } + + if err := opts.sleep(ctx, backoff); err != nil { + return flinkgatewayv1.SqlV1Statement{}, err + } + backoff = min(backoff*2, maxBackoff) + } +} + +// drain pulls result pages until there's no next page and the statement is +// terminal. +// +// Phase is read before each page fetch, not after: reading it after could miss a +// completion that happened during the fetch, making a genuinely short page look +// complete. Reading it first means a token-less page is only "final" once nothing +// more could have been added before the fetch — one extra status call per page. +// +// page_token is a positional offset, not a real cursor, so nothing can advance +// past a token-less page. An empty page while still RUNNING just means "not yet"; +// retry the same offset. A page with rows but no token means the gateway can't +// give one — Incomplete says so rather than guessing the run finished. +func drain(ctx context.Context, opts Options, statementName string, schema flinkgatewayv1.SqlV1ResultSchema, result *Result) error { + pageToken := "" + backoff := initialBackoff + + for { + if err := ctx.Err(); err != nil { + return err + } + + statement, err := opts.authenticatedClient().GetStatement(opts.EnvironmentId, statementName, opts.OrganizationId) + if err != nil { + return err + } + result.Statement = statement + terminalBeforeFetch := IsTerminal(types.PHASE(statement.Status.GetPhase())) + + page, err := opts.authenticatedClient().GetStatementResults(opts.EnvironmentId, statementName, opts.OrganizationId, pageToken) + if err != nil { + return &ResultsFetchError{Err: err} + } + + pageResults := page.GetResults() + converted, err := results.ConvertToInternalResults(pageResults.GetData(), schema) + if err != nil { + return err + } + pageRows := converted.GetRows() + result.Rows = append(result.Rows, pageRows...) + + if opts.MaxRows > 0 && len(result.Rows) > opts.MaxRows { + result.Rows = result.Rows[:opts.MaxRows] + result.Truncated = true + return nil + } + + metadata := page.GetMetadata() + var nextPageToken string + if nextUrl := metadata.GetNext(); nextUrl != "" { + nextPageToken, err = ccloudv2.ExtractPageToken(nextUrl) + if err != nil { + return err + } + } + + if nextPageToken != "" { + pageToken = nextPageToken + backoff = initialBackoff + continue + } + + if terminalBeforeFetch { + return nil + } + + if len(pageRows) > 0 { + // terminalBeforeFetch predates the results call, so it could miss a + // completion that happened during it. Re-check once before conceding; + // a failed re-check falls back to Incomplete. + if statement, err := opts.authenticatedClient().GetStatement(opts.EnvironmentId, statementName, opts.OrganizationId); err == nil { + result.Statement = statement + if IsTerminal(types.PHASE(statement.Status.GetPhase())) { + return nil + } + } + result.Incomplete = true + return nil + } + + if err := opts.sleep(ctx, backoff); err != nil { + return err + } + backoff = min(backoff*2, maxBackoff) + } +} + +// IsTerminal reports whether phase is one the statement cannot leave on its own. +// Exported so callers holding a Result can check this without re-deriving the list. +func IsTerminal(phase types.PHASE) bool { + switch phase { + case types.COMPLETED, types.FAILED, types.STOPPED, types.DELETING: + return true + } + return false +} + +func sleepContext(ctx context.Context, duration time.Duration) error { + timer := time.NewTimer(duration) + defer timer.Stop() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} diff --git a/pkg/flink/query/query_test.go b/pkg/flink/query/query_test.go new file mode 100644 index 0000000000..f18a994b9b --- /dev/null +++ b/pkg/flink/query/query_test.go @@ -0,0 +1,383 @@ +package query + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + flinkgatewayv1 "github.com/confluentinc/ccloud-sdk-go-v2/flink-gateway/v1" + + "github.com/confluentinc/cli/v4/pkg/flink/test/mock" + "github.com/confluentinc/cli/v4/pkg/flink/types" +) + +const ( + testEnvironmentId = "env-123456" + testOrganizationId = "org-123456" + testStatementName = "test-statement" +) + +// testOptions builds Options whose sleep is a no-op, so backoff never costs wall time. +func testOptions(client *mock.MockGatewayClientInterface) Options { + return Options{ + Client: client, + EnvironmentId: testEnvironmentId, + OrganizationId: testOrganizationId, + sleep: func(context.Context, time.Duration) error { return nil }, + } +} + +func schema(columnNames ...string) *flinkgatewayv1.SqlV1ResultSchema { + columns := make([]flinkgatewayv1.ColumnDetails, len(columnNames)) + for i, name := range columnNames { + columns[i] = flinkgatewayv1.ColumnDetails{ + Name: name, + Type: flinkgatewayv1.DataType{Type: "VARCHAR", Nullable: false}, + } + } + return &flinkgatewayv1.SqlV1ResultSchema{Columns: &columns} +} + +func statement(phase string, traits *flinkgatewayv1.SqlV1StatementTraits) flinkgatewayv1.SqlV1Statement { + return flinkgatewayv1.SqlV1Statement{ + Name: flinkgatewayv1.PtrString(testStatementName), + Spec: &flinkgatewayv1.SqlV1StatementSpec{Statement: flinkgatewayv1.PtrString("SELECT * FROM t;")}, + Status: &flinkgatewayv1.SqlV1StatementStatus{ + Phase: phase, + Traits: traits, + }, + } +} + +func boundedTraits(columnNames ...string) *flinkgatewayv1.SqlV1StatementTraits { + return &flinkgatewayv1.SqlV1StatementTraits{ + SqlKind: flinkgatewayv1.PtrString("SELECT"), + IsBounded: flinkgatewayv1.PtrBool(true), + IsAppendOnly: flinkgatewayv1.PtrBool(true), + Schema: schema(columnNames...), + } +} + +// page builds a result page. nextPageToken of "" means the gateway reported no next +// page. +func page(nextPageToken string, rows ...[]any) flinkgatewayv1.SqlV1StatementResult { + data := make([]any, len(rows)) + for i, row := range rows { + data[i] = map[string]any{"op": float64(0), "row": row} + } + + metadata := flinkgatewayv1.ResultListMeta{} + if nextPageToken != "" { + next := fmt.Sprintf("https://flink.us-east-1.aws.confluent.cloud/sql/v1/organizations/%s/environments/%s/statements/%s/results?page_token=%s", + testOrganizationId, testEnvironmentId, testStatementName, nextPageToken) + metadata.Next = &next + } + + return flinkgatewayv1.SqlV1StatementResult{ + Metadata: metadata, + Results: &flinkgatewayv1.SqlV1StatementResultResults{Data: &data}, + } +} + +func rowValues(t *testing.T, result *Result) [][]string { + t.Helper() + values := make([][]string, len(result.Rows)) + for i, row := range result.Rows { + fields := make([]string, len(row.GetFields())) + for j, field := range row.GetFields() { + fields[j] = field.ToString() + } + values[i] = fields + } + return values +} + +func TestRunDrainsASinglePage(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + completed := statement("COMPLETED", boundedTraits("id", "status")) + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil).Times(2) + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1", "SHIPPED"}, []any{"2", "PENDING"}), nil) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.Equal(t, [][]string{{"1", "SHIPPED"}, {"2", "PENDING"}}, rowValues(t, result)) + require.False(t, result.Truncated) + require.False(t, result.Incomplete) + require.Equal(t, types.COMPLETED, result.Phase()) +} + +func TestRunDrainsEveryPage(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + running := statement("RUNNING", boundedTraits("id")) + completed := statement("COMPLETED", boundedTraits("id")) + + // Leaves PENDING, produces, completes after the last page — phase read before + // each page, not after. + gomock.InOrder( + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("10", []any{"1"}), nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, "10"). + Return(page("20", []any{"2"}), nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, "20"). + Return(page("", []any{"3"}), nil), + ) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.Equal(t, [][]string{{"1"}, {"2"}, {"3"}}, rowValues(t, result)) + require.False(t, result.Incomplete) +} + +// Unlike the shell (missing token = done), a run must not silently succeed when +// the statement is still running and there's no token left to advance with. +func TestRunFlagsIncompleteWhenPagesStopBeforeStatementDoes(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + running := statement("RUNNING", boundedTraits("id")) + + gomock.InOrder( + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}), nil), + // The re-check before conceding Incomplete: still RUNNING, so it stays Incomplete. + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + ) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.True(t, result.Incomplete) + require.Equal(t, [][]string{{"1"}}, rowValues(t, result)) +} + +// The statement can reach a terminal phase during the GetStatementResults call +// itself (reproduced against a real gateway) — terminalBeforeFetch is stale by +// then, so drain must re-read phase before concluding the read was short. +func TestRunDoesNotFlagIncompleteWhenStatementCompletesDuringTheFinalFetch(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + running := statement("RUNNING", boundedTraits("id")) + completed := statement("COMPLETED", boundedTraits("id")) + + gomock.InOrder( + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}), nil), + // The statement finished during the GetStatementResults call above. + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil), + ) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.False(t, result.Incomplete) + require.Equal(t, [][]string{{"1"}}, rowValues(t, result)) +} + +// An empty page with no token is the gateway saying "nothing yet". Re-requesting the +// same offset is safe, so the loop should keep waiting rather than declaring victory. +func TestRunRetriesEmptyPagesUntilStatementIsTerminal(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + running := statement("RUNNING", boundedTraits("id")) + completed := statement("COMPLETED", boundedTraits("id")) + + gomock.InOrder( + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, "").Return(page(""), nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}), nil), + ) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.False(t, result.Incomplete) + require.Equal(t, [][]string{{"1"}}, rowValues(t, result)) +} + +// Phase must be read before fetching a page, not after — reading it after could +// miss a completion happening in the gap. Pinning the call order here makes a +// regression fail loudly instead of occasionally under-counting rows in prod. +func TestRunReadsStatementStateBeforeFetchingResults(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + running := statement("RUNNING", boundedTraits("id")) + completed := statement("COMPLETED", boundedTraits("id")) + + gomock.InOrder( + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(running, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}), nil), + ) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.False(t, result.Incomplete) + require.Equal(t, [][]string{{"1"}}, rowValues(t, result)) +} + +func TestRunWaitsForPendingStatement(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + pending := statement("PENDING", nil) + completed := statement("COMPLETED", boundedTraits("id")) + + gomock.InOrder( + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(pending, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(pending, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil), + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil), + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}), nil), + ) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.Equal(t, [][]string{{"1"}}, rowValues(t, result)) +} + +func TestRunTruncatesAtMaxRowsAndSaysSo(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + completed := statement("COMPLETED", boundedTraits("id")) + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil).Times(2) + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("10", []any{"1"}, []any{"2"}, []any{"3"}), nil) + + options := testOptions(client) + options.MaxRows = 2 + + result, err := Run(context.Background(), options, testStatementName) + require.NoError(t, err) + require.True(t, result.Truncated) + require.Equal(t, [][]string{{"1"}, {"2"}}, rowValues(t, result)) +} + +// Exactly MaxRows rows is not a truncation — nothing was dropped. +func TestRunDoesNotReportTruncationWhenRowCountEqualsMaxRows(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + completed := statement("COMPLETED", boundedTraits("id")) + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil).Times(2) + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}, []any{"2"}), nil) + + options := testOptions(client) + options.MaxRows = 2 + + result, err := Run(context.Background(), options, testStatementName) + require.NoError(t, err) + require.False(t, result.Truncated) + require.Len(t, result.Rows, 2) +} + +func TestRunRejectsAnUnboundedStatement(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + traits := boundedTraits("id") + traits.IsBounded = flinkgatewayv1.PtrBool(false) + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId). + Return(statement("RUNNING", traits), nil) + + options := testOptions(client) + options.RequireBounded = true + + _, err := Run(context.Background(), options, testStatementName) + var unbounded *UnboundedError + require.ErrorAs(t, err, &unbounded) + require.Equal(t, testStatementName, unbounded.StatementName) +} + +// Without RequireBounded the caller opted into draining whatever comes back. +func TestRunAllowsUnboundedStatementWhenNotRequired(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + traits := boundedTraits("id") + traits.IsBounded = flinkgatewayv1.PtrBool(false) + completed := statement("COMPLETED", traits) + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil).Times(2) + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}), nil) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.Len(t, result.Rows, 1) +} + +// A row whose width doesn't match the schema must fail the run — unlike the shell, +// which drops such rows silently. +func TestRunFailsOnRowSchemaMismatch(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + completed := statement("COMPLETED", boundedTraits("id", "status")) + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil).Times(2) + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(page("", []any{"1"}), nil) + + _, err := Run(context.Background(), testOptions(client), testStatementName) + require.ErrorContains(t, err, "does not match the provided schema") +} + +// DDL and INSERT INTO have no result schema; there is nothing to poll. +func TestRunSkipsResultsForStatementWithoutSchema(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + traits := &flinkgatewayv1.SqlV1StatementTraits{SqlKind: flinkgatewayv1.PtrString("CREATE_TABLE")} + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId). + Return(statement("COMPLETED", traits), nil) + + result, err := Run(context.Background(), testOptions(client), testStatementName) + require.NoError(t, err) + require.Empty(t, result.Rows) + require.Empty(t, result.Columns) +} + +func TestRunStopsOnCancelledContext(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId). + Return(statement("PENDING", nil), nil).AnyTimes() + + ctx, cancel := context.WithCancel(context.Background()) + options := testOptions(client) + options.sleep = func(context.Context, time.Duration) error { + cancel() + return context.Canceled + } + + _, err := Run(ctx, options, testStatementName) + require.ErrorIs(t, err, context.Canceled) +} + +// A page-fetch failure is wrapped so the caller can tell it apart from a failure +// reading the statement itself. +func TestRunWrapsResultsFetchErrors(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + completed := statement("COMPLETED", boundedTraits("id")) + + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId).Return(completed, nil).Times(2) + client.EXPECT().GetStatementResults(testEnvironmentId, testStatementName, testOrganizationId, ""). + Return(flinkgatewayv1.SqlV1StatementResult{}, errors.New("page not found")) + + _, err := Run(context.Background(), testOptions(client), testStatementName) + var resultsErr *ResultsFetchError + require.ErrorAs(t, err, &resultsErr) + require.ErrorContains(t, err, "page not found") +} + +func TestRunPropagatesGatewayErrors(t *testing.T) { + client := mock.NewMockGatewayClientInterface(gomock.NewController(t)) + client.EXPECT().GetStatement(testEnvironmentId, testStatementName, testOrganizationId). + Return(flinkgatewayv1.SqlV1Statement{}, errors.New("unauthorized")) + + _, err := Run(context.Background(), testOptions(client), testStatementName) + require.ErrorContains(t, err, "unauthorized") +} diff --git a/pkg/flink/types/processed_statement.go b/pkg/flink/types/processed_statement.go index f1e3861cd3..a3747c2c2b 100644 --- a/pkg/flink/types/processed_statement.go +++ b/pkg/flink/types/processed_statement.go @@ -19,6 +19,8 @@ const ( RUNNING PHASE = "RUNNING" // More results are available (pagination) COMPLETED PHASE = "COMPLETED" // All results were fetched FAILED PHASE = "FAILED" + STOPPED PHASE = "STOPPED" // The statement was stopped; no more results are coming + DELETING PHASE = "DELETING" // The statement is being deleted; no more results are coming ) // ProcessedStatement Custom Internal type that shall be used internally by the client diff --git a/pkg/flink/types/result_fields.go b/pkg/flink/types/result_fields.go index ff76dd9feb..93cb61a183 100644 --- a/pkg/flink/types/result_fields.go +++ b/pkg/flink/types/result_fields.go @@ -94,6 +94,8 @@ type StatementResultField interface { GetType() StatementResultFieldType ToString() string ToSDKType() any + // ToSerializedValue is implemented in result_fields_serialized.go. + ToSerializedValue() any } type AtomicStatementResultField struct { diff --git a/pkg/flink/types/result_fields_serialized.go b/pkg/flink/types/result_fields_serialized.go new file mode 100644 index 0000000000..8952733b24 --- /dev/null +++ b/pkg/flink/types/result_fields_serialized.go @@ -0,0 +1,104 @@ +package types + +import ( + "math" + "strconv" +) + +// ToSerializedValue renders a field as a value that carries its SQL type into +// `-o json`/`-o yaml`, so a number reads as a number and NULL as null. +// +// Not ToSDKType: that produces the gateway's wire shape (every atom a string, MAP +// as pairs) — correct for the API, wrong for a script. +// +// Two constraints: only native Go types are used (json.Number renders as a bare +// literal in encoding/json but quoted in yaml.v3, which would make the formats +// disagree), and a value that can't be represented natively keeps the gateway's +// exact text rather than becoming a zero or an approximation. + +func (f AtomicStatementResultField) ToSerializedValue() any { + // A NULL arrives as Type Null carrying the literal text "NULL", which is otherwise + // indistinguishable from a VARCHAR containing that word. + if f.Type == Null { + return nil + } + + switch f.Type { + case Boolean: + if value, err := strconv.ParseBool(f.Value); err == nil { + return value + } + case Tinyint, Smallint, Integer: + // These top out at 2^31-1, well inside float64's exact range. BIGINT isn't + // here — see below. + if value, err := strconv.ParseInt(f.Value, 10, 64); err == nil { + return value + } + case Float, Double: + // NaN/±Inf are legal in a DOUBLE column but encoding/json refuses them; failing + // here would take down an otherwise-successful drain. + if value, err := strconv.ParseFloat(f.Value, 64); err == nil && !math.IsNaN(value) && !math.IsInf(value, 0) { + return value + } + } + + // Everything else stays text on purpose: + // - BIGINT exceeds float64's exact range (2^53); JS/jq parse JSON numbers as + // float64, so a bare literal would silently corrupt large values. Same + // reasoning keeps DECIMAL a string. + // - DECIMAL is arbitrary-precision; float64 would round it. + // - DATE/TIME/TIMESTAMP/INTERVAL have no native JSON form. + // - CHAR/VARCHAR/BINARY/VARBINARY are already text. + return f.Value +} + +func (f ArrayStatementResultField) ToSerializedValue() any { + // Length rather than nil, so an empty array serializes as [] and not null. + values := make([]any, len(f.Values)) + for idx, value := range f.Values { + values[idx] = value.ToSerializedValue() + } + return values +} + +func (f MapStatementResultField) ToSerializedValue() any { + // A map with textual keys becomes a JSON object; anything else keeps an explicit + // key/value list instead of ToSDKType's bare positional pairs. + if f.KeyType == Char || f.KeyType == Varchar { + entries := make(map[string]any, len(f.Entries)) + for _, entry := range f.Entries { + entries[entry.Key.ToString()] = entry.Value.ToSerializedValue() + } + return entries + } + + entries := make([]any, len(f.Entries)) + for idx, entry := range f.Entries { + entries[idx] = map[string]any{ + "key": entry.Key.ToSerializedValue(), + "value": entry.Value.ToSerializedValue(), + } + } + return entries +} + +func (f RowStatementResultField) ToSerializedValue() any { + // A ROW carries no field names, so it stays positional. + values := make([]any, len(f.Values)) + for idx, value := range f.Values { + values[idx] = value.ToSerializedValue() + } + return values +} + +func (f StructuredTypeStatementResultField) ToSerializedValue() any { + values := make(map[string]any, len(f.Values)) + for idx, value := range f.Values { + // A short FieldNames would panic here; a partial object beats losing the row. + if idx >= len(f.FieldNames) { + break + } + values[f.FieldNames[idx]] = value.ToSerializedValue() + } + return values +} diff --git a/pkg/flink/types/result_fields_serialized_test.go b/pkg/flink/types/result_fields_serialized_test.go new file mode 100644 index 0000000000..b933dc6816 --- /dev/null +++ b/pkg/flink/types/result_fields_serialized_test.go @@ -0,0 +1,164 @@ +package types + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAtomicToSerializedValue(t *testing.T) { + tests := []struct { + name string + field AtomicStatementResultField + want any + }{ + { + // A NULL and a CHAR "NULL" arrive as the same text, told apart only by type. + name: "a NULL becomes JSON null", + field: AtomicStatementResultField{Type: Null, Value: "NULL"}, + want: nil, + }, + { + name: "a string holding the word NULL stays a string", + field: AtomicStatementResultField{Type: Varchar, Value: "NULL"}, + want: "NULL", + }, + {name: "integer", field: AtomicStatementResultField{Type: Integer, Value: "3065"}, want: int64(3065)}, + {name: "negative integer", field: AtomicStatementResultField{Type: Integer, Value: "-7"}, want: int64(-7)}, + {name: "tinyint", field: AtomicStatementResultField{Type: Tinyint, Value: "12"}, want: int64(12)}, + {name: "smallint", field: AtomicStatementResultField{Type: Smallint, Value: "300"}, want: int64(300)}, + {name: "double", field: AtomicStatementResultField{Type: Double, Value: "52.48"}, want: 52.48}, + {name: "float", field: AtomicStatementResultField{Type: Float, Value: "0.5"}, want: 0.5}, + {name: "boolean true", field: AtomicStatementResultField{Type: Boolean, Value: "true"}, want: true}, + {name: "boolean false", field: AtomicStatementResultField{Type: Boolean, Value: "false"}, want: false}, + { + // float64 can't represent every BIGINT; JS/jq both parse JSON numbers as float64. + name: "bigint stays text so it survives a float64 parser", + field: AtomicStatementResultField{Type: Bigint, Value: "9007199254740993"}, + want: "9007199254740993", + }, + { + name: "decimal stays text so precision is not rounded away", + field: AtomicStatementResultField{Type: Decimal, Value: "1.23"}, + want: "1.23", + }, + { + // encoding/json refuses NaN/±Inf; falling back to text avoids failing the marshal. + name: "NaN falls back to text rather than failing the marshal", + field: AtomicStatementResultField{Type: Double, Value: "NaN"}, + want: "NaN", + }, + { + name: "infinity falls back to text", + field: AtomicStatementResultField{Type: Double, Value: "Infinity"}, + want: "Infinity", + }, + { + name: "a number that does not parse keeps its text instead of becoming zero", + field: AtomicStatementResultField{Type: Integer, Value: "not-a-number"}, + want: "not-a-number", + }, + {name: "timestamp keeps the gateway's rendering", field: AtomicStatementResultField{Type: TimestampWithoutTimeZone, Value: "2026-08-13 10:00:00"}, want: "2026-08-13 10:00:00"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + require.Equal(t, test.want, test.field.ToSerializedValue()) + }) + } +} + +// A BIGINT past 2^53 has to survive the round trip a JSON consumer actually performs. +func TestBigintSurvivesAFloat64Parser(t *testing.T) { + const exact = "9007199254740993" + field := AtomicStatementResultField{Type: Bigint, Value: exact} + + encoded, err := json.Marshal(map[string]any{"big": field.ToSerializedValue()}) + require.NoError(t, err) + + // json.Unmarshal into `any` decodes every number as float64, which is what a + // JavaScript client does too. + var decoded map[string]any + require.NoError(t, json.Unmarshal(encoded, &decoded)) + require.Equal(t, exact, decoded["big"], "BIGINT lost precision through a float64 parser") +} + +func TestArrayToSerializedValue(t *testing.T) { + field := ArrayStatementResultField{ + ElementType: Integer, + Values: []StatementResultField{ + AtomicStatementResultField{Type: Integer, Value: "1"}, + AtomicStatementResultField{Type: Null, Value: "NULL"}, + }, + } + require.Equal(t, []any{int64(1), nil}, field.ToSerializedValue()) +} + +// An empty array must serialize as [] rather than null. +func TestEmptyArraySerializesAsEmptyList(t *testing.T) { + field := ArrayStatementResultField{ElementType: Integer} + encoded, err := json.Marshal(field.ToSerializedValue()) + require.NoError(t, err) + require.Equal(t, "[]", string(encoded)) +} + +func TestMapWithTextKeysBecomesAnObject(t *testing.T) { + field := MapStatementResultField{ + KeyType: Varchar, + ValueType: Integer, + Entries: []MapStatementResultFieldEntry{{ + Key: AtomicStatementResultField{Type: Varchar, Value: "a"}, + Value: AtomicStatementResultField{Type: Integer, Value: "1"}, + }}, + } + require.Equal(t, map[string]any{"a": int64(1)}, field.ToSerializedValue()) +} + +// A JSON object cannot have a non-textual key, so those keep an explicit key/value list. +func TestMapWithNonTextKeysKeepsPairs(t *testing.T) { + field := MapStatementResultField{ + KeyType: Integer, + ValueType: Varchar, + Entries: []MapStatementResultFieldEntry{{ + Key: AtomicStatementResultField{Type: Integer, Value: "1"}, + Value: AtomicStatementResultField{Type: Varchar, Value: "a"}, + }}, + } + require.Equal(t, []any{map[string]any{"key": int64(1), "value": "a"}}, field.ToSerializedValue()) +} + +func TestRowStaysPositional(t *testing.T) { + field := RowStatementResultField{ + ElementTypes: []StatementResultFieldType{Integer, Varchar}, + Values: []StatementResultField{ + AtomicStatementResultField{Type: Integer, Value: "1"}, + AtomicStatementResultField{Type: Varchar, Value: "a"}, + }, + } + require.Equal(t, []any{int64(1), "a"}, field.ToSerializedValue()) +} + +func TestStructuredTypeUsesFieldNames(t *testing.T) { + field := StructuredTypeStatementResultField{ + FieldNames: []string{"id", "label"}, + FieldTypes: []StatementResultFieldType{Integer, Varchar}, + Values: []StatementResultField{ + AtomicStatementResultField{Type: Integer, Value: "1"}, + AtomicStatementResultField{Type: Varchar, Value: "a"}, + }, + } + require.Equal(t, map[string]any{"id": int64(1), "label": "a"}, field.ToSerializedValue()) +} + +// Fewer names than values must not panic; a partial object beats losing the row. +func TestStructuredTypeWithShortFieldNamesDoesNotPanic(t *testing.T) { + field := StructuredTypeStatementResultField{ + FieldNames: []string{"id"}, + Values: []StatementResultField{ + AtomicStatementResultField{Type: Integer, Value: "1"}, + AtomicStatementResultField{Type: Varchar, Value: "dropped"}, + }, + } + require.Equal(t, map[string]any{"id": int64(1)}, field.ToSerializedValue()) +} diff --git a/pkg/flink/types/statement_traits.go b/pkg/flink/types/statement_traits.go index 18c5451ad7..cb165602df 100644 --- a/pkg/flink/types/statement_traits.go +++ b/pkg/flink/types/statement_traits.go @@ -28,6 +28,29 @@ func (s *StatementTraits) GetUpsertColumns() *[]int32 { return nil } +// GetIsBounded reports whether the statement produces a finite result set. The +// second return value is false when the trait is absent (before leaving PENDING). +func (s *StatementTraits) GetIsBounded() (bool, bool) { + if s.FlinkGatewayV1StatementTraits != nil && s.FlinkGatewayV1StatementTraits.IsBounded != nil { + return s.FlinkGatewayV1StatementTraits.GetIsBounded(), true + } else if s.CmfStatementTraits != nil && s.CmfStatementTraits.IsBounded != nil { + return s.CmfStatementTraits.GetIsBounded(), true + } + return false, false +} + +// GetIsAppendOnly reports whether the statement only ever emits insertions (the +// changelog and materialized table are then the same thing). The second return +// value is false when the trait is absent. +func (s *StatementTraits) GetIsAppendOnly() (bool, bool) { + if s.FlinkGatewayV1StatementTraits != nil && s.FlinkGatewayV1StatementTraits.IsAppendOnly != nil { + return s.FlinkGatewayV1StatementTraits.GetIsAppendOnly(), true + } else if s.CmfStatementTraits != nil && s.CmfStatementTraits.IsAppendOnly != nil { + return s.CmfStatementTraits.GetIsAppendOnly(), true + } + return false, false +} + func (s *StatementTraits) GetColumnNames() []string { var columnNames []string if s.FlinkGatewayV1StatementTraits != nil { diff --git a/test/fixtures/input/query/select.sql b/test/fixtures/input/query/select.sql new file mode 100644 index 0000000000..3553d42c52 --- /dev/null +++ b/test/fixtures/input/query/select.sql @@ -0,0 +1 @@ +SELECT order_id, status FROM orders LIMIT 2; \ No newline at end of file diff --git a/test/fixtures/output/help.golden b/test/fixtures/output/help.golden index ab91e2ed10..3f176f5929 100644 --- a/test/fixtures/output/help.golden +++ b/test/fixtures/output/help.golden @@ -32,6 +32,7 @@ Available Commands: plugin Manage Confluent plugins. prompt Add Confluent CLI context to your terminal prompt. provider-integration Manage Confluent Cloud provider integrations. + query Run a bounded Flink SQL query and print its results. rtce Manage Real Time Context Engine. schema-registry Manage Schema Registry. service-quota Look up Confluent Cloud service quota limits. diff --git a/test/fixtures/output/query/changelog.golden b/test/fixtures/output/query/changelog.golden new file mode 100644 index 0000000000..ea981b9b7c --- /dev/null +++ b/test/fixtures/output/query/changelog.golden @@ -0,0 +1,7 @@ +Warning: this statement emits updates and deletions. The rows below are the raw changelog, not a materialized table. ++-----------+----+ +| Operation | id | ++-----------+----+ +| +I | 1 | +| +U | 1 | ++-----------+----+ diff --git a/test/fixtures/output/query/failed.golden b/test/fixtures/output/query/failed.golden new file mode 100644 index 0000000000..09f9b8f7d3 --- /dev/null +++ b/test/fixtures/output/query/failed.golden @@ -0,0 +1,4 @@ +Error:\ statement\ "cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"\ failed:\ Something\ went\ wrong\ compiling\ the\ statement\ +\ +Suggestions:\ +\ \ \ \ Inspect\ the\ failure\ with\ `confluent\ flink\ statement\ exception\ list\ cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}`\.\ diff --git a/test/fixtures/output/query/help.golden b/test/fixtures/output/query/help.golden new file mode 100644 index 0000000000..d84e5e0af0 --- /dev/null +++ b/test/fixtures/output/query/help.golden @@ -0,0 +1,50 @@ +Run a bounded (snapshot) Flink SQL query, block until it finishes, and print the complete result set. + +The SQL can be given as "--sql" or as a positional argument, but not both. + +Unlike statement creation, which submits a statement and returns immediately, this command waits for every result page and exits with a non-zero status if the statement fails. It is intended for scripting and one-shot queries against a bounded (point-in-time) result set. + +With "-o json" or "-o yaml", output defaults to an envelope carrying the column schema alongside the rows, since the rows on their own carry no type information. Pass "--raw" for a bare array of row objects instead. + +Usage: + confluent query [sql] [flags] + +Examples: +Run a bounded query in the current compute pool and print the rows as a table. + + $ confluent query --sql "SELECT * FROM orders LIMIT 10;" + +Run a bounded query against Kafka cluster "my-cluster" and emit JSON for a script to consume. + + $ confluent query --sql "SELECT status, COUNT(*) FROM orders GROUP BY status;" --compute-pool lfcp-123456 --database my-cluster --output json + +Emit a bare JSON array of rows, with no envelope, for a script that only wants the data. + + $ confluent query --sql "SELECT * FROM orders LIMIT 10;" --output json --raw + +Pass the SQL as a positional argument instead of "--sql". + + $ confluent query "SELECT * FROM orders LIMIT 10;" + +Flags: + --sql string The Flink SQL statement. Alternatively, pass it as a positional argument or with "-f". + -f, --file string Path to a file containing the Flink SQL statement. Alternatively, pass the SQL with "--sql" or as a positional argument. + --compute-pool string Flink compute pool ID. + --service-account string Service account ID. + --database string The database which will be used as the default database. When using Kafka, this is the cluster ID. + --cluster string Alias for "--database". Unlike "--database", this also sets the CLI's active Kafka cluster context, the same as it does on every other command. + --property strings A mechanism to pass properties in the form key=value when creating a Flink statement. + --wait-timeout duration Maximum time to wait for the query to finish before giving up. (default 10m0s) + --max-rows int Stop fetching and discard the rest after this many rows, or 0 to fetch every row. Client-side only: rows past the limit are still produced by the query. + --raw Emit the rows as a bare array with no envelope. Requires "-o json" or "-o yaml". + --environment string Environment ID. + --catalog string Alias for "--environment". + --context string CLI context name. + -o, --output string Specify the output format as "human", "json", or "yaml". (default "human") + --cloud string Specify the cloud provider as "aws", "azure", or "gcp". + --region string Cloud region for Flink (use "confluent flink region list" to see all). + +Global Flags: + -h, --help Show help for this command. + --unsafe-trace Equivalent to -vvvv, but also log HTTP requests and responses which might contain plaintext secrets. + -v, --verbose count Increase verbosity (-v for warn, -vv for info, -vvv for debug, -vvvv for trace). diff --git a/test/fixtures/output/query/max-rows.golden b/test/fixtures/output/query/max-rows.golden new file mode 100644 index 0000000000..d91719447c --- /dev/null +++ b/test/fixtures/output/query/max-rows.golden @@ -0,0 +1,8 @@ +Warning:\ stopped\ after\ 2\ rows\ because\ of\ the\ "\-\-max\-rows"\ flag\.\ The\ result\ set\ below\ is\ truncated\.\ +\+\-\-\-\-\+\ +\|\ id\ \|\ +\+\-\-\-\-\+\ +\|\ 1\ \ \|\ +\|\ 2\ \ \|\ +\+\-\-\-\-\+\ +Stopped\ statement\ "cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"\.\ diff --git a/test/fixtures/output/query/missing-sql.golden b/test/fixtures/output/query/missing-sql.golden new file mode 100644 index 0000000000..db76588c14 --- /dev/null +++ b/test/fixtures/output/query/missing-sql.golden @@ -0,0 +1 @@ +Error: the SQL statement is required: pass it as a positional argument, with the `--sql` flag, or with the `--file` flag diff --git a/test/fixtures/output/query/multi-page.golden b/test/fixtures/output/query/multi-page.golden new file mode 100644 index 0000000000..c54cbe080a --- /dev/null +++ b/test/fixtures/output/query/multi-page.golden @@ -0,0 +1,7 @@ ++----+ +| id | ++----+ +| 1 | +| 2 | +| 3 | ++----+ diff --git a/test/fixtures/output/query/no-rows.golden b/test/fixtures/output/query/no-rows.golden new file mode 100644 index 0000000000..7771d9a8e1 --- /dev/null +++ b/test/fixtures/output/query/no-rows.golden @@ -0,0 +1 @@ +The\ query\ returned\ no\ rows\.\ Statement\ "cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"\ is\ in\ phase\ COMPLETED\.\ diff --git a/test/fixtures/output/query/select-json.golden b/test/fixtures/output/query/select-json.golden new file mode 100644 index 0000000000..bff6489fbb --- /dev/null +++ b/test/fixtures/output/query/select-json.golden @@ -0,0 +1,29 @@ +\{\ +\ \ "statement_name":\ "cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}",\ +\ \ "engine":\ "snapshot",\ +\ \ "phase":\ "COMPLETED",\ +\ \ "columns":\ \[\ +\ \ \ \ \{\ +\ \ \ \ \ \ "name":\ "order_id",\ +\ \ \ \ \ \ "type":\ "INTEGER"\ +\ \ \ \ \},\ +\ \ \ \ \{\ +\ \ \ \ \ \ "name":\ "status",\ +\ \ \ \ \ \ "type":\ "VARCHAR"\ +\ \ \ \ \}\ +\ \ \],\ +\ \ "rows":\ \[\ +\ \ \ \ \{\ +\ \ \ \ \ \ "order_id":\ 1021,\ +\ \ \ \ \ \ "status":\ "SHIPPED"\ +\ \ \ \ \},\ +\ \ \ \ \{\ +\ \ \ \ \ \ "order_id":\ 1044,\ +\ \ \ \ \ \ "status":\ "PENDING"\ +\ \ \ \ \}\ +\ \ \],\ +\ \ "row_count":\ 2,\ +\ \ "truncated":\ false,\ +\ \ "incomplete":\ false,\ +\ \ "append_only":\ true\ +\}\ diff --git a/test/fixtures/output/query/select-raw.golden b/test/fixtures/output/query/select-raw.golden new file mode 100644 index 0000000000..5773b17f01 --- /dev/null +++ b/test/fixtures/output/query/select-raw.golden @@ -0,0 +1,10 @@ +[ + { + "order_id": 1021, + "status": "SHIPPED" + }, + { + "order_id": 1044, + "status": "PENDING" + } +] diff --git a/test/fixtures/output/query/select-yaml.golden b/test/fixtures/output/query/select-yaml.golden new file mode 100644 index 0000000000..1e4ec96451 --- /dev/null +++ b/test/fixtures/output/query/select-yaml.golden @@ -0,0 +1,17 @@ +statement_name:\ cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}\ +engine:\ snapshot\ +phase:\ COMPLETED\ +columns:\ +\ \ \ \ \-\ name:\ order_id\ +\ \ \ \ \ \ type:\ INTEGER\ +\ \ \ \ \-\ name:\ status\ +\ \ \ \ \ \ type:\ VARCHAR\ +rows:\ +\ \ \ \ \-\ order_id:\ 1021\ +\ \ \ \ \ \ status:\ SHIPPED\ +\ \ \ \ \-\ order_id:\ 1044\ +\ \ \ \ \ \ status:\ PENDING\ +row_count:\ 2\ +truncated:\ false\ +incomplete:\ false\ +append_only:\ true\ diff --git a/test/fixtures/output/query/select.golden b/test/fixtures/output/query/select.golden new file mode 100644 index 0000000000..06d6f6bae6 --- /dev/null +++ b/test/fixtures/output/query/select.golden @@ -0,0 +1,6 @@ ++----------+---------+ +| order_id | status | ++----------+---------+ +| 1021 | SHIPPED | +| 1044 | PENDING | ++----------+---------+ diff --git a/test/fixtures/output/query/snapshot-mode-override.golden b/test/fixtures/output/query/snapshot-mode-override.golden new file mode 100644 index 0000000000..32c4f59a5d --- /dev/null +++ b/test/fixtures/output/query/snapshot-mode-override.golden @@ -0,0 +1,4 @@ +Error: the `--property` flag must not set `sql.snapshot.mode` + +Suggestions: + This command only supports a snapshot (point-in-time, append-only) read, so this property is fixed and cannot be overridden. diff --git a/test/fixtures/output/query/unbounded.golden b/test/fixtures/output/query/unbounded.golden new file mode 100644 index 0000000000..f7b9be15c3 --- /dev/null +++ b/test/fixtures/output/query/unbounded.golden @@ -0,0 +1,5 @@ +Stopped\ statement\ "cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"\.\ +Error:\ statement\ "cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"\ produces\ an\ unbounded\ result\ and\ cannot\ be\ run\ as\ a\ snapshot\ query\ +\ +Suggestions:\ +\ \ \ \ Bound\ the\ query\ with\ a\ LIMIT\ clause\ or\ a\ time\ predicate\ and\ run\ `confluent\ query`\ again\.\ Statement\ "cli-\d{4}-\d{2}-\d{2}-\d{6}-[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"\ was\ stopped\.\ diff --git a/test/query_test.go b/test/query_test.go new file mode 100644 index 0000000000..03fd237b49 --- /dev/null +++ b/test/query_test.go @@ -0,0 +1,60 @@ +package test + +// TestQuery covers `confluent query` end to end against the mock Flink gateway. +// Scenarios are keyed off the `--sql` text; see buildQueryTestFixture in +// test/test-server/flink_gateway_router.go for what each one returns. +func (s *CLITestSuite) TestQuery() { + tests := []CLITest{ + {args: "query --help", fixture: "query/help.golden"}, + + {args: `query --sql "SELECT order_id, status FROM orders LIMIT 2;" --compute-pool lfcp-123456 --service-account sa-123456 --database lkc-123456`, fixture: "query/select.golden"}, + + // Positional SQL instead of --sql, and a multi-page result set. + {args: `query "SELECT id FROM multi_page_table;" --compute-pool lfcp-123456 --service-account sa-123456`, fixture: "query/multi-page.golden"}, + + // -o json / -o yaml default to the schema+rows envelope; the statement name it + // carries is random per run (types.GenerateStatementName), so these are regexes. + {args: `query --sql "SELECT order_id, status FROM orders LIMIT 2;" --compute-pool lfcp-123456 --service-account sa-123456 -o json`, fixture: "query/select-json.golden", regex: true}, + {args: `query --sql "SELECT order_id, status FROM orders LIMIT 2;" --compute-pool lfcp-123456 --service-account sa-123456 -o yaml`, fixture: "query/select-yaml.golden", regex: true}, + + // --raw drops the envelope (and the statement name with it), so this one is exact. + {args: `query --sql "SELECT order_id, status FROM orders LIMIT 2;" --compute-pool lfcp-123456 --service-account sa-123456 -o json --raw`, fixture: "query/select-raw.golden"}, + + // --max-rows stops the drain early. Truncated is one of the two conditions that + // makes runQuery's deferred cleanup stop the statement, so the name shows up again + // in the "Stopped statement" message. + {args: `query --sql "SELECT id FROM many_rows;" --compute-pool lfcp-123456 --service-account sa-123456 --max-rows 2`, fixture: "query/max-rows.golden", regex: true}, + + // Non-append-only: an Operation column and a changelog warning, no stop (the + // statement already reached a terminal phase on its own). + {args: `query --sql "SELECT * FROM changelog;" --compute-pool lfcp-123456 --service-account sa-123456`, fixture: "query/changelog.golden"}, + + // RequireBounded rejects an unbounded statement before ever touching results. + {args: `query --sql "SELECT * FROM unbounded_stream;" --compute-pool lfcp-123456 --service-account sa-123456`, fixture: "query/unbounded.golden", regex: true, exitCode: 1}, + + // The statement itself fails server-side. + {args: `query --sql "SELECT * FROM will_fail;" --compute-pool lfcp-123456 --service-account sa-123456`, fixture: "query/failed.golden", regex: true, exitCode: 1}, + + // DDL has no result schema, so Run() returns before ever calling GetStatementResults. + {args: `query --sql "CREATE TABLE t (id INT);" --compute-pool lfcp-123456 --service-account sa-123456`, fixture: "query/no-rows.golden", regex: true}, + + // sql.snapshot.mode can't be overridden; this fails in buildQueryProperties, before + // a statement is ever created, so there's no random name in the output. + {args: `query --sql "SELECT 1;" --compute-pool lfcp-123456 --service-account sa-123456 --property sql.snapshot.mode=at-earliest`, fixture: "query/snapshot-mode-override.golden", exitCode: 1}, + + // -f/--file reads the same SQL as the happy path, so it produces the same table. + {args: "query -f test/fixtures/input/query/select.sql --compute-pool lfcp-123456 --service-account sa-123456 --database lkc-123456", fixture: "query/select.golden"}, + + // --catalog/--cluster are aliases for --environment/--database and don't change + // what's printed. + {args: `query --sql "SELECT order_id, status FROM orders LIMIT 2;" --compute-pool lfcp-123456 --service-account sa-123456 --catalog env-596 --cluster lkc-123456`, fixture: "query/select.golden"}, + + // resolveSQL requires exactly one of --sql, --file or the positional argument. + {args: "query --compute-pool lfcp-123456 --service-account sa-123456", fixture: "query/missing-sql.golden", exitCode: 1}, + } + + for _, test := range tests { + test.login = "cloud" + s.runIntegrationTest(test) + } +} diff --git a/test/test-server/flink_gateway_router.go b/test/test-server/flink_gateway_router.go index affabe13a4..b09ad67740 100644 --- a/test/test-server/flink_gateway_router.go +++ b/test/test-server/flink_gateway_router.go @@ -4,7 +4,9 @@ import ( "encoding/json" "fmt" "net/http" + "strconv" "strings" + "sync" "testing" "time" @@ -24,6 +26,7 @@ var flinkGatewayRoutes = []route{ {"/sql/v1/organizations/{organization_id}/environments/{environment}/statements", handleSqlEnvironmentsEnvironmentStatements}, {"/sql/v1/organizations/{organization_id}/environments/{environment}/statements/{statement}", handleSqlEnvironmentsEnvironmentStatementsStatement}, {"/sql/v1/organizations/{organization_id}/environments/{environment}/statements/{statement}/exceptions", handleSqlEnvironmentsEnvironmentStatementExceptions}, + {"/sql/v1/organizations/{organization_id}/environments/{environment}/statements/{statement}/results", handleSqlEnvironmentsEnvironmentStatementsStatementResults}, {"/sql/v1/organizations/{organization_id}/environments/{environment_id}/connections", handleSqlEnvironmentsEnvironmentConnections}, {"/sql/v1/organizations/{organization_id}/environments/{environment_id}/connections/{connection}", handleSqlEnvironmentsEnvironmentConnectionsConnection}, {"/sql/v1/organizations/{organization_id}/environments/{environment_id}/databases/{kafka_cluster_id}/materialized-tables", handleSqlMaterializedTables}, @@ -184,6 +187,13 @@ func handleSqlEnvironmentsEnvironmentStatements(t *testing.T) http.HandlerFunc { statement.Status = &flinkgatewayv1.SqlV1StatementStatus{Phase: "PENDING"} + if strings.HasPrefix(statement.GetName(), queryTestStatementPrefix) { + fixture := buildQueryTestFixture(statement.GetName(), statement.Spec.GetStatement()) + queryTestFixturesMu.Lock() + queryTestFixtures[statement.GetName()] = fixture + queryTestFixturesMu.Unlock() + } + err = json.NewEncoder(w).Encode(statement) require.NoError(t, err) } @@ -209,6 +219,134 @@ func handleSqlEnvironmentsEnvironmentStatementExceptions(t *testing.T) http.Hand } } +// queryTestStatementPrefix marks a statement created by `confluent query` +// (types.GenerateStatementName always produces this prefix). Fixtures below are +// keyed by statement name and only apply to names with this prefix, so they never +// affect the fixed-name statements the rest of this file's tests use. +const queryTestStatementPrefix = "cli-" + +// queryTestFixture is the mock's stand-in for a real gateway statement plus its +// paginated result set. Built once from the submitted SQL text (there is no other +// per-invocation signal available, since `confluent query` never lets a test choose +// the statement name) and served back unchanged for every subsequent GetStatement/ +// GetStatementResults call the drain loop makes. +type queryTestFixture struct { + statement flinkgatewayv1.SqlV1Statement + pages [][]map[string]any +} + +var ( + queryTestFixturesMu sync.Mutex + queryTestFixtures = make(map[string]*queryTestFixture) +) + +func queryColumn(name, sqlType string) flinkgatewayv1.ColumnDetails { + return flinkgatewayv1.ColumnDetails{Name: name, Type: flinkgatewayv1.DataType{Type: sqlType}} +} + +func queryRow(op int, values ...string) map[string]any { + row := make([]any, len(values)) + for i, v := range values { + row[i] = v + } + return map[string]any{"op": op, "row": row} +} + +// buildQueryTestFixture maps a fixed set of `--sql` values, used by test/query_test.go, +// to a scripted statement lifecycle. Add a case here for every new integration test +// scenario that needs the drain loop to actually run. +func buildQueryTestFixture(name, sql string) *queryTestFixture { + traits := &flinkgatewayv1.SqlV1StatementTraits{IsBounded: flinkgatewayv1.PtrBool(true), IsAppendOnly: flinkgatewayv1.PtrBool(true)} + phase := "COMPLETED" + detail := "SQL statement is completed" + var pages [][]map[string]any + + switch sql { + case "SELECT order_id, status FROM orders LIMIT 2;": + traits.Schema = &flinkgatewayv1.SqlV1ResultSchema{Columns: &[]flinkgatewayv1.ColumnDetails{ + queryColumn("order_id", "INTEGER"), + queryColumn("status", "VARCHAR"), + }} + pages = [][]map[string]any{{queryRow(0, "1021", "SHIPPED"), queryRow(0, "1044", "PENDING")}} + case "SELECT id FROM multi_page_table;": + traits.Schema = &flinkgatewayv1.SqlV1ResultSchema{Columns: &[]flinkgatewayv1.ColumnDetails{queryColumn("id", "INTEGER")}} + pages = [][]map[string]any{ + {queryRow(0, "1"), queryRow(0, "2")}, + {queryRow(0, "3")}, + } + case "SELECT id FROM many_rows;": + traits.Schema = &flinkgatewayv1.SqlV1ResultSchema{Columns: &[]flinkgatewayv1.ColumnDetails{queryColumn("id", "INTEGER")}} + pages = [][]map[string]any{{ + queryRow(0, "1"), queryRow(0, "2"), queryRow(0, "3"), queryRow(0, "4"), queryRow(0, "5"), + }} + case "SELECT * FROM changelog;": + traits.IsAppendOnly = flinkgatewayv1.PtrBool(false) + traits.Schema = &flinkgatewayv1.SqlV1ResultSchema{Columns: &[]flinkgatewayv1.ColumnDetails{queryColumn("id", "INTEGER")}} + pages = [][]map[string]any{{queryRow(0, "1"), queryRow(2, "1")}} + case "SELECT * FROM unbounded_stream;": + traits.IsBounded = flinkgatewayv1.PtrBool(false) + case "SELECT * FROM will_fail;": + traits = nil + phase = "FAILED" + detail = "Something went wrong compiling the statement" + case "CREATE TABLE t (id INT);": + traits = nil + default: + traits.Schema = &flinkgatewayv1.SqlV1ResultSchema{Columns: &[]flinkgatewayv1.ColumnDetails{queryColumn("id", "INTEGER")}} + pages = [][]map[string]any{{queryRow(0, "1")}} + } + + status := &flinkgatewayv1.SqlV1StatementStatus{Phase: phase, Detail: flinkgatewayv1.PtrString(detail), Traits: traits} + return &queryTestFixture{ + statement: flinkgatewayv1.SqlV1Statement{ + Name: flinkgatewayv1.PtrString(name), + Spec: &flinkgatewayv1.SqlV1StatementSpec{ + Statement: flinkgatewayv1.PtrString(sql), + ComputePoolId: flinkgatewayv1.PtrString(validFlinkStatementComputePoolId), + }, + Status: status, + Metadata: &flinkgatewayv1.StatementObjectMeta{CreatedAt: flinkgatewayv1.PtrTime(time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC))}, + }, + pages: pages, + } +} + +func handleSqlEnvironmentsEnvironmentStatementsStatementResults(t *testing.T) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + name := mux.Vars(r)["statement"] + + queryTestFixturesMu.Lock() + fixture, ok := queryTestFixtures[name] + queryTestFixturesMu.Unlock() + if !ok || len(fixture.pages) == 0 { + w.WriteHeader(http.StatusNotFound) + err := writeError(w, fmt.Sprintf(`statement "%s" has no results`, name)) + require.NoError(t, err) + return + } + + pageIndex := 0 + if token := r.URL.Query().Get("page_token"); token != "" { + parsed, err := strconv.Atoi(token) + require.NoError(t, err) + pageIndex = parsed + } + + data := make([]any, len(fixture.pages[pageIndex])) + for i, row := range fixture.pages[pageIndex] { + data[i] = row + } + + result := flinkgatewayv1.SqlV1StatementResult{Results: &flinkgatewayv1.SqlV1StatementResultResults{Data: &data}} + if pageIndex+1 < len(fixture.pages) { + result.Metadata.SetNext(fmt.Sprintf("%s?page_token=%d", r.URL.Path, pageIndex+1)) + } + + err := json.NewEncoder(w).Encode(result) + require.NoError(t, err) + } +} + // Handler for "/sql/v1/organizations/{organization_id}/environments/{environment_id}/statements/{statement_name}" func handleSqlEnvironmentsEnvironmentStatementsStatement(t *testing.T) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { @@ -226,6 +364,18 @@ func handleSqlEnvironmentsEnvironmentStatementsStatement(t *testing.T) http.Hand func handleStatementGet(t *testing.T) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { + name := mux.Vars(r)["statement"] + if strings.HasPrefix(name, queryTestStatementPrefix) { + queryTestFixturesMu.Lock() + fixture, ok := queryTestFixtures[name] + queryTestFixturesMu.Unlock() + if ok { + err := json.NewEncoder(w).Encode(fixture.statement) + require.NoError(t, err) + return + } + } + statement := flinkgatewayv1.SqlV1Statement{ Name: flinkgatewayv1.PtrString(mux.Vars(r)["statement"]), Spec: &flinkgatewayv1.SqlV1StatementSpec{ @@ -263,6 +413,16 @@ func handleStatementUpdate(t *testing.T) http.HandlerFunc { principal := req.Spec.GetPrincipal() computePool := req.Spec.GetComputePoolId() + // The real gateway rejects a body omitting the SQL text; a mock more + // permissive than that let a broken stop path pass every test and only + // fail against staging. + if req.Spec.GetStatement() == "" { + w.WriteHeader(http.StatusBadRequest) + err = writeError(w, "Request is malformed: Violations: Statement is nil or empty") + require.NoError(t, err) + return + } + // Handle the stop case, principal and computerPool shouldn't matter if stopped { w.WriteHeader(http.StatusAccepted)