From aa871290e48192ac8e7f0f6c374ba4b921a055fc Mon Sep 17 00:00:00 2001
From: Rafi Shamim <rafi@cockroachlabs.com>
Date: Mon, 13 Feb 2023 14:27:12 -0500
Subject: [PATCH] sql: handle AOST in multi-stmt implicit txn

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.
---
 .../logictestccl/testdata/logic_test/as_of    |  2 +-
 pkg/sql/conn_executor_exec.go                 | 20 +++++++++++--------
 pkg/sql/logictest/testdata/logic_test/as_of   | 11 ++++++++++
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/pkg/ccl/logictestccl/testdata/logic_test/as_of b/pkg/ccl/logictestccl/testdata/logic_test/as_of
index 722406252bef..fced0b3d8039 100644
--- a/pkg/ccl/logictestccl/testdata/logic_test/as_of
+++ b/pkg/ccl/logictestccl/testdata/logic_test/as_of
@@ -29,7 +29,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
diff --git a/pkg/sql/conn_executor_exec.go b/pkg/sql/conn_executor_exec.go
index ed0a117bb417..a87a52ed1f1f 100644
--- a/pkg/sql/conn_executor_exec.go
+++ b/pkg/sql/conn_executor_exec.go
@@ -672,11 +672,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
@@ -760,7 +759,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 {
@@ -801,9 +803,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
diff --git a/pkg/sql/logictest/testdata/logic_test/as_of b/pkg/sql/logictest/testdata/logic_test/as_of
index b9facf210f2e..b4f98a5ce168 100644
--- a/pkg/sql/logictest/testdata/logic_test/as_of
+++ b/pkg/sql/logictest/testdata/logic_test/as_of
@@ -139,3 +139,14 @@ SET TRANSACTION AS OF system time '-1s'
 
 statement ok
 ROLLBACK
+
+# 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;