Merge #35316

35316: sql: SET TRANSACTION improvements r=andreimatei a=andreimatei

sql: fix wild unlocking

Release note (bug fix): Fix a crash on SET TRANSACTION AS OF SYSTEM TIME
with invalid expressions.


sql: move the evaluation of SET TRANSACTION ... AS OF <expr>  

Before this patch, the way this evaluation worked was that the planner
in charge of running the whole statement was calling back into the
connExecutor, which was evaluating <expr> using connEx.planner. The two
planners were likely the same, but this was ugly. I'm trying to reduce
reliance on ex.planner; instead each statement should know what its
planner is.
This patch moves the evaluation of <expr> before the callback into the
connEx.

Release note: None

Co-authored-by: Andrei Matei <[email protected]>

pkg/sql/conn_executor.go

@@ -623,7 +623,7 @@ func (s *Server) newConnExecutorWithTxn(
 		ctx,
 		explicitTxn,
 		txn.OrigTimestamp().GoTime(),
-		/* historicalTimestamp */ nil,
+		nil, /* historicalTimestamp */
 		txn.UserPriority(),
 		tree.ReadWrite,
 		txn,
@@ -1727,7 +1727,9 @@ func (ex *connExecutor) synchronizeParallelStmts(ctx context.Context) error {
 }
 
 // setTransactionModes implements the txnModesSetter interface.
-func (ex *connExecutor) setTransactionModes(modes tree.TransactionModes) error {
+func (ex *connExecutor) setTransactionModes(
+	modes tree.TransactionModes, asOfTs hlc.Timestamp,
+) error {
 	// This method cheats and manipulates ex.state directly, not through an event.
 	// The alternative would be to create a special event, but it's unclear how
 	// that'd work given that this method is called while executing a statement.
@@ -1747,19 +1749,15 @@ func (ex *connExecutor) setTransactionModes(modes tree.TransactionModes) error {
 		return errors.Errorf("unknown isolation level: %s", modes.Isolation)
 	}
 	rwMode := modes.ReadWriteMode
-	if modes.AsOf.Expr != nil {
-
-		ts, err := ex.planner.EvalAsOfTimestamp(modes.AsOf)
-		if err != nil {
-			ex.state.mu.Unlock()
-			return err
-		}
+	if modes.AsOf.Expr != nil && (asOfTs == hlc.Timestamp{}) {
+		return pgerror.NewAssertionErrorf("expected an evaluated AS OF timestamp")
+	}
+	if (asOfTs != hlc.Timestamp{}) {
+		ex.state.setHistoricalTimestamp(ex.Ctx(), asOfTs)
+		ex.state.sqlTimestamp = asOfTs.GoTime()
 		if rwMode == tree.UnspecifiedReadWriteMode {
 			rwMode = tree.ReadOnly
 		}
-		ex.state.setHistoricalTimestamp(ex.Ctx(), ts)
-		ex.planner.semaCtx.AsOfTimestamp = &ts
-		ex.state.sqlTimestamp = ts.GoTime()
 	}
 	return ex.state.setReadOnlyMode(rwMode)
 }

pkg/sql/logictest/testdata/logic_test/txn_as_of

@@ -258,3 +258,18 @@ RELEASE SAVEPOINT cockroach_restart
 
 statement ok
 COMMIT
+
+# Ensure that errors evaluating AOST clauses in BEGIN and SET TRANSACTION do not
+# cause problems.
+
+statement error pq: AS OF SYSTEM TIME: zero timestamp is invalid
+BEGIN AS OF SYSTEM TIME '0'
+
+statement ok
+BEGIN
+
+statement error pq: AS OF SYSTEM TIME: zero timestamp is invalid
+SET TRANSACTION AS OF SYSTEM TIME '0'
+
+statement ok
+ROLLBACK

pkg/sql/planner.go

@@ -592,7 +592,10 @@ func (p *planner) runWithDistSQL(
 // txnModesSetter is an interface used by SQL execution to influence the current
 // transaction.
 type txnModesSetter interface {
-	setTransactionModes(modes tree.TransactionModes) error
+	// setTransactionModes updates some characteristics of the current
+	// transaction.
+	// asOfTs, if not empty, is the evaluation of modes.AsOf.
+	setTransactionModes(modes tree.TransactionModes, asOfTs hlc.Timestamp) error
 }
 
 // sqlStatsCollector is the interface used by SQL execution, through the

pkg/sql/set_transaction.go

@@ -14,11 +14,26 @@
 
 package sql
 
-import "github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+import (
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+)
 
 // SetTransaction sets a transaction's isolation level, priority, ro/rw state,
 // and as of timestamp.
 func (p *planner) SetTransaction(n *tree.SetTransaction) (planNode, error) {
-	return newZeroNode(nil /* columns */),
-		p.extendedEvalCtx.TxnModesSetter.setTransactionModes(n.Modes)
+	var asOfTs hlc.Timestamp
+	if n.Modes.AsOf.Expr != nil {
+		var err error
+		asOfTs, err = p.EvalAsOfTimestamp(n.Modes.AsOf)
+		if err != nil {
+			return nil, err
+		}
+		p.semaCtx.AsOfTimestamp = &asOfTs
+	}
+
+	if err := p.extendedEvalCtx.TxnModesSetter.setTransactionModes(n.Modes, asOfTs); err != nil {
+		return nil, err
+	}
+	return newZeroNode(nil /* columns */), nil
 }