Merge #50924 #51065 #51121 #51163 #51194 #51218 #51226

50924: cli: add commands for managing statement diagnostics r=RaduBerinde a=RaduBerinde

This change adds a `statement-diag` command, with the following subcommands:
```
  list        list available bundles and outstanding activation requests
  download    download statement diagnostics bundle into a zip file
  delete      delete a statement diagnostics bundle
  cancel      cancel an outstanding activation request
```

Fixes #48597.

Release note (cli change): A new set of `statement-diag` CLI commands that can
be used to manage statement diagnostics.

51065: importccl: fix target column ordering bug for PGDUMP import r=mjibson a=Anzoteh96

The current implementation assumes that the target columns of a PGDUMP query is the same as how they are created in the case where target columns are declared in PGDUMP file. This PR addresses it by detecting the target columns in the PGDUMP statement itself if this is the case. In addition, given that the target columns may not be well-determined at the formation of a new `DatumRowConverter`, so the check of unsupported default column expression is also moved to the `DatumRowConverter.Row()` function. 

Fixed #51159 

Release note: None

51121: sql: move parallelize scans control in the execbuilder r=RaduBerinde a=RaduBerinde

Parallel scans refers to disabling scan batch limits, which allows the
distsender to issue requests to multiple ranges in parallel. This is
only safe to use when there is a known upper bound for the number of
results.

Currently we plumb maxResults to the scanNode and TableReader, and
each execution component runs similar logic to decide whether to
parallelize.

This change cleans this up by centralizing this decision inside the
execbuilder. In the future, we may want the coster to be aware of this
parallelization as well.

For simplicity, we drop the cluster setting that controls this (it was
added for fear of problems but it has been on by default for a long
time).

Release note: None

51163: jobs: do not include cancel jobs as running for scheduled jobs r=pbardea a=pbardea

Previously, the query that the job scheduling system would use to detect
if there were any already running jobs would include canceled jobs due
to using a the alternative spelling: 'cancelled'.

This commit changes the query to use the enums instead of manually
listing their state. It also adds a test to ensure that jobs are still
run, regardless of the wait policy when all previous runs are in a
terminal state.

Release note: None

51194: pgwire: de-strictify extended protocol handling r=jordanlewis a=jordanlewis

Fixes #33693.
Fixes #41511.

Previously, the protocol handler didn't permit simple queries during the
extended protocol mode. I think this was done because the spec vaguely
says that extended protocol commands must be ended with a SYNC command:
https://www.postgresql.org/docs/9.3/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY

However, an examination of the PostgreSQL source code shows that the
extended protocol mode bit is only used for error recovery. If an error
is encountered during extended protocol mode in Postgres, all commands
are skipped until the next Sync.

CockroachDB never implemented that behavior - instead, it used the
extended protocol mode bit only to reject simple queries if they came
before the Sync.

This commit removes the extended protocol mode bit, as the use case we
used it for was incorrect. It's unclear to me whether we need to re-add
the bit for dealing with error cases, but we haven't needed it yet.
Adding that might be follow-on work, and won't come in this PR.

Release note (bug fix): prevent spurious "SimpleQuery not allowed while
in extended protocol mode" errors.

51218: Fix typo in txn.commits metric help text r=nvanbenschoten a=a-robinson

Sorry for the mostly useless change, but it bothered me :)

51226: opt: fix error in case of casted NULL arguments to AddGeometryColumn r=rytaft a=rytaft

This commit fixes an error that occurred when `AddGeometryColumn` was
called with `NULL` arguments that were cast to the type specified by the
function signature. #50992 already fixed the case when `AddGeometryColumn`
was called with bare `NULL` arguments, since those were detected by `TypeCheck`.
`TypeCheck` does not detect `NULL` arguments if they are cast to the correct
type.

This commit fixes the error by adding an explicit check in the `optbuilder`
that each argument is not null before calling the `SQLFn` of the
`AddGeometryColumn` overload.

Informs #50296

Release note (bug fix): Fixed an internal error that occurred when
AddGeometryColumn was called with NULL arguments. This now results in
a no-op and returns NULL.

Co-authored-by: Radu Berinde <[email protected]>
Co-authored-by: anzoteh96 <[email protected]>
Co-authored-by: Paul Bardea <[email protected]>
Co-authored-by: Jordan Lewis <[email protected]>
Co-authored-by: Alex Robinson <[email protected]>
Co-authored-by: Rebecca Taft <[email protected]>

pkg/ccl/importccl/import_stmt_test.go

@@ -2605,6 +2605,13 @@ func TestImportIntoCSV(t *testing.T) {
 				fmt.Sprintf(`non-constant default expression .* for non-targeted column "b" is not supported by IMPORT INTO`),
 				fmt.Sprintf(`IMPORT INTO t (a) CSV DATA ("%s")`, srv.URL))
 		})
+		t.Run("pgdump", func(t *testing.T) {
+			data = "INSERT INTO t VALUES (1, 2), (3, 4)"
+			sqlDB.Exec(t, `CREATE TABLE t (a INT, b INT DEFAULT 42, c INT)`)
+			sqlDB.Exec(t, "IMPORT INTO t (c, a) PGDUMP DATA ($1)", srv.URL)
+			defer sqlDB.Exec(t, `DROP TABLE t`)
+			sqlDB.CheckQueryResults(t, `SELECT * from t`, [][]string{{"2", "42", "1"}, {"4", "42", "3"}})
+		})
 	})
 
 	t.Run("import-not-targeted-not-null", func(t *testing.T) {
@@ -4029,6 +4036,21 @@ func TestImportPgDump(t *testing.T) {
 			}
 		})
 	}
+	t.Run("target-cols-reordered", func(t *testing.T) {
+		data := `
+				CREATE TABLE "t" ("a" INT, "b" INT DEFAULT 42, "c" INT);
+				INSERT INTO "t" ("c", "a") VALUES ('1', '2'), ('3', '4');
+			`
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.Method == "GET" {
+				_, _ = w.Write([]byte(data))
+			}
+		}))
+		defer srv.Close()
+		defer sqlDB.Exec(t, "DROP TABLE t")
+		sqlDB.Exec(t, "IMPORT PGDUMP ($1)", srv.URL)
+		sqlDB.CheckQueryResults(t, `SELECT * from t`, [][]string{{"2", "42", "1"}, {"4", "42", "3"}})
+	})
 }
 
 // TestImportPgDumpGeo tests that a file with SQLFn classes can be

pkg/ccl/importccl/read_import_pgdump.go

@@ -483,6 +483,7 @@ type pgDumpReader struct {
 	descs  map[string]*execinfrapb.ReadImportDataSpec_ImportTable
 	kvCh   chan row.KVBatch
 	opts   roachpb.PgDumpOptions
+	colMap map[*row.DatumRowConverter](map[string]int)
 }
 
 var _ inputConverter = &pgDumpReader{}
@@ -496,20 +497,31 @@ func newPgDumpReader(
 	evalCtx *tree.EvalContext,
 ) (*pgDumpReader, error) {
 	converters := make(map[string]*row.DatumRowConverter, len(descs))
+	colMap := make(map[*row.DatumRowConverter](map[string]int))
 	for name, table := range descs {
 		if table.Desc.IsTable() {
-			conv, err := row.NewDatumRowConverter(ctx, table.Desc, nil /* targetColNames */, evalCtx, kvCh)
+			colSubMap := make(map[string]int, len(table.TargetCols))
+			targetCols := make(tree.NameList, len(table.TargetCols))
+			for i, colName := range table.TargetCols {
+				targetCols[i] = tree.Name(colName)
+			}
+			for i, col := range table.Desc.VisibleColumns() {
+				colSubMap[col.Name] = i
+			}
+			conv, err := row.NewDatumRowConverter(ctx, table.Desc, targetCols, evalCtx, kvCh)
 			if err != nil {
 				return nil, err
 			}
 			converters[name] = conv
+			colMap[conv] = colSubMap
 		}
 	}
 	return &pgDumpReader{
 		kvCh:   kvCh,
 		tables: converters,
 		descs:  descs,
 		opts:   opts,
+		colMap: colMap,
 	}, nil
 }
 
@@ -567,22 +579,46 @@ func (m *pgDumpReader) readFile(
 			if ok && conv == nil {
 				return errors.Errorf("missing schema info for requested table %q", name)
 			}
+			expectedColLen := len(i.Columns)
+			if expectedColLen == 0 {
+				// Case where the targeted columns are not specified in the PGDUMP file, but in
+				// the command "IMPORT INTO table (targetCols) PGDUMP DATA (filename)"
+				expectedColLen = len(conv.VisibleCols)
+			}
 			values, ok := i.Rows.Select.(*tree.ValuesClause)
 			if !ok {
 				return errors.Errorf("unsupported: %s", i.Rows.Select)
 			}
 			inserts++
 			startingCount := count
+			var targetColMapInd []int
+			if len(i.Columns) != 0 {
+				targetColMapInd = make([]int, len(i.Columns))
+				conv.IsTargetCol = make(map[int]struct{}, len(i.Columns))
+				for j := range i.Columns {
+					colName := i.Columns[j].String()
+					ind, ok := m.colMap[conv][colName]
+					if !ok {
+						return errors.Newf("targeted column %q not found", colName)
+					}
+					conv.IsTargetCol[ind] = struct{}{}
+					targetColMapInd[j] = ind
+				}
+			}
 			for _, tuple := range values.Rows {
 				count++
 				if count <= resumePos {
 					continue
 				}
-				if expected, got := len(conv.VisibleCols), len(tuple); expected != got {
-					return errors.Errorf("expected %d values, got %d: %v", expected, got, tuple)
+				if got := len(tuple); expectedColLen != got {
+					return errors.Errorf("expected %d values, got %d: %v", expectedColLen, got, tuple)
 				}
-				for i, expr := range tuple {
-					typed, err := expr.TypeCheck(ctx, &semaCtx, conv.VisibleColTypes[i])
+				for j, expr := range tuple {
+					ind := j
+					if len(i.Columns) != 0 {
+						ind = targetColMapInd[j]
+					}
+					typed, err := expr.TypeCheck(ctx, &semaCtx, conv.VisibleColTypes[ind])
 					if err != nil {
 						return errors.Wrapf(err, "reading row %d (%d in insert statement %d)",
 							count, count-startingCount, inserts)
@@ -592,7 +628,7 @@ func (m *pgDumpReader) readFile(
 						return errors.Wrapf(err, "reading row %d (%d in insert statement %d)",
 							count, count-startingCount, inserts)
 					}
-					conv.Datums[i] = converted
+					conv.Datums[ind] = converted
 				}
 				if err := conv.Row(ctx, inputIdx, count); err != nil {
 					return err

pkg/cli/cli.go

@@ -195,6 +195,7 @@ func init() {
 		quitCmd,
 
 		sqlShellCmd,
+		stmtDiagCmd,
 		authCmd,
 		nodeCmd,
 		dumpCmd,

pkg/cli/cli_test.go

@@ -297,7 +297,7 @@ func isSQLCommand(args []string) bool {
 		return false
 	}
 	switch args[0] {
-	case "sql", "dump", "workload", "nodelocal":
+	case "sql", "dump", "workload", "nodelocal", "statement-diag":
 		return true
 	case "node":
 		if len(args) == 0 {
@@ -1405,6 +1405,7 @@ Available Commands:
   init              initialize a cluster
   cert              create ca, node, and client certs
   sql               open a sql shell
+  statement-diag    commands for managing statement diagnostics bundles
   auth-session      log in and out of HTTP sessions
   node              list, inspect, drain or remove nodes
   dump              dump sql tables

pkg/cli/cliflags/flags.go

@@ -1201,4 +1201,14 @@ other items retrieved by the zip command may still consider
 confidential data or PII.
 `,
 	}
+
+	StmtDiagDeleteAll = FlagInfo{
+		Name:        "all",
+		Description: `Delete all bundles.`,
+	}
+
+	StmtDiagCancelAll = FlagInfo{
+		Name:        "all",
+		Description: `Cancel all outstanding requests.`,
+	}
 )

pkg/cli/context.go

@@ -49,6 +49,7 @@ func initCLIDefaults() {
 	setNetworkBenchContextDefaults()
 	setSqlfmtContextDefaults()
 	setDemoContextDefaults()
+	setStmtDiagContextDefaults()
 	setAuthContextDefaults()
 
 	initPreFlagsDefaults()
@@ -506,6 +507,16 @@ func setDemoContextDefaults() {
 	demoCtx.insecure = false
 }
 
+// stmtDiagCtx captures the command-line parameters of the 'statement-diag'
+// command.
+var stmtDiagCtx struct {
+	all bool
+}
+
+func setStmtDiagContextDefaults() {
+	stmtDiagCtx.all = false
+}
+
 // GetServerCfgStores provides direct public access to the StoreSpecList inside
 // serverCfg. This is used by CCL code to populate some fields.
 //

pkg/cli/flags.go

@@ -539,6 +539,7 @@ func init() {
 	clientCmds = append(clientCmds, nodeCmds...)
 	clientCmds = append(clientCmds, systemBenchCmds...)
 	clientCmds = append(clientCmds, nodeLocalCmds...)
+	clientCmds = append(clientCmds, stmtDiagCmds...)
 	for _, cmd := range clientCmds {
 		f := cmd.PersistentFlags()
 		varFlag(f, addrSetter{&cliCtx.clientConnHost, &cliCtx.clientConnPort}, cliflags.ClientHost)
@@ -649,6 +650,7 @@ func init() {
 	sqlCmds := []*cobra.Command{sqlShellCmd, dumpCmd, demoCmd}
 	sqlCmds = append(sqlCmds, authCmds...)
 	sqlCmds = append(sqlCmds, demoCmd.Commands()...)
+	sqlCmds = append(sqlCmds, stmtDiagCmds...)
 	sqlCmds = append(sqlCmds, nodeLocalCmds...)
 	for _, cmd := range sqlCmds {
 		f := cmd.Flags()
@@ -736,6 +738,12 @@ func init() {
 		stringFlag(f, &startCtx.geoLibsDir, cliflags.GeoLibsDir)
 	}
 
+	// statement-diag command.
+	{
+		boolFlag(stmtDiagDeleteCmd.Flags(), &stmtDiagCtx.all, cliflags.StmtDiagDeleteAll)
+		boolFlag(stmtDiagCancelCmd.Flags(), &stmtDiagCtx.all, cliflags.StmtDiagCancelAll)
+	}
+
 	// sqlfmt command.
 	{
 		f := sqlfmtCmd.Flags()

pkg/cli/main_test.go

@@ -13,6 +13,7 @@ package cli_test
 import (
 	"os"
 	"testing"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/build"
 	"github.com/cockroachdb/cockroach/pkg/security"
@@ -30,6 +31,10 @@ func TestMain(m *testing.M) {
 	// a version injected. Pretend to be a very up-to-date version.
 	defer build.TestingOverrideTag("v999.0.0")()
 
+	// The times for Example_statement_diag are reported in the local timezone.
+	// Fix it to UTC so the output is always the same.
+	time.Local = time.UTC
+
 	serverutils.InitTestServerFactory(server.TestServerFactory)
 	os.Exit(m.Run())
 }