fix: avoid implicit transaction when query is pipelined with client-side statements (#4976)

* fix: avoid implicit transaction when query is pipelined with client-side statements

Previously, when a query was sent in the same pipeline (before a Sync message)
along with client-side or session statements (such as SHOW spanner.read_timestamp
or SET timezone = 'UTC'), PGAdapter treated the batch as multiple statements
and began an implicit read-only transaction.

This caused two issues:
1. An unnecessary BeginTransaction RPC was executed.
2. Queries failed when spanner.read_only_staleness was set to bounded staleness
   (e.g., MAX_STALENESS or MIN_READ_TIMESTAMP), because Spanner only allows
   bounded staleness on single-use reads outside of a multi-statement transaction.

This change:
- Updates BackendConnection.maybeBeginImplicitTransaction() to check if all
  subsequent statements in the pipeline are client-side or session-management
  statements. If so, it keeps auto-commit mode and executes the query as a
  single-use read.
- Unifies the detection of client-side statements with getSessionManagementStatement()
  to recognize SET, SHOW, and RESET ALL commands.
- Generalizes the Describe + Execute statement check to be index-relative,
  preventing pipelines that start with a client-side statement from triggering an
  unintended transaction.

* fix: transaction handling for mixed client-side + DML

Author: olavloite

src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java

@@ -242,6 +242,10 @@ boolean isBatchingPossible() {
 
     abstract boolean isUpdate();
 
+    boolean isClientSide() {
+      return false;
+    }
+
     void execute() {
       Span span = createSpan(toString(), statement);
       try (Scope ignore = span.makeCurrent()) {
@@ -272,6 +276,10 @@ private final class Execute extends BufferedStatement<StatementResult> {
     private final String command;
     private final Function<Statement, Statement> statementBinder;
     private final boolean analyze;
+    private final LocalStatement localStatement;
+    private final SessionStatement sessionStatement;
+    private final Exception parseException;
+    private final boolean clientSide;
 
     Execute(
         String command,
@@ -291,6 +299,29 @@ private final class Execute extends BufferedStatement<StatementResult> {
       this.command = command;
       this.statementBinder = statementBinder;
       this.analyze = analyze;
+      LocalStatement localStatement = getLocalStatement(this.statement);
+      if (localStatement != null && !localStatement.hasReplacementStatement()) {
+        this.localStatement = localStatement;
+        this.sessionStatement = null;
+        this.parseException = null;
+        this.clientSide = true;
+      } else {
+        this.localStatement = null;
+        SessionStatement sessionStatement = null;
+        Exception parseException = null;
+        try {
+          sessionStatement = getSessionManagementStatement(this.statement, this.parsedStatement);
+        } catch (Exception exception) {
+          parseException = exception;
+        }
+        this.sessionStatement = sessionStatement;
+        this.parseException = parseException;
+        this.clientSide =
+            !isTransactionStatement(this.parsedStatement)
+                && (this.parsedStatement.getType() == StatementType.CLIENT_SIDE
+                    || this.sessionStatement != null
+                    || this.parseException != null);
+      }
     }
 
     @Override
@@ -308,23 +339,28 @@ boolean isUpdate() {
       return this.parsedStatement.isUpdate();
     }
 
+    @Override
+    boolean isClientSide() {
+      return this.clientSide;
+    }
+
     @Override
     void doExecute() {
       Statement updatedStatement = statement;
       try {
         checkConnectionState();
+        if (this.parseException != null) {
+          throw SpannerExceptionFactory.asSpannerException(this.parseException);
+        }
         // TODO(b/235719478): If the statement is a BEGIN statement and there is a COMMIT statement
         //  at a later point in the batch, and all the statements in the transaction block are
         //  SELECT statements, then we should create a read-only transaction. Also, if a transaction
         //  block always ends with a ROLLBACK, PGAdapter should skip the entire execution of that
         //  block.
-        SessionStatement sessionStatement =
-            getSessionManagementStatement(updatedStatement, parsedStatement);
-        LocalStatement localStatement = getLocalStatement(statement);
-        if (localStatement != null && !localStatement.hasReplacementStatement()) {
-          result.set(localStatement.execute(BackendConnection.this, statement));
-        } else if (sessionStatement != null) {
-          result.set(sessionStatement.execute(sessionState, spannerConnection));
+        if (this.localStatement != null) {
+          result.set(this.localStatement.execute(BackendConnection.this, statement));
+        } else if (this.sessionStatement != null) {
+          result.set(this.sessionStatement.execute(sessionState, spannerConnection));
         } else if (connectionState == ConnectionState.ABORTED
             && !spannerConnection.isInTransaction()
             && (isRollback(parsedStatement) || isCommit(parsedStatement))) {
@@ -539,20 +575,6 @@ boolean isUnsupportedConcurrencyModeException(SpannerException spannerException)
               .startsWith(
                   "INVALID_ARGUMENT: io.grpc.StatusRuntimeException: INVALID_ARGUMENT: Unsupported concurrency mode in query using INFORMATION_SCHEMA.");
     }
-
-    @Nullable
-    SessionStatement getSessionManagementStatement(
-        Statement statement, ParsedStatement parsedStatement) {
-      if (parsedStatement.getType() == StatementType.UNKNOWN
-          || (parsedStatement.getType() == StatementType.QUERY
-              && isCommand(SHOW, statement.getSql()))
-          || (parsedStatement.getType() == StatementType.CLIENT_SIDE
-              && parsedStatement.getClientSideStatementType()
-                  == ClientSideStatementType.RESET_ALL)) {
-        return SessionStatementParser.parse(parsedStatement, statement.getSql());
-      }
-      return null;
-    }
   }
 
   public static final int MAX_PARTITIONS =
@@ -1299,9 +1321,10 @@ private void maybeBeginImplicitTransaction(int index, boolean isSync) {
       return;
     }
 
-    // Only start an implicit transaction if we have more than one statement left. Otherwise, just
-    // let the Spanner connection execute the statement in auto-commit mode.
-    if (isSync && index == bufferedStatements.size() - 1) {
+    // Only start an implicit transaction if we have more than one statement left that can be
+    // executed in a transaction. Otherwise, just let the Spanner connection execute the statement
+    // in auto-commit mode.
+    if (isSync && hasOnlyClientSideStatementsAfter(index)) {
       return;
     }
     // Don't start an implicit transaction if this is already a transaction statement.
@@ -1310,7 +1333,7 @@ private void maybeBeginImplicitTransaction(int index, boolean isSync) {
     }
     // No need to start a transaction for DDL or client side statements.
     if (bufferedStatements.get(index).parsedStatement.getType() == StatementType.DDL
-        || bufferedStatements.get(index).parsedStatement.getType() == StatementType.CLIENT_SIDE) {
+        || bufferedStatements.get(index).isClientSide()) {
       return;
     }
     // If there are only DML statements left, those can be executed as an auto-commit dml batch.
@@ -1322,18 +1345,20 @@ private void maybeBeginImplicitTransaction(int index, boolean isSync) {
       return;
     }
     // Do not start an implicit transaction if all that is in the buffer is a DESCRIBE and an
-    // EXECUTE message for the same statement.
+    // EXECUTE message for the same statement (optionally surrounded only by client-side
+    // statements).
     if (isSync
-        && bufferedStatements.size() == 2
-        && bufferedStatements.get(0) instanceof Execute
-        && bufferedStatements.get(1) instanceof Execute
-        && ((Execute) bufferedStatements.get(0)).analyze
-        && !((Execute) bufferedStatements.get(1)).analyze
+        && bufferedStatements.size() > index + 1
+        && bufferedStatements.get(index) instanceof Execute
+        && bufferedStatements.get(index + 1) instanceof Execute
+        && ((Execute) bufferedStatements.get(index)).analyze
+        && !((Execute) bufferedStatements.get(index + 1)).analyze
         && bufferedStatements
-            .get(0)
+            .get(index)
             .statement
             .getSql()
-            .equals(bufferedStatements.get(1).statement.getSql())) {
+            .equals(bufferedStatements.get(index + 1).statement.getSql())
+        && hasOnlyClientSideStatementsAfter(index + 1)) {
       return;
     }
 
@@ -1505,9 +1530,36 @@ private boolean isTransactionStatement(int index) {
     return isBegin(index) || isCommit(index) || isRollback(index);
   }
 
-  private boolean hasOnlyDmlStatementsAfter(int index) {
-    return bufferedStatements.subList(index, bufferedStatements.size()).stream()
-        .allMatch(statement -> statement.parsedStatement.getType() == StatementType.UPDATE);
+  @VisibleForTesting
+  boolean hasOnlyDmlStatementsAfter(int index) {
+    boolean seenExecute = false;
+    boolean seenClientSide = false;
+    for (int i = index; i < bufferedStatements.size(); i++) {
+      BufferedStatement<?> statement = bufferedStatements.get(i);
+      if (statement.isClientSide()) {
+        seenClientSide = true;
+      } else if (statement.isUpdate() && !statement.parsedStatement.hasReturningClause()) {
+        if (seenClientSide) {
+          // A DML statement appeared after a client-side statement.
+          // These cannot be batched together in a single auto-commit DML batch.
+          return false;
+        }
+        boolean isAnalyze = statement instanceof Execute && ((Execute) statement).analyze;
+        if (isAnalyze) {
+          if (seenExecute || i > index) {
+            // An analyze/describe statement appeared after an execute statement,
+            // or there are multiple analyze statements. These cannot be executed
+            // in a single auto-commit DML batch.
+            return false;
+          }
+        } else {
+          seenExecute = true;
+        }
+      } else {
+        return false;
+      }
+    }
+    return true;
   }
 
   @VisibleForTesting
@@ -1516,7 +1568,43 @@ boolean hasUpdateStatementsAfter(int index) {
         .anyMatch(BufferedStatement::isUpdate);
   }
 
-  private int getStatementCount() {
+  @Nullable
+  SessionStatement getSessionManagementStatement(
+      Statement statement, ParsedStatement parsedStatement) {
+    if (parsedStatement.getType() == StatementType.UNKNOWN
+        || (parsedStatement.getType() == StatementType.QUERY && isCommand(SHOW, statement.getSql()))
+        || (parsedStatement.getType() == StatementType.CLIENT_SIDE
+            && parsedStatement.getClientSideStatementType() == ClientSideStatementType.RESET_ALL)) {
+      return SessionStatementParser.parse(parsedStatement, statement.getSql());
+    }
+    return null;
+  }
+
+  @VisibleForTesting
+  boolean hasOnlyClientSideStatementsAfter(int index) {
+    if (index >= bufferedStatements.size() - 1) {
+      return true;
+    }
+    for (int i = index + 1; i < bufferedStatements.size(); i++) {
+      if (!bufferedStatements.get(i).isClientSide()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  @VisibleForTesting
+  boolean isClientSide(BufferedStatement<?> statement) {
+    return statement.isClientSide();
+  }
+
+  @VisibleForTesting
+  BufferedStatement<?> getBufferedStatement(int index) {
+    return bufferedStatements.get(index);
+  }
+
+  @VisibleForTesting
+  int getStatementCount() {
     return bufferedStatements.size();
   }