Merge #155824 #155867

155824: sql: fix top-level query stats when "inner" plans are present r=yuzefovich a=yuzefovich

Previously, whenever we ran an "inner" plan (via `runPlanInsidePlan` helper), we ignored the top-level query stats for the "inner" plan. As a result, reads and writes performed by the inner plan weren't reflected in the "outer" top-level query stats. This is now fixed by adjusting the routines, apply joins, and recursive CTEs to propagate their metrics as ProducerMetadata objects.

Note that for routines the access to the DistSQL infra is rather difficult, so we plumbed the access via the planner straight into the DistSQLReceiver, and even though it's ugly, it should work in practice. The only alternative I see is adding the necessary reference into the `eval.Context`, but then it gets tricky to actually set that and ensure the right copy of the eval context is observed by all routines (plus we'd need to make the copy or restore the original state somehow), so I chose to not pursue it.

Epic: None

Release note (bug fix): Previously, CockroachDB didn't include reads and writes performed by routines (user-defined functions and stored procedures) as well as apply joins into `bytes read`, `rows read`, and `rows written` statement execution statistics, and this is now fixed. The bug has been present since before 23.2 version.

155867: sql: fix REPLACE routine logic with multiple DEFAULT expressions r=yuzefovich a=yuzefovich

When we added support for DEFAULT expressions in routine signatures in 41cb974, we introduced a bug for REPLACE functionality. Namely, we need to store the type-checked serialized representation of the default expressions in the proto, so the expressions in the existing Overload must also have been type-checked (so that we can find the correct signature in the proto to update). However, we had the incorrect short-circuiting behavior when type-checking existing default expressions where we'd stop once a "signature change" is discovered. The fix is simple - we need to type-check all default expressions separately from checking whether the signature has changed.

Fixes: #155859.

Release note (bug fix): Previously, CockroachDB could encounter an internal error when replacing (via CREATE OR REPLACE stmt) a user-defined function or a stored procedure when multiple DEFAULT expressions are included in the old signature. The bug has been present since v24.2 (when support for DEFAULT expressions was added) and is now fixed.

Co-authored-by: Yahor Yuzefovich <yahor@cockroachlabs.com>

Author: craig[bot]

pkg/ccl/logictestccl/testdata/logic_test/procedure_params

@@ -880,4 +880,23 @@ $$ LANGUAGE PLpgSQL;
 statement ok
 DROP PROCEDURE p;
 
+# Regression test for not type-checking all default expressions when replacing
+# the procedure (#155859).
+statement ok
+CREATE PROCEDURE p (s STRING(1) DEFAULT NULL, d DECIMAL(1, 1) DEFAULT NULL) AS $$
+BEGIN
+  NULL;
+END;
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+CREATE OR REPLACE PROCEDURE p (s STRING(2) DEFAULT NULL, d DECIMAL(2, 2) DEFAULT NULL) AS $$
+BEGIN
+  NULL;
+END;
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+DROP PROCEDURE p;
+
 subtest end

pkg/ccl/logictestccl/testdata/logic_test/udf_params

@@ -1048,4 +1048,23 @@ $$ LANGUAGE PLpgSQL;
 statement ok
 DROP PROCEDURE p;
 
+# Regression test for not type-checking all default expressions when replacing
+# the UDF (#155859).
+statement ok
+CREATE FUNCTION f (s STRING(1) DEFAULT NULL, d DECIMAL(1, 1) DEFAULT NULL) RETURNS VOID AS $$
+BEGIN
+  NULL;
+END;
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+CREATE OR REPLACE FUNCTION f (s STRING(2) DEFAULT NULL, d DECIMAL(2, 2) DEFAULT NULL) RETURNS VOID AS $$
+BEGIN
+  NULL;
+END;
+$$ LANGUAGE PLpgSQL;
+
+statement ok
+DROP FUNCTION f;
+
 subtest end

pkg/sql/apply_join.go

@@ -35,6 +35,10 @@ type applyJoinNode struct {
 	// The data source with no outer columns.
 	singleInputPlanNode
 
+	// forwarder allows propagating the ProducerMetadata towards the
+	// DistSQLReceiver.
+	forwarder metadataForwarder
+
 	// pred represents the join predicate.
 	pred *joinPredicate
 
@@ -248,12 +252,13 @@ func (a *applyJoinNode) runNextRightSideIteration(params runParams, leftRow tree
 	}
 	plan := p.(*planComponents)
 	rowResultWriter := NewRowResultWriter(&a.run.rightRows)
-	if err := runPlanInsidePlan(
-		ctx, params, plan, rowResultWriter,
-		nil /* deferredRoutineSender */, "", /* stmtForDistSQLDiagram */
-	); err != nil {
+	queryStats, err := runPlanInsidePlan(
+		ctx, params, plan, rowResultWriter, nil /* deferredRoutineSender */, "", /* stmtForDistSQLDiagram */
+	)
+	if err != nil {
 		return err
 	}
+	forwardInnerQueryStats(a.forwarder, queryStats)
 	a.run.rightRowsIterator = newRowContainerIterator(ctx, a.run.rightRows)
 	return nil
 }
@@ -267,7 +272,7 @@ func runPlanInsidePlan(
 	resultWriter rowResultWriter,
 	deferredRoutineSender eval.DeferredRoutineSender,
 	stmtForDistSQLDiagram string,
-) error {
+) (topLevelQueryStats, error) {
 	defer plan.close(ctx)
 	execCfg := params.ExecCfg()
 	recv := MakeDistSQLReceiver(
@@ -286,6 +291,11 @@ func runPlanInsidePlan(
 	// before we can produce any "outer" rows to be returned to the client, so
 	// we make sure to unset pausablePortal field on the planner.
 	plannerCopy.pausablePortal = nil
+	// Avoid any possible metadata confusion by unsetting the
+	// routineMetadataForwarder (if there is a routine in the inner plan that
+	// needs it, then the plannerCopy will be updated during the inner plan
+	// setup).
+	plannerCopy.routineMetadataForwarder = nil
 
 	// planner object embeds the extended eval context, so we will modify that
 	// (which won't affect the outer planner's extended eval context), and we'll
@@ -301,6 +311,8 @@ func runPlanInsidePlan(
 		// return from this method (after the main query is executed).
 		subqueryResultMemAcc := params.p.Mon().MakeBoundAccount()
 		defer subqueryResultMemAcc.Close(ctx)
+		// Note that planAndRunSubquery updates recv.stats with top-level
+		// subquery stats.
 		if !execCfg.DistSQLPlanner.PlanAndRunSubqueries(
 			ctx,
 			&plannerCopy,
@@ -311,7 +323,7 @@ func runPlanInsidePlan(
 			false, /* skipDistSQLDiagramGeneration */
 			params.p.mustUseLeafTxn(),
 		) {
-			return resultWriter.Err()
+			return recv.stats, resultWriter.Err()
 		}
 	}
 
@@ -338,10 +350,10 @@ func runPlanInsidePlan(
 
 	// Check if there was an error interacting with the resultWriter.
 	if recv.commErr != nil {
-		return recv.commErr
+		return recv.stats, recv.commErr
 	}
 	if resultWriter.Err() != nil {
-		return resultWriter.Err()
+		return recv.stats, resultWriter.Err()
 	}
 
 	plannerCopy.autoCommit = false
@@ -358,10 +370,10 @@ func runPlanInsidePlan(
 	// need to update the plan for cleanup purposes before proceeding.
 	*plan = plannerCopy.curPlan.planComponents
 	if recv.commErr != nil {
-		return recv.commErr
+		return recv.stats, recv.commErr
 	}
 
-	return resultWriter.Err()
+	return recv.stats, resultWriter.Err()
 }
 
 func (a *applyJoinNode) Values() tree.Datums {

pkg/sql/conn_executor.go

@@ -3964,6 +3964,7 @@ func (ex *connExecutor) initPlanner(ctx context.Context, p *planner) {
 	p.noticeSender = nil
 	p.preparedStatements = ex.getPrepStmtsAccessor()
 	p.sqlCursors = ex.getCursorAccessor()
+	p.routineMetadataForwarder = nil
 	p.storedProcTxnState = ex.getStoredProcTxnStateAccessor()
 	p.createdSequences = ex.getCreatedSequencesAccessor()
 

pkg/sql/conn_executor_exec.go

@@ -3288,6 +3288,8 @@ type topLevelQueryStats struct {
 	// client receiving the PGWire protocol messages (as well as construcing
 	// those messages).
 	clientTime time.Duration
+	// NB: when adding another field here, consider whether
+	// forwardInnerQueryStats method needs an adjustment.
 }
 
 func (s *topLevelQueryStats) add(other *topLevelQueryStats) {

pkg/sql/create_function.go

@@ -309,20 +309,26 @@ func (n *createFunctionNode) replaceFunction(
 		return err
 	}
 
+	// We allow two types of "signature changes":
+	// - reordering OUT parameters in respect to input ones, and
+	// - changing the DEFAULT expression.
 	signatureChanged := len(existing.OutParamOrdinals) != len(outParamOrdinals) ||
 		len(existing.DefaultExprs) != len(defaultExprs)
 	for i := 0; !signatureChanged && i < len(outParamOrdinals); i++ {
 		signatureChanged = existing.OutParamOrdinals[i] != outParamOrdinals[i] ||
 			!existing.OutParamTypes.GetAt(i).Equivalent(outParamTypes[i])
 	}
-	for i := 0; !signatureChanged && i < len(defaultExprs); i++ {
-		typ := existing.Types.GetAt(i + existing.Types.Length() - len(defaultExprs))
-		// Update the overload to store the type-checked expression since this
-		// is what is stored in the FunctionSignature proto.
+	// Additionally, we must type-check all DEFAULT expressions since we need to
+	// store the type-checked serialized form in the FunctionSignature proto.
+	// (This is also assumed in ReplaceOverload.)
+	for i := 0; i < len(existing.DefaultExprs); i++ {
+		typ := existing.Types.GetAt(i + existing.Types.Length() - len(existing.DefaultExprs))
 		existing.DefaultExprs[i], err = tree.TypeCheck(params.ctx, existing.DefaultExprs[i], params.p.SemaCtx(), typ)
 		if err != nil {
 			return err
 		}
+	}
+	for i := 0; !signatureChanged && i < len(defaultExprs); i++ {
 		signatureChanged = tree.Serialize(existing.DefaultExprs[i]) != defaultExprs[i]
 	}
 	if signatureChanged {

pkg/sql/distsql_physical_planner.go

@@ -4343,14 +4343,17 @@ func (dsp *DistSQLPlanner) wrapPlan(
 
 	// Copy the evalCtx.
 	evalCtx := *planCtx.ExtendedEvalCtx
-	wrapper := newPlanNodeToRowSource(
+	wrapper, err := newPlanNodeToRowSource(
 		n,
 		runParams{
 			extendedEvalCtx: &evalCtx,
 			p:               planCtx.planner,
 		},
 		firstNotWrapped,
 	)
+	if err != nil {
+		return nil, err
+	}
 	if !wrapper.rowsAffected && planCtx.planDepth == 1 && planCtx.stmtType == tree.RowsAffected {
 		// Return an error if the receiver expects to get the number of rows
 		// affected, but the planNode returns something else.

pkg/sql/distsql_running.go

@@ -1014,6 +1014,15 @@ func (dsp *DistSQLPlanner) Run(
 		return
 	}
 
+	if len(flows) == 1 && evalCtx.Txn != nil && evalCtx.Txn.Type() == kv.RootTxn {
+		// If we have a fully local plan and a RootTxn, we don't expect any
+		// concurrency, so it's safe to use the DistSQLReceiver to push the
+		// metadata into directly from routines.
+		if planCtx.planner != nil {
+			planCtx.planner.routineMetadataForwarder = recv
+		}
+	}
+
 	if finishedSetupFn != nil {
 		finishedSetupFn(flow)
 	}
@@ -1144,6 +1153,8 @@ type DistSQLReceiver struct {
 	}
 }
 
+var _ metadataForwarder = &DistSQLReceiver{}
+
 // rowResultWriter is a subset of CommandResult to be used with the
 // DistSQLReceiver. It's implemented by RowResultWriter.
 type rowResultWriter interface {
@@ -1570,6 +1581,35 @@ func (r *DistSQLReceiver) checkConcurrentError() {
 	}
 }
 
+type metadataForwarder interface {
+	forwardMetadata(metadata *execinfrapb.ProducerMetadata)
+}
+
+// forwardInnerQueryStats propagates the query stats of "inner" plans as
+// metadata via the forwarder.
+func forwardInnerQueryStats(f metadataForwarder, stats topLevelQueryStats) {
+	if !buildutil.CrdbTestBuild && f == nil {
+		// Safety measure in production builds in case the forwarder is nil for
+		// some reason.
+		return
+	}
+	meta := execinfrapb.GetProducerMeta()
+	meta.Metrics = execinfrapb.GetMetricsMeta()
+	meta.Metrics.BytesRead = stats.bytesRead
+	meta.Metrics.RowsRead = stats.rowsRead
+	meta.Metrics.RowsWritten = stats.rowsWritten
+	// stats.networkEgressEstimate and stats.clientTime are ignored since they
+	// only matter at the "true" top-level query (and actually should be zero
+	// here anyway).
+	f.forwardMetadata(meta)
+}
+
+func (r *DistSQLReceiver) forwardMetadata(metadata *execinfrapb.ProducerMetadata) {
+	// Note that we don't use pushMeta method directly in order to go through
+	// the testing callback path.
+	r.Push(nil /* row */, metadata)
+}
+
 // pushMeta takes in non-empty metadata object and pushes it to the result
 // writer. Possibly updated status is returned.
 func (r *DistSQLReceiver) pushMeta(meta *execinfrapb.ProducerMetadata) execinfra.ConsumerStatus {

pkg/sql/distsql_running_test.go

@@ -1214,3 +1214,127 @@ SELECT id, details FROM jobs AS j INNER JOIN cte1 ON id = job_id WHERE id = 1;
 	require.Equal(t, 1, id)
 	require.Equal(t, 1, details)
 }
+
+// TestTopLevelQueryStats verifies that top-level query stats are collected
+// correctly, including when the query executes "plans inside plans".
+func TestTopLevelQueryStats(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// testQuery will be updated throughout the test to the current target.
+	var testQuery atomic.Value
+	// The callback will send number of rows read and rows written (for each
+	// ProducerMetadata.Metrics object) on these channels, respectively.
+	rowsReadCh, rowsWrittenCh := make(chan int64), make(chan int64)
+	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			SQLExecutor: &ExecutorTestingKnobs{
+				DistSQLReceiverPushCallbackFactory: func(_ context.Context, query string) func(rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) (rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) {
+					if target := testQuery.Load(); target == nil || target.(string) != query {
+						return nil
+					}
+					return func(row rowenc.EncDatumRow, batch coldata.Batch, meta *execinfrapb.ProducerMetadata) (rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) {
+						if meta != nil && meta.Metrics != nil {
+							rowsReadCh <- meta.Metrics.RowsRead
+							rowsWrittenCh <- meta.Metrics.RowsWritten
+						}
+						return row, batch, meta
+					}
+				},
+			},
+		},
+	})
+	defer s.Stopper().Stop(context.Background())
+
+	if _, err := sqlDB.Exec(`
+CREATE TABLE t (k INT PRIMARY KEY);
+INSERT INTO t SELECT generate_series(1, 10);
+CREATE FUNCTION no_reads() RETURNS INT AS 'SELECT 1' LANGUAGE SQL;
+CREATE FUNCTION reads() RETURNS INT AS 'SELECT count(*) FROM t' LANGUAGE SQL;
+CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x' LANGUAGE SQL;
+`); err != nil {
+		t.Fatal(err)
+	}
+
+	for _, tc := range []struct {
+		name           string
+		query          string
+		expRowsRead    int64
+		expRowsWritten int64
+	}{
+		{
+			name:           "simple read",
+			query:          "SELECT k FROM t",
+			expRowsRead:    10,
+			expRowsWritten: 0,
+		},
+		{
+			name:           "simple write",
+			query:          "INSERT INTO t SELECT generate_series(11, 42)",
+			expRowsRead:    0,
+			expRowsWritten: 32,
+		},
+		{
+			name: "read with apply join",
+			query: `SELECT (
+    WITH foo AS MATERIALIZED (SELECT k FROM t AS x WHERE x.k = y.k)
+    SELECT * FROM foo
+  ) FROM t AS y`,
+			expRowsRead:    84, // scanning the table twice
+			expRowsWritten: 0,
+		},
+		{
+			name:           "routine, no reads",
+			query:          "SELECT no_reads()",
+			expRowsRead:    0,
+			expRowsWritten: 0,
+		},
+		{
+			name:           "routine, reads",
+			query:          "SELECT reads()",
+			expRowsRead:    42,
+			expRowsWritten: 0,
+		},
+		{
+			name:           "routine, write",
+			query:          "SELECT write(43)",
+			expRowsRead:    0,
+			expRowsWritten: 1,
+		},
+		{
+			name:           "routine, multiple reads and writes",
+			query:          "SELECT reads(), write(44), reads(), write(45), write(46), reads()",
+			expRowsRead:    133, // first read is 43 rows, second is 44, third is 46
+			expRowsWritten: 3,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			testQuery.Store(tc.query)
+			errCh := make(chan error)
+			// Spin up the worker goroutine which will actually execute the
+			// query.
+			go func() {
+				defer close(errCh)
+				_, err := sqlDB.Exec(tc.query)
+				errCh <- err
+			}()
+			// In the main goroutine, loop until the query is completed while
+			// accumulating the top-level query stats.
+			var rowsRead, rowsWritten int64
+		LOOP:
+			for {
+				select {
+				case read := <-rowsReadCh:
+					rowsRead += read
+				case written := <-rowsWrittenCh:
+					rowsWritten += written
+				case err := <-errCh:
+					require.NoError(t, err)
+					break LOOP
+				}
+			}
+			require.Equal(t, tc.expRowsRead, rowsRead)
+			require.Equal(t, tc.expRowsWritten, rowsWritten)
+		})
+	}
+}