Merge #151153 #151338 #151491

151153: sql: disallow pausable portal for procedure call statement r=ZhouXing19 a=ZhouXing19

Fixes #147568

Statement that is a transaction control statement is no longer allowed with pausable portals. Specifically, the following won't be allowed:

```
	CALL
	PAUSE JOB
	PAUSE SCHEDULE
	PAUSE JOBS
	PAUSE JOBS
	COMMIT PREPARED
	DEALLOCATE
	DISCARD
	EXECUTE
	ROLLBACK PREPARED
	SET TRANSACTION
	UNLISTEN
```

Ideally, we should only disallow the CALL statements for procedures that contains mutations. However, the `CanMutate` is not fully propagated from the function body yet. (Issue #147568) For now, we just disallow any TCL statements with pausable portal.

Release note (sql change): Previously, using pausable portal with procedure call caused panic, depending on the function body. Now statements that are transaction control statement, such as procedure call(e.g. `CALL myfunc()`) are not allowed with pausable portal.

151338: opt: return RelExpr from (*memo).RootExpr r=mgartner a=mgartner

`(*memo).RootExpr` now returns a `RelExpr` instead of an `Expr`, helping
eliminate some type assertions.

Release note: None


151491: colexecerror: further optimize BenchmarkSQLCatchVectorizedRuntimeError r=yuzefovich a=yuzefovich

Rather than creating and closing connections for each test case, we can just reuse the same connections in the benchmark, which speeds it up significantly (from about 4 min to 1 min).

Epic: None
Release note: None

Co-authored-by: ZhouXing19 <[email protected]>
Co-authored-by: Marcus Gartner <[email protected]>
Co-authored-by: Yahor Yuzefovich <[email protected]>

Author: 4 people

pkg/sql/colexecerror/error_test.go

@@ -184,22 +184,32 @@ func BenchmarkSQLCatchVectorizedRuntimeError(b *testing.B) {
 	// crdb_test-build behavior.
 	defer colexecerror.ProductionBehaviorForTests()()
 
-	for _, parallelism := range []int{1, 8, 32} {
+	const maxParallelism = 32
+	maxNumConns := runtime.GOMAXPROCS(0) * maxParallelism
+	// Create as many warm connections as we will need for the benchmark. These
+	// will be reused across all test cases to avoid the connection churn.
+	connsPool := make([]*gosql.DB, maxNumConns)
+	for i := range connsPool {
+		conn := s.ApplicationLayer().SQLConn(b, serverutils.DBName(""))
+		// Make sure we're using local, vectorized execution.
+		sqlDB := sqlutils.MakeSQLRunner(conn)
+		sqlDB.Exec(b, "SET distsql = off")
+		sqlDB.Exec(b, "SET vectorize = on")
+		connsPool[i] = conn
+	}
+
+	for _, parallelism := range []int{1, 8, maxParallelism} {
 		numConns := runtime.GOMAXPROCS(0) * parallelism
 		b.Run(fmt.Sprintf("conns=%d", numConns), func(b *testing.B) {
 			for _, tc := range cases {
 				stmt := fmt.Sprintf(sqlFmt, tc.builtin)
 				b.Run(tc.name, func(b *testing.B) {
-					// Create as many warm connections as we will need for the benchmark.
 					conns := make(chan *gosql.DB, numConns)
 					for i := 0; i < numConns; i++ {
-						conn := s.ApplicationLayer().SQLConn(b, serverutils.DBName(""))
-						// Make sure we're using local, vectorized execution.
-						sqlDB := sqlutils.MakeSQLRunner(conn)
-						sqlDB.Exec(b, "SET distsql = off")
-						sqlDB.Exec(b, "SET vectorize = on")
-						// Warm up the connection by executing the statement once. We should
-						// always go through the query plan cache after this.
+						conn := connsPool[i]
+						// Warm up the connection by executing the statement
+						// once. We should always go through the query plan
+						// cache after this.
 						_, _ = conn.Exec(stmt)
 						conns <- conn
 					}
@@ -209,7 +219,6 @@ func BenchmarkSQLCatchVectorizedRuntimeError(b *testing.B) {
 						var conn *gosql.DB
 						select {
 						case conn = <-conns:
-							defer conn.Close()
 						default:
 							b.Fatal("not enough warm connections")
 						}

pkg/sql/conn_executor_exec.go

@@ -322,6 +322,7 @@ func (ex *connExecutor) execPortal(
 			return nil, nil, nil
 		}
 		ev, payload, retErr = ex.execStmt(ctx, portal.Stmt.Statement, &portal, pinfo, stmtRes, canAutoCommit)
+
 		// For a non-pausable portal, it is considered exhausted regardless of the
 		// fact whether an error occurred or not - if it did, we still don't want
 		// to re-execute the portal from scratch.
@@ -2936,12 +2937,25 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	if ppInfo := getPausablePortalInfo(); ppInfo != nil {
 		if !ppInfo.dispatchToExecutionEngine.cleanup.isComplete {
 			ctx, err = ex.makeExecPlan(ctx, planner)
-			if flags := planner.curPlan.flags; err == nil && (flags.IsSet(planFlagContainsMutation) || flags.IsSet(planFlagIsDDL)) {
+			// 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)
+				}
 				defer planner.curPlan.close(ctx)
 			} else {
 				ppInfo.dispatchToExecutionEngine.planTop = planner.curPlan

pkg/sql/opt/exec/execbuilder/scalar.go

@@ -675,7 +675,7 @@ func (b *Builder) buildExistsSubquery(
 		stmtProps := []*physical.Required{{Presentation: physical.Presentation{aliasedCol}}}
 
 		// Create an wrapRootExprFn that wraps input in a Limit and a Project.
-		wrapRootExpr := func(f *norm.Factory, e memo.RelExpr) opt.Expr {
+		wrapRootExpr := func(f *norm.Factory, e memo.RelExpr) memo.RelExpr {
 			return f.ConstructProject(
 				f.ConstructLimit(
 					e,
@@ -1115,7 +1115,7 @@ func (b *Builder) initRoutineExceptionHandler(
 	blockState.ExceptionHandler = exceptionHandler
 }
 
-type wrapRootExprFn func(f *norm.Factory, e memo.RelExpr) opt.Expr
+type wrapRootExprFn func(f *norm.Factory, e memo.RelExpr) memo.RelExpr
 
 // buildRoutinePlanGenerator returns a tree.RoutinePlanFn that can plan the
 // statements in a routine that has one or more arguments.
@@ -1249,7 +1249,7 @@ func (b *Builder) buildRoutinePlanGenerator(
 			f.CopyAndReplace(originalMemo, stmt, props, replaceFn)
 
 			if wrapRootExpr != nil {
-				wrapped := wrapRootExpr(f, f.Memo().RootExpr().(memo.RelExpr)).(memo.RelExpr)
+				wrapped := wrapRootExpr(f, f.Memo().RootExpr())
 				f.Memo().SetRoot(wrapped, props)
 			}
 

pkg/sql/opt/memo/memo.go

@@ -132,7 +132,7 @@ type Memo struct {
 	// rootExpr is the root expression of the memo expression forest. It is set
 	// via a call to SetRoot. After optimization, it is set to be the root of the
 	// lowest cost tree in the forest.
-	rootExpr opt.Expr
+	rootExpr RelExpr
 
 	// rootProps are the physical properties required of the root memo expression.
 	// It is set via a call to SetRoot.
@@ -366,7 +366,7 @@ func (m *Memo) Metadata() *opt.Metadata {
 
 // RootExpr returns the root memo expression previously set via a call to
 // SetRoot.
-func (m *Memo) RootExpr() opt.Expr {
+func (m *Memo) RootExpr() RelExpr {
 	return m.rootExpr
 }
 
@@ -395,12 +395,7 @@ func (m *Memo) SetRoot(e RelExpr, phys *physical.Required) {
 // HasPlaceholders returns true if the memo contains at least one placeholder
 // operator.
 func (m *Memo) HasPlaceholders() bool {
-	rel, ok := m.rootExpr.(RelExpr)
-	if !ok {
-		panic(errors.AssertionFailedf("placeholders only supported when memo root is relational"))
-	}
-
-	return rel.Relational().HasPlaceholder
+	return m.rootExpr.Relational().HasPlaceholder
 }
 
 // IsStale returns true if the memo has been invalidated by changes to any of
@@ -555,8 +550,7 @@ func (m *Memo) ResetCost(e RelExpr, cost Cost) {
 func (m *Memo) IsOptimized() bool {
 	// The memo is optimized once the root expression has its physical properties
 	// assigned.
-	rel, ok := m.rootExpr.(RelExpr)
-	return ok && rel.RequiredPhysical() != nil
+	return m.rootExpr != nil && m.rootExpr.RequiredPhysical() != nil
 }
 
 // OptimizationCost returns a rough estimate of the cost of optimization of the

pkg/sql/opt/memo/memo_test.go

@@ -691,15 +691,15 @@ func TestStatsAvailable(t *testing.T) {
 
 	// Stats should not be available for any expression.
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT * FROM t WHERE a=1")
-	testNotAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testNotAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT sum(a), b FROM t GROUP BY b")
-	testNotAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testNotAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx,
 		"SELECT * FROM t AS t1, t AS t2 WHERE t1.a = t2.a AND t1.b = 5",
 	)
-	testNotAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testNotAvailable(o.Memo().RootExpr())
 
 	if _, err := catalog.ExecuteDDL(
 		`ALTER TABLE t INJECT STATISTICS '[
@@ -729,15 +729,15 @@ func TestStatsAvailable(t *testing.T) {
 
 	// Stats should be available for all expressions.
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT * FROM t WHERE a=1")
-	testAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx, "SELECT sum(a), b FROM t GROUP BY b")
-	testAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testAvailable(o.Memo().RootExpr())
 
 	opttestutils.BuildQuery(t, &o, catalog, &evalCtx,
 		"SELECT * FROM t AS t1, t AS t2 WHERE t1.a = t2.a AND t1.b = 5",
 	)
-	testAvailable(o.Memo().RootExpr().(memo.RelExpr))
+	testAvailable(o.Memo().RootExpr())
 }
 
 // traverseExpr is a helper function to recursively traverse a relational

pkg/sql/opt/norm/factory.go

@@ -420,7 +420,7 @@ func (f *Factory) AssignPlaceholders(from *memo.Memo) (err error) {
 		}
 		return f.CopyAndReplaceDefault(e, replaceFn)
 	}
-	f.CopyAndReplace(from, from.RootExpr().(memo.RelExpr), from.RootProps(), replaceFn)
+	f.CopyAndReplace(from, from.RootExpr(), from.RootProps(), replaceFn)
 
 	return nil
 }

pkg/sql/opt/norm/factory_test.go

@@ -101,7 +101,7 @@ func TestCopyAndReplace(t *testing.T) {
 		}
 		return o.Factory().CopyAndReplaceDefault(e, replaceFn)
 	}
-	o.Factory().CopyAndReplace(m, m.RootExpr().(memo.RelExpr), m.RootProps(), replaceFn)
+	o.Factory().CopyAndReplace(m, m.RootExpr(), m.RootProps(), replaceFn)
 
 	if e, err := o.Optimize(); err != nil {
 		t.Fatal(err)
@@ -146,7 +146,7 @@ func TestCopyAndReplaceWithScan(t *testing.T) {
 			replaceFn = func(e opt.Expr) opt.Expr {
 				return o.Factory().CopyAndReplaceDefault(e, replaceFn)
 			}
-			o.Factory().CopyAndReplace(m, m.RootExpr().(memo.RelExpr), m.RootProps(), replaceFn)
+			o.Factory().CopyAndReplace(m, m.RootExpr(), m.RootProps(), replaceFn)
 
 			if _, err := o.Optimize(); err != nil {
 				t.Fatal(err)

pkg/sql/opt/xform/optimizer.go

@@ -261,7 +261,7 @@ func (o *Optimizer) Optimize() (_ opt.Expr, err error) {
 	o.optimizeRootWithProps()
 
 	// Now optimize the entire expression tree.
-	root := o.mem.RootExpr().(memo.RelExpr)
+	root := o.mem.RootExpr()
 	rootProps := o.mem.RootProps()
 	o.optimizeGroup(root, rootProps)
 
@@ -922,10 +922,7 @@ func (o *Optimizer) ensureOptState(grp memo.RelExpr, required *physical.Required
 // properties required of it. This may trigger the creation of a new root and
 // new properties.
 func (o *Optimizer) optimizeRootWithProps() {
-	root, ok := o.mem.RootExpr().(memo.RelExpr)
-	if !ok {
-		panic(errors.AssertionFailedf("Optimize can only be called on relational root expressions"))
-	}
+	root := o.mem.RootExpr()
 	rootProps := o.mem.RootProps()
 
 	// [SimplifyRootOrdering]

pkg/sql/opt/xform/optimizer_test.go

@@ -113,7 +113,7 @@ func TestDetachMemoRace(t *testing.T) {
 			// Rewrite the filter to use a different column, which will trigger creation
 			// of new table statistics. If the statistics object is aliased, this will
 			// be racy.
-			f.CopyAndReplace(mem, mem.RootExpr().(memo.RelExpr), mem.RootProps(), replaceFn)
+			f.CopyAndReplace(mem, mem.RootExpr(), mem.RootProps(), replaceFn)
 			wg.Done()
 		}()
 	}

pkg/sql/opt/xform/placeholder_fast_path.go

@@ -45,7 +45,7 @@ func (o *Optimizer) TryPlaceholderFastPath() (ok bool, err error) {
 		}
 	}()
 
-	root := o.mem.RootExpr().(memo.RelExpr)
+	root := o.mem.RootExpr()
 
 	rootRelProps := root.Relational()
 	// We are dealing with a memo that still contains placeholders. The statistics