Merge #86484

86484: sql,cli: misc statement bundle fixes r=yuzefovich,cucaroach a=knz

Fixes #86472.
Fixes #83546.

Replaces #86303.

Release justification: bug fixes


Co-authored-by: Raphael 'kena' Poss <[email protected]>

pkg/cli/clisqlclient/conn.go

@@ -380,6 +380,12 @@ func (c *sqlConn) checkServerMetadata(ctx context.Context) error {
 		return nil
 	}
 
+	// The checks below use statements that don't support being
+	// prepared. So disable result type inference for the duration
+	// of these checks.
+	defer func(prev bool) { c.alwaysInferResultTypes = prev }(c.alwaysInferResultTypes)
+	c.alwaysInferResultTypes = false
+
 	_, newServerVersion, newClusterID, err := c.GetServerMetadata(ctx)
 	if c.conn.IsClosed() {
 		return MarkWithConnectionClosed(err)

pkg/cli/clisqlclient/statement_diag.go

@@ -16,6 +16,8 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/errors"
@@ -60,10 +62,32 @@ func stmtDiagListBundlesInternal(ctx context.Context, conn Conn) ([]StmtDiagBund
 		} else if err != nil {
 			return nil, err
 		}
+		i, ok := vals[0].(int64)
+		if !ok {
+			// We're arriving in this function via the interactive shell,
+			// with result type inference disabled.
+			// The value has been read as a string.
+			i, err = strconv.ParseInt(vals[0].(string), 10, 64)
+			if err != nil {
+				return nil, err
+			}
+		}
+		d, ok := vals[2].(time.Time)
+		if !ok {
+			// We're arriving in this function via the interactive shell,
+			// with result type inference disabled.
+			// The value has been read as a string.
+			ts := vals[2].(string)
+			ts = strings.TrimSuffix(ts, "+00")
+			d, err = time.Parse("2006-01-02 15:04:05.999999", ts)
+			if err != nil {
+				return nil, err
+			}
+		}
 		info := StmtDiagBundleInfo{
-			ID:          vals[0].(int64),
+			ID:          i,
 			Statement:   vals[1].(string),
-			CollectedAt: vals[2].(time.Time),
+			CollectedAt: d,
 		}
 		result = append(result, info)
 	}

pkg/cli/demo.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cli/clierrorplus"
 	"github.com/cockroachdb/cockroach/pkg/cli/cliflagcfg"
 	"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
+	"github.com/cockroachdb/cockroach/pkg/cli/clisqlclient"
 	"github.com/cockroachdb/cockroach/pkg/cli/democluster"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -176,7 +177,7 @@ func checkDemoConfiguration(
 		}
 
 		// If the geo-partitioned replicas flag was given and the nodes have changed, throw an error.
-		if f.Lookup(cliflags.DemoNodes.Name).Changed {
+		if fl := f.Lookup(cliflags.DemoNodes.Name); fl != nil && fl.Changed {
 			if demoCtx.NumNodes != 9 {
 				return nil, errors.Newf("--nodes with a value different from 9 cannot be used with %s", geoFlag)
 			}
@@ -205,7 +206,15 @@ func checkDemoConfiguration(
 	return gen, nil
 }
 
-func runDemo(cmd *cobra.Command, gen workload.Generator) (resErr error) {
+func runDemo(cmd *cobra.Command, gen workload.Generator) error {
+	return runDemoInternal(cmd, gen, func(context.Context, clisqlclient.Conn) error { return nil })
+}
+
+func runDemoInternal(
+	cmd *cobra.Command,
+	gen workload.Generator,
+	extraInit func(context.Context, clisqlclient.Conn) error,
+) (resErr error) {
 	closeFn, err := sqlCtx.Open(os.Stdin)
 	if err != nil {
 		return err
@@ -355,5 +364,10 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) (resErr error) {
 	defer func() { resErr = errors.CombineErrors(resErr, conn.Close()) }()
 
 	sqlCtx.ShellCtx.ParseURL = clienturl.MakeURLParserFn(cmd, cliCtx.clientOpts)
+
+	if err := extraInit(ctx, conn); err != nil {
+		return err
+	}
+
 	return sqlCtx.Run(ctx, conn)
 }

pkg/cli/flags.go

@@ -515,7 +515,7 @@ func init() {
 	}
 
 	// Flags that apply to commands that start servers.
-	telemetryEnabledCmds := append(serverCmds, demoCmd)
+	telemetryEnabledCmds := append(serverCmds, demoCmd, statementBundleRecreateCmd)
 	telemetryEnabledCmds = append(telemetryEnabledCmds, demoCmd.Commands()...)
 	for _, cmd := range telemetryEnabledCmds {
 		// Report flag usage for server commands in telemetry. We do this
@@ -698,7 +698,10 @@ func init() {
 	sqlCmds = append(sqlCmds, importCmds...)
 	sqlCmds = append(sqlCmds, userFileCmds...)
 	for _, cmd := range sqlCmds {
-		clientflags.AddSQLFlags(cmd, &cliCtx.clientOpts, sqlCtx, cmd == sqlShellCmd, cmd == demoCmd)
+		clientflags.AddSQLFlags(cmd, &cliCtx.clientOpts, sqlCtx,
+			cmd == sqlShellCmd, /* isShell */
+			cmd == demoCmd || cmd == statementBundleRecreateCmd, /* isDemo */
+		)
 	}
 
 	// Make the non-SQL client commands also recognize --url in strict SSL mode
@@ -744,11 +747,11 @@ func init() {
 	}
 
 	// demo command.
-	{
+	for _, cmd := range []*cobra.Command{demoCmd, statementBundleRecreateCmd} {
 		// We use the persistent flag set so that the flags apply to every
 		// workload sub-command. This enables e.g.
 		// ./cockroach demo movr --nodes=3.
-		f := demoCmd.PersistentFlags()
+		f := cmd.PersistentFlags()
 
 		cliflagcfg.IntFlag(f, &demoCtx.NumNodes, cliflags.DemoNodes)
 		cliflagcfg.BoolFlag(f, &demoCtx.RunWorkload, cliflags.RunDemoWorkload)
@@ -772,11 +775,6 @@ func init() {
 		_ = f.MarkHidden(cliflags.DemoMultitenant.Name)
 
 		cliflagcfg.BoolFlag(f, &demoCtx.SimulateLatency, cliflags.Global)
-		// The --empty flag is only valid for the top level demo command,
-		// so we use the regular flag set.
-		cliflagcfg.BoolFlag(demoCmd.Flags(), &demoCtx.NoExampleDatabase, cliflags.UseEmptyDatabase)
-		_ = f.MarkDeprecated(cliflags.UseEmptyDatabase.Name, "use --no-workload-database")
-		cliflagcfg.BoolFlag(demoCmd.Flags(), &demoCtx.NoExampleDatabase, cliflags.NoExampleDatabase)
 		// We also support overriding the GEOS library path for 'demo'.
 		// Even though the demoCtx uses mostly different configuration
 		// variables from startCtx, this is one case where we afford
@@ -788,6 +786,15 @@ func init() {
 		cliflagcfg.StringFlag(f, &demoCtx.ListeningURLFile, cliflags.ListeningURLFile)
 	}
 
+	{
+		// The --empty flag is only valid for the top level demo command,
+		// so we use the regular flag set.
+		f := demoCmd.Flags()
+		cliflagcfg.BoolFlag(f, &demoCtx.NoExampleDatabase, cliflags.UseEmptyDatabase)
+		_ = f.MarkDeprecated(cliflags.UseEmptyDatabase.Name, "use --no-example-database")
+		cliflagcfg.BoolFlag(f, &demoCtx.NoExampleDatabase, cliflags.NoExampleDatabase)
+	}
+
 	// statement-diag command.
 	{
 		cliflagcfg.BoolFlag(stmtDiagDeleteCmd.Flags(), &stmtDiagCtx.all, cliflags.StmtDiagDeleteAll)

pkg/cli/interactive_tests/test_explain_analyze_debug.tcl

@@ -18,6 +18,9 @@ spawn $argv sql
 set client_spawn_id $spawn_id
 eexpect root@
 
+# Note: we need to use SELECT 1 (i.e. do not select a table)
+# so that the "recreate" test below is a proper regression test
+# for issue https://github.com/cockroachdb/cockroach/issues/86472.
 send "EXPLAIN ANALYZE (DEBUG) SELECT 1;\r"
 eexpect "Statement diagnostics bundle generated."
 expect -re "SQL shell: \\\\statement-diag download (\\d+)" {
@@ -37,8 +40,15 @@ expect {
   }
 }
 
+send "\\statement-diag list\r"
+eexpect "Statement diagnostics bundles:"
+eexpect "$id"
+eexpect "EXPLAIN"
+eexpect root@
+
 send "\\statement-diag download $id\r"
 eexpect "Bundle saved to"
+eexpect root@
 
 file_exists "stmt-bundle-$id.zip"
 
@@ -47,6 +57,27 @@ eexpect eof
 
 end_test
 
+stop_server $argv
+
+start_test "Ensure that a bundle can be restarted from."
+
+set python "python2.7"
+set pyfile [file join [file dirname $argv0] unzip.py]
+system "mkdir bundle"
+system "$python $pyfile stmt-bundle-$id.zip bundle"
+
+spawn $argv debug statement-bundle recreate bundle
+eexpect "Statement was:"
+eexpect "SELECT"
+eexpect root@
+
+send_eof
+eexpect eof
+
+end_test
+
+start_server $argv
+
 start_test "Ensure that EXPLAIN ANALYZE (DEBUG) works for a tenant"
 
 start_tenant 5 $argv
@@ -76,8 +107,15 @@ expect {
   }
 }
 
+send "\\statement-diag list\r"
+eexpect "Statement diagnostics bundles:"
+eexpect "$id"
+eexpect "EXPLAIN"
+eexpect root@
+
 send "\\statement-diag download $id\r"
 eexpect "Bundle saved to"
+eexpect root@
 
 file_exists "stmt-bundle-$id.zip"
 

pkg/cli/interactive_tests/unzip.py

@@ -0,0 +1,5 @@
+import zipfile
+import sys
+
+z = zipfile.ZipFile(sys.argv[1])
+z.extractall(path=sys.argv[2])