Merge #97063

97063: sql: handle AOST in multi-stmt implicit txn r=ZhouXing19 a=rafiss

fixes #85066

Release note (bug fix): Fixed a bug where the AS OF SYSTEM TIME clause was handled correctly in an implicit transaction that had multiple statements.

Co-authored-by: Rafi Shamim <[email protected]>

pkg/ccl/logictestccl/testdata/logic_test/as_of

@@ -30,7 +30,7 @@ BEGIN AS OF SYSTEM TIME follower_read_timestamp(); SELECT * FROM t
 statement ok
 ROLLBACK
 
-statement error pgcode 0A000 cannot specify AS OF SYSTEM TIME with different timestamps
+statement error pgcode 0A000 inconsistent AS OF SYSTEM TIME timestamp
 SELECT * from t AS OF SYSTEM TIME '-1μs'; SELECT * from t AS OF SYSTEM TIME '-2μs'
 
 statement ok

pkg/sql/conn_executor_exec.go

@@ -720,11 +720,10 @@ func (ex *connExecutor) execStmtInOpenState(
 	p.stmt = stmt
 	p.cancelChecker.Reset(ctx)
 
-	// Auto-commit is disallowed during statement execution, if we previously executed any DDL.
-	// This is because may potentially create jobs and do other operations rather than
-	// a KV commit. Insteadand carry out any extra operations needed for DDL.he auto-connection executor will commit after this statement,
-	// in this scenario.
-	// This prevents commit during statement execution, but the connection executor,
+	// Auto-commit is disallowed during statement execution if we previously
+	// executed any DDL. This is because may potentially create jobs and do other
+	// operations rather than a KV commit.
+	// This prevents commit during statement execution, but the conn_executor
 	// will still commit this transaction after this statement executes.
 	p.autoCommit = canAutoCommit &&
 		!ex.server.cfg.TestingKnobs.DisableAutoCommitDuringExec && ex.extraTxnState.numDDL == 0
@@ -811,7 +810,10 @@ func (ex *connExecutor) handleAOST(ctx context.Context, stmt tree.Statement) err
 	if asOf == nil {
 		return nil
 	}
-	if ex.implicitTxn() {
+
+	// Implicit transactions can have multiple statements, so we need to check
+	// if one has already been executed.
+	if ex.implicitTxn() && !ex.extraTxnState.firstStmtExecuted {
 		if p.extendedEvalCtx.AsOfSystemTime == nil {
 			p.extendedEvalCtx.AsOfSystemTime = asOf
 			if !asOf.BoundedStaleness {
@@ -852,9 +854,11 @@ func (ex *connExecutor) handleAOST(ctx context.Context, stmt tree.Statement) err
 		)
 	}
 	if readTs := ex.state.getReadTimestamp(); asOf.Timestamp != readTs {
-		err = pgerror.Newf(pgcode.Syntax,
+		err = pgerror.Newf(pgcode.FeatureNotSupported,
 			"inconsistent AS OF SYSTEM TIME timestamp; expected: %s, got: %s", readTs, asOf.Timestamp)
-		err = errors.WithHint(err, "try SET TRANSACTION AS OF SYSTEM TIME")
+		if !ex.implicitTxn() {
+			err = errors.WithHint(err, "try SET TRANSACTION AS OF SYSTEM TIME")
+		}
 		return err
 	}
 	p.extendedEvalCtx.AsOfSystemTime = asOf

pkg/sql/logictest/testdata/logic_test/as_of

@@ -147,3 +147,14 @@ ROLLBACK
 # internal error.
 statement error pq: AS OF SYSTEM TIME: expected timestamp, decimal, or interval, got bool
 SELECT * FROM t AS OF SYSTEM TIME ('' BETWEEN '' AND '')
+
+# Check that AOST in multi-stmt implicit transaction has the proper error.
+statement error inconsistent AS OF SYSTEM TIME timestamp
+insert into t values(1); select * from t as of system time '-1s';
+
+statement error inconsistent AS OF SYSTEM TIME timestamp
+select * from t as of system time '-1s'; select * from t as of system time '-2s';
+
+# Specifying the AOST in the first statement (and no others) is allowed.
+statement ok
+select * from t as of system time '-1s'; select * from t;