Merge #155670 #155721

155670: sql: minor cleanup around portal pausability r=yuzefovich a=yuzefovich

This PR contains several commits that cleanup and harden code around pausable portals. See each commit for details.

Informs: https://github.com/cockroachlabs/support/issues/3463.
Epic: None
Release note: None

155721: sql: add comment about skipped FKs r=yuzefovich a=yuzefovich

We just made a change to include skipped FKs (those that we deemed "irrelevant") in the commented out form. This commit improves that logic by adding a comment for why these FKs are commented out, to reduce possible confusion.

Epic: None
Release note: None

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

Author: craig[bot]

pkg/sql/conn_executor.go

@@ -2455,7 +2455,7 @@ func (ex *connExecutor) execCmd() (retErr error) {
 			// If this is the first-time execution of a portal without a limit set,
 			// it means all rows will be exhausted, so no need to pause this portal.
 			if tcmd.Limit == 0 && portal.pauseInfo != nil && portal.pauseInfo.curRes == nil {
-				portal.pauseInfo = nil
+				ex.disablePortalPausability(&portal)
 			}
 
 			stmtRes := ex.clientComm.CreateStatementResult(
@@ -2467,7 +2467,7 @@ func (ex *connExecutor) execCmd() (retErr error) {
 				ex.sessionData().DataConversionConfig,
 				ex.sessionData().GetLocation(),
 				tcmd.Limit,
-				portalName,
+				portal.Name,
 				ex.implicitTxn(),
 				portal.portalPausablity,
 			)
@@ -2483,7 +2483,7 @@ func (ex *connExecutor) execCmd() (retErr error) {
 			// followed by Sync (which is the common case), then we still can auto-commit,
 			// which allows the 1PC txn fast path to be used.
 			canAutoCommit := ex.implicitTxn() && tcmd.FollowedBySync
-			ev, payload, err = ex.execPortal(ctx, portal, portalName, stmtRes, pinfo, canAutoCommit)
+			ev, payload, err = ex.execPortal(ctx, portal, stmtRes, pinfo, canAutoCommit)
 			return err
 		}()
 		// Note: we write to ex.statsCollector.phaseTimes, instead of ex.phaseTimes,

pkg/sql/conn_executor_exec.go

@@ -281,7 +281,6 @@ func (ex *connExecutor) recordFailure(p eventNonRetryableErrPayload) {
 func (ex *connExecutor) execPortal(
 	ctx context.Context,
 	portal PreparedPortal,
-	portalName string,
 	stmtRes CommandResult,
 	pinfo *tree.PlaceholderInfo,
 	canAutoCommit bool,
@@ -290,7 +289,7 @@ func (ex *connExecutor) execPortal(
 		if portal.isPausable() {
 			if !portal.pauseInfo.exhaustPortal.cleanup.isComplete {
 				portal.pauseInfo.exhaustPortal.cleanup.appendFunc(func(_ context.Context) {
-					ex.exhaustPortal(portalName)
+					ex.exhaustPortal(portal.Name)
 				})
 				portal.pauseInfo.exhaustPortal.cleanup.isComplete = true
 			}
@@ -330,8 +329,8 @@ func (ex *connExecutor) execPortal(
 		// to re-execute the portal from scratch.
 		// The current statement may have just closed and deleted the portal,
 		// so only exhaust it if it still exists.
-		if _, ok := ex.extraTxnState.prepStmtsNamespace.portals[portalName]; ok && !portal.isPausable() {
-			defer ex.exhaustPortal(portalName)
+		if _, ok := ex.extraTxnState.prepStmtsNamespace.portals[portal.Name]; ok && !portal.isPausable() {
+			defer ex.exhaustPortal(portal.Name)
 		}
 		return ev, payload, retErr
 
@@ -1695,6 +1694,9 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 			}
 		}
 
+		// Note that here we process the cleanup function (which will append it
+		// to the cleanup queue) without a defer since there is no more code
+		// relevant to pausable portals model below.
 		processCleanupFunc(func() {
 			cancelQueryCtx := ctx
 			if portal.isPausable() {
@@ -2526,6 +2528,13 @@ func (ex *connExecutor) rollbackSQLTransaction(
 	return eventTxnFinishAborted{}, nil
 }
 
+func getPausablePortalInfo(p *planner) *portalPauseInfo {
+	if p != nil && p.pausablePortal != nil {
+		return p.pausablePortal.pauseInfo
+	}
+	return nil
+}
+
 // Each statement in an explicit READ COMMITTED transaction has a SAVEPOINT.
 // This allows for TransactionRetry errors to be retried automatically. We don't
 // do this for implicit transactions because the conn_executor state machine
@@ -2543,13 +2552,7 @@ func (ex *connExecutor) dispatchReadCommittedStmtToExecutionEngine(
 		)
 	}
 
-	getPausablePortalInfo := func() *portalPauseInfo {
-		if p != nil && p.pausablePortal != nil {
-			return p.pausablePortal.pauseInfo
-		}
-		return nil
-	}
-	if ppInfo := getPausablePortalInfo(); ppInfo != nil {
+	if ppInfo := getPausablePortalInfo(p); ppInfo != nil {
 		p.autoRetryStmtReason = ppInfo.dispatchReadCommittedStmtToExecutionEngine.autoRetryStmtReason
 		p.autoRetryStmtCounter = ppInfo.dispatchReadCommittedStmtToExecutionEngine.autoRetryStmtCounter
 	}
@@ -2645,7 +2648,7 @@ func (ex *connExecutor) dispatchReadCommittedStmtToExecutionEngine(
 		}
 		p.autoRetryStmtCounter++
 		p.autoRetryStmtReason = maybeRetryableErr
-		if ppInfo := getPausablePortalInfo(); ppInfo != nil {
+		if ppInfo := getPausablePortalInfo(p); ppInfo != nil {
 			ppInfo.dispatchReadCommittedStmtToExecutionEngine.autoRetryStmtReason = p.autoRetryStmtReason
 			ppInfo.dispatchReadCommittedStmtToExecutionEngine.autoRetryStmtCounter = p.autoRetryStmtCounter
 		}
@@ -2672,14 +2675,8 @@ func (ex *connExecutor) dispatchReadCommittedStmtToExecutionEngine(
 func (ex *connExecutor) dispatchToExecutionEngine(
 	ctx context.Context, planner *planner, res RestrictedCommandResult,
 ) (retErr error) {
-	getPausablePortalInfo := func() *portalPauseInfo {
-		if planner != nil && planner.pausablePortal != nil {
-			return planner.pausablePortal.pauseInfo
-		}
-		return nil
-	}
 	defer func() {
-		if ppInfo := getPausablePortalInfo(); ppInfo != nil {
+		if ppInfo := getPausablePortalInfo(planner); ppInfo != nil {
 			if !ppInfo.dispatchToExecutionEngine.cleanup.isComplete {
 				ppInfo.dispatchToExecutionEngine.cleanup.isComplete = true
 			}
@@ -2743,34 +2740,50 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	}
 
 	var err error
-	if ppInfo := getPausablePortalInfo(); ppInfo != nil {
+	if ppInfo := getPausablePortalInfo(planner); ppInfo != nil {
 		if !ppInfo.dispatchToExecutionEngine.cleanup.isComplete {
 			ctx, err = ex.makeExecPlan(ctx, planner)
-			// TODO(janexing): This is a temporary solution to disallow procedure
-			// call statements that contain mutations for pausable portals. Since
-			// relational.CanMutate is not yet propagated from the function body
-			// via builder.BuildCall(), we must temporarily disallow all
-			// TCL statements, which includes the CALL statements.
-			// This should be removed once CanMutate is fully propagated.
-			// (pending https://github.com/cockroachdb/cockroach/issues/147568)
-			isTCL := planner.curPlan.stmt.AST.StatementType() == tree.TypeTCL
-			if flags := planner.curPlan.flags; err == nil && (isTCL || flags.IsSet(planFlagContainsMutation) || flags.IsSet(planFlagIsDDL)) {
-				telemetry.Inc(sqltelemetry.NotReadOnlyStmtsTriedWithPausablePortals)
-				// We don't allow mutations in a pausable portal. Set it back to
-				// an un-pausable (normal) portal.
-				planner.pausablePortal.pauseInfo = nil
-				err = res.RevokePortalPausability()
-				// If this plan is a transaction control statement, we don't
-				// even execute it but just early exit.
-				if isTCL {
-					err = errors.CombineErrors(err, ErrStmtNotSupportedForPausablePortal)
+			if err == nil {
+				// TODO(janexing): This is a temporary solution to disallow procedure
+				// call statements that contain mutations for pausable portals. Since
+				// relational.CanMutate is not yet propagated from the function body
+				// via builder.BuildCall(), we must temporarily disallow all
+				// TCL statements, which includes the CALL statements.
+				// This should be removed once CanMutate is fully propagated.
+				// (pending https://github.com/cockroachdb/cockroach/issues/147568)
+				isTCL := planner.curPlan.stmt.AST.StatementType() == tree.TypeTCL
+				// We don't allow mutations in a pausable portal.
+				notReadOnly := isTCL || planner.curPlan.flags.IsSet(planFlagContainsMutation) || planner.curPlan.flags.IsSet(planFlagIsDDL)
+				// We don't allow sub / post queries for pausable portal.
+				hasSubOrPostQuery := len(planner.curPlan.subqueryPlans) != 0 || len(planner.curPlan.cascades) != 0 ||
+					len(planner.curPlan.checkPlans) != 0 || len(planner.curPlan.triggers) != 0
+				if notReadOnly || hasSubOrPostQuery {
+					if notReadOnly {
+						telemetry.Inc(sqltelemetry.NotReadOnlyStmtsTriedWithPausablePortals)
+					} else {
+						telemetry.Inc(sqltelemetry.SubOrPostQueryStmtsTriedWithPausablePortals)
+					}
+					// This stmt is not supported via the pausable portals model
+					// - set it back to an un-pausable (normal) portal.
+					ex.disablePortalPausability(planner.pausablePortal)
+					planner.pausablePortal = nil
+					err = res.RevokePortalPausability()
+					// If this plan is a transaction control statement, we don't
+					// even execute it but just early exit.
+					if isTCL {
+						err = errors.CombineErrors(err, ErrStmtNotSupportedForPausablePortal)
+					}
+					defer planner.curPlan.close(ctx)
+				} else {
+					ppInfo.dispatchToExecutionEngine.planTop = planner.curPlan
+					defer func() {
+						ppInfo.dispatchToExecutionEngine.cleanup.appendFunc(func(ctx context.Context) {
+							ppInfo.dispatchToExecutionEngine.planTop.close(ctx)
+						})
+					}()
 				}
-				defer planner.curPlan.close(ctx)
 			} else {
-				ppInfo.dispatchToExecutionEngine.planTop = planner.curPlan
-				ppInfo.dispatchToExecutionEngine.cleanup.appendFunc(func(ctx context.Context) {
-					ppInfo.dispatchToExecutionEngine.planTop.close(ctx)
-				})
+				defer planner.curPlan.close(ctx)
 			}
 		} else {
 			planner.curPlan = ppInfo.dispatchToExecutionEngine.planTop
@@ -2825,32 +2838,12 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	ex.sessionTracing.TracePlanCheckStart(ctx)
 
 	var afterGetPlanDistribution func()
-	if planner.pausablePortal != nil {
-		if len(planner.curPlan.subqueryPlans) == 0 &&
-			len(planner.curPlan.cascades) == 0 &&
-			len(planner.curPlan.checkPlans) == 0 &&
-			len(planner.curPlan.triggers) == 0 {
-			// We don't allow a distributed plan for pausable portals.
-			origDistSQLMode := ex.sessionData().DistSQLMode
-			ex.sessionData().DistSQLMode = sessiondatapb.DistSQLOff
-			afterGetPlanDistribution = func() {
-				ex.sessionData().DistSQLMode = origDistSQLMode
-			}
-		} else {
-			telemetry.Inc(sqltelemetry.SubOrPostQueryStmtsTriedWithPausablePortals)
-			// We don't allow sub / post queries for pausable portal. Set it back to an
-			// un-pausable (normal) portal.
-			// With pauseInfo is nil, no cleanup function will be added to the stack
-			// and all clean-up steps will be performed as for normal portals.
-			// TODO(#115887): We may need to move resetting pauseInfo before we add
-			// the pausable portal cleanup step above.
-			planner.pausablePortal.pauseInfo = nil
-			// We need this so that the result consumption for this portal cannot be
-			// paused either.
-			if err := res.RevokePortalPausability(); err != nil {
-				res.SetError(err)
-				return nil
-			}
+	if getPausablePortalInfo(planner) != nil {
+		// We don't allow a distributed plan for pausable portals.
+		origDistSQLMode := ex.sessionData().DistSQLMode
+		ex.sessionData().DistSQLMode = sessiondatapb.DistSQLOff
+		afterGetPlanDistribution = func() {
+			ex.sessionData().DistSQLMode = origDistSQLMode
 		}
 	}
 	distributePlan, distSQLProhibitedErr := planner.getPlanDistribution(ctx, planner.curPlan.main)
@@ -2903,7 +2896,7 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	stats, err := ex.execWithDistSQLEngine(
 		ctx, planner, stmt.AST.StatementReturnType(), res, distribute, progAtomic, distSQLProhibitedErr,
 	)
-	if ppInfo := getPausablePortalInfo(); ppInfo != nil {
+	if ppInfo := getPausablePortalInfo(planner); ppInfo != nil {
 		// For pausable portals, we log the stats when closing the portal, so we need
 		// to aggregate the stats for all executions.
 		ppInfo.dispatchToExecutionEngine.queryStats.add(&stats)
@@ -2930,10 +2923,12 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	ex.extraTxnState.bytesRead += stats.bytesRead
 	ex.extraTxnState.rowsWritten += stats.rowsWritten
 
-	if ppInfo := getPausablePortalInfo(); ppInfo != nil && !ppInfo.dispatchToExecutionEngine.cleanup.isComplete {
+	if ppInfo := getPausablePortalInfo(planner); ppInfo != nil && !ppInfo.dispatchToExecutionEngine.cleanup.isComplete {
 		// We need to ensure that we're using the planner bound to the first-time
 		// execution of a portal.
 		curPlanner := *planner
+		// Note that here we append the cleanup function without a defer since
+		// there is no more code relevant to pausable portals model below.
 		ppInfo.dispatchToExecutionEngine.cleanup.appendFunc(func(ctx context.Context) {
 			populateQueryLevelStats(ctx, &curPlanner, ex.server.cfg, ppInfo.dispatchToExecutionEngine.queryStats, &ex.cpuStatsCollector)
 			ppInfo.dispatchToExecutionEngine.stmtFingerprintID = ex.recordStatementSummary(

pkg/sql/explain_bundle.go

@@ -894,9 +894,12 @@ func (b *stmtBundleBuilder) addEnv(ctx context.Context) {
 		for _, addFK := range addFKs {
 			fmt.Fprintf(&buf, "%s;\n", addFK)
 		}
-		// Include FK constraints that were skipped in commented out form.
-		for _, skipFK := range skipFKs {
-			fmt.Fprintf(&buf, "-- %s;\n", skipFK)
+		if len(skipFKs) > 0 {
+			// Include FK constraints that were skipped in commented out form.
+			fmt.Fprintf(&buf, "-- NOTE: these FKs are active and are only commented out for ease of bundle recreation.\n--\n")
+			for _, skipFK := range skipFKs {
+				fmt.Fprintf(&buf, "-- %s;\n", skipFK)
+			}
 		}
 	}
 	for i := range views {