feat: support DML with Returning clause in Connection API

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java

@@ -140,6 +140,7 @@ public static class ParsedStatement {
     private final ClientSideStatementImpl clientSideStatement;
     private final Statement statement;
     private final String sqlWithoutComments;
+    private final boolean returningClause;
 
     private static ParsedStatement clientSideStatement(
         ClientSideStatementImpl clientSideStatement,
@@ -155,11 +156,13 @@ private static ParsedStatement ddl(Statement statement, String sqlWithoutComment
     private static ParsedStatement query(
         Statement statement, String sqlWithoutComments, QueryOptions defaultQueryOptions) {
       return new ParsedStatement(
-          StatementType.QUERY, statement, sqlWithoutComments, defaultQueryOptions);
+          StatementType.QUERY, statement, sqlWithoutComments, defaultQueryOptions, false);
     }
 
-    private static ParsedStatement update(Statement statement, String sqlWithoutComments) {
-      return new ParsedStatement(StatementType.UPDATE, statement, sqlWithoutComments);
+    private static ParsedStatement update(
+        Statement statement, String sqlWithoutComments, boolean returningClause) {
+      return new ParsedStatement(
+          StatementType.UPDATE, statement, sqlWithoutComments, returningClause);
     }
 
     private static ParsedStatement unknown(Statement statement, String sqlWithoutComments) {
@@ -176,23 +179,34 @@ private ParsedStatement(
       this.clientSideStatement = clientSideStatement;
       this.statement = statement;
       this.sqlWithoutComments = sqlWithoutComments;
+      this.returningClause = false;
+    }
+
+    private ParsedStatement(
+        StatementType type,
+        Statement statement,
+        String sqlWithoutComments,
+        boolean returningClause) {
+      this(type, statement, sqlWithoutComments, null, returningClause);
     }
 
     private ParsedStatement(StatementType type, Statement statement, String sqlWithoutComments) {
-      this(type, statement, sqlWithoutComments, null);
+      this(type, statement, sqlWithoutComments, null, false);
     }
 
     private ParsedStatement(
         StatementType type,
         Statement statement,
         String sqlWithoutComments,
-        QueryOptions defaultQueryOptions) {
+        QueryOptions defaultQueryOptions,
+        boolean returningClause) {
       Preconditions.checkNotNull(type);
       Preconditions.checkNotNull(statement);
       this.type = type;
       this.clientSideStatement = null;
       this.statement = mergeQueryOptions(statement, defaultQueryOptions);
       this.sqlWithoutComments = sqlWithoutComments;
+      this.returningClause = returningClause;
     }
 
     @Override
@@ -219,6 +233,12 @@ public StatementType getType() {
       return type;
     }
 
+    /** Returns whether the statement has a returning clause or not. * */
+    @InternalApi
+    public boolean hasReturningClause() {
+      return this.returningClause;
+    }
+
     /**
      * Returns true if the statement is a query that will return a {@link
      * com.google.cloud.spanner.ResultSet}.
@@ -354,7 +374,7 @@ ParsedStatement parse(Statement statement, QueryOptions defaultQueryOptions) {
     } else if (isQuery(sql)) {
       return ParsedStatement.query(statement, sql, defaultQueryOptions);
     } else if (isUpdateStatement(sql)) {
-      return ParsedStatement.update(statement, sql);
+      return ParsedStatement.update(statement, sql, checkReturningClause(sql));
     } else if (isDdlStatement(sql)) {
       return ParsedStatement.ddl(statement, sql);
     }
@@ -521,4 +541,20 @@ static int countOccurrencesOf(char c, String string) {
     }
     return res;
   }
+
+  /**
+   * Checks if the given SQL string contains a Returning clause. This method is used only in case of
+   * a DML statement.
+   *
+   * @param sql The sql string without comments that has to be evaluated.
+   * @return A boolean indicating whether the sql string has a Returning clause or not.
+   * @throws SpannerException If the input sql string contains an unclosed string/byte literal.
+   */
+  @InternalApi
+  abstract boolean checkReturningClauseInternal(String sql);
+
+  @InternalApi
+  public boolean checkReturningClause(String sql) {
+    return checkReturningClauseInternal(sql);
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java

@@ -853,6 +853,9 @@ public StatementResult execute(Statement statement) {
       case QUERY:
         return StatementResultImpl.of(internalExecuteQuery(parsedStatement, AnalyzeMode.NONE));
       case UPDATE:
+        if (parsedStatement.hasReturningClause()) {
+          return StatementResultImpl.of(internalExecuteQuery(parsedStatement, AnalyzeMode.NONE));
+        }
         return StatementResultImpl.of(get(internalExecuteUpdateAsync(parsedStatement)));
       case DDL:
         get(executeDdlAsync(parsedStatement));
@@ -881,6 +884,10 @@ public AsyncStatementResult executeAsync(Statement statement) {
         return AsyncStatementResultImpl.of(
             internalExecuteQueryAsync(parsedStatement, AnalyzeMode.NONE));
       case UPDATE:
+        if (parsedStatement.hasReturningClause()) {
+          return AsyncStatementResultImpl.of(
+              internalExecuteQueryAsync(parsedStatement, AnalyzeMode.NONE));
+        }
         return AsyncStatementResultImpl.of(internalExecuteUpdateAsync(parsedStatement));
       case DDL:
         return AsyncStatementResultImpl.noResult(executeDdlAsync(parsedStatement));
@@ -918,7 +925,7 @@ private ResultSet parseAndExecuteQuery(
     Preconditions.checkNotNull(analyzeMode);
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
     ParsedStatement parsedStatement = getStatementParser().parse(query, this.queryOptions);
-    if (parsedStatement.isQuery()) {
+    if (parsedStatement.isQuery() || parsedStatement.isUpdate()) {
       switch (parsedStatement.getType()) {
         case CLIENT_SIDE:
           return parsedStatement
@@ -928,6 +935,19 @@ private ResultSet parseAndExecuteQuery(
         case QUERY:
           return internalExecuteQuery(parsedStatement, analyzeMode, options);
         case UPDATE:
+          if (parsedStatement.hasReturningClause()) {
+            // Cannot execute in read-only mode or in READ_ONLY_TRANSACTION transaction mode.
+            if (this.isReadOnly()
+                || (this.isInTransaction()
+                    && this.getTransactionMode() == TransactionMode.READ_ONLY_TRANSACTION)) {
+              throw SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.FAILED_PRECONDITION,
+                  "DML statement with returning clause cannot be executed in read-only mode: "
+                      + parsedStatement.getSqlWithoutComments());
+            }
+            return internalExecuteQuery(parsedStatement, analyzeMode, options);
+          }
+
         case DDL:
         case UNKNOWN:
         default:
@@ -943,7 +963,7 @@ private AsyncResultSet parseAndExecuteQueryAsync(
     Preconditions.checkNotNull(query);
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
     ParsedStatement parsedStatement = getStatementParser().parse(query, this.queryOptions);
-    if (parsedStatement.isQuery()) {
+    if (parsedStatement.isQuery() || parsedStatement.isUpdate()) {
       switch (parsedStatement.getType()) {
         case CLIENT_SIDE:
           return ResultSets.toAsyncResultSet(
@@ -956,6 +976,19 @@ private AsyncResultSet parseAndExecuteQueryAsync(
         case QUERY:
           return internalExecuteQueryAsync(parsedStatement, analyzeMode, options);
         case UPDATE:
+          if (parsedStatement.hasReturningClause()) {
+            // Cannot execute in read-only mode or in READ_ONLY_TRANSACTION transaction mode.
+            if (this.isReadOnly()
+                || (this.isInTransaction()
+                    && this.getTransactionMode() == TransactionMode.READ_ONLY_TRANSACTION)) {
+              throw SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.FAILED_PRECONDITION,
+                  "DML statement with returning clause cannot be executed in read-only mode: "
+                      + parsedStatement.getSqlWithoutComments());
+            }
+            return internalExecuteQueryAsync(parsedStatement, analyzeMode, options);
+          }
+
         case DDL:
         case UNKNOWN:
         default:
@@ -974,6 +1007,13 @@ public long executeUpdate(Statement update) {
     if (parsedStatement.isUpdate()) {
       switch (parsedStatement.getType()) {
         case UPDATE:
+          if (parsedStatement.hasReturningClause()) {
+            throw SpannerExceptionFactory.newSpannerException(
+                ErrorCode.FAILED_PRECONDITION,
+                "DML statement with returning clause cannot be executed over executeUpdate: "
+                    + parsedStatement.getSqlWithoutComments()
+                    + ". Please use executeQuery instead.");
+          }
           return get(internalExecuteUpdateAsync(parsedStatement));
         case CLIENT_SIDE:
         case QUERY:
@@ -995,6 +1035,13 @@ public ApiFuture<Long> executeUpdateAsync(Statement update) {
     if (parsedStatement.isUpdate()) {
       switch (parsedStatement.getType()) {
         case UPDATE:
+          if (parsedStatement.hasReturningClause()) {
+            throw SpannerExceptionFactory.newSpannerException(
+                ErrorCode.FAILED_PRECONDITION,
+                "DML statement with returning clause cannot be executed over executeUpdateAsync: "
+                    + parsedStatement.getSqlWithoutComments()
+                    + ". Please use executeQueryAsync instead.");
+          }
           return internalExecuteUpdateAsync(parsedStatement);
         case CLIENT_SIDE:
         case QUERY:
@@ -1141,8 +1188,9 @@ private ResultSet internalExecuteQuery(
       final QueryOption... options) {
     Preconditions.checkArgument(
         statement.getType() == StatementType.QUERY
-            || (statement.getType() == StatementType.UPDATE && analyzeMode != AnalyzeMode.NONE),
-        "Statement must either be a query or a DML mode with analyzeMode!=NONE");
+            || (statement.getType() == StatementType.UPDATE
+                && (analyzeMode != AnalyzeMode.NONE || statement.hasReturningClause())),
+        "Statement must either be a query or a DML mode with analyzeMode!=NONE or returning clause");
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
     return get(
         transaction.executeQueryAsync(
@@ -1154,7 +1202,9 @@ private AsyncResultSet internalExecuteQueryAsync(
       final AnalyzeMode analyzeMode,
       final QueryOption... options) {
     Preconditions.checkArgument(
-        statement.getType() == StatementType.QUERY, "Statement must be a query");
+        (statement.getType() == StatementType.QUERY)
+            || (statement.getType() == StatementType.UPDATE && statement.hasReturningClause()),
+        "Statement must be a query or DML with returning clause.");
     UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
     return ResultSets.toAsyncResultSet(
         transaction.executeQueryAsync(

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/PostgreSQLStatementParser.java

@@ -25,6 +25,7 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
 @InternalApi
@@ -65,6 +66,8 @@ String removeCommentsAndTrimInternal(String sql) {
     Preconditions.checkNotNull(sql);
     boolean isInSingleLineComment = false;
     int multiLineCommentLevel = 0;
+    boolean whitespaceBeforeOrAfterMultiLineComment = false;
+    int multiLineCommentStartIdx = -1;
     StringBuilder res = new StringBuilder(sql.length());
     int index = 0;
     while (index < sql.length()) {
@@ -78,6 +81,19 @@ String removeCommentsAndTrimInternal(String sql) {
       } else if (multiLineCommentLevel > 0) {
         if (sql.length() > index + 1 && c == ASTERISK && sql.charAt(index + 1) == SLASH) {
           multiLineCommentLevel--;
+          if (multiLineCommentLevel == 0) {
+            if (!whitespaceBeforeOrAfterMultiLineComment && (sql.length() > index + 2)) {
+              whitespaceBeforeOrAfterMultiLineComment =
+                  Character.isWhitespace(sql.charAt(index + 2));
+            }
+            // If the multiline comment does not have any whitespace before or after it, and it is
+            // neither at the start nor at the end of SQL string, append an extra space.
+            if (!whitespaceBeforeOrAfterMultiLineComment
+                && (multiLineCommentStartIdx != 0)
+                && (index != sql.length() - 2)) {
+              res.append(' ');
+            }
+          }
           index++;
         } else if (sql.length() > index + 1 && c == SLASH && sql.charAt(index + 1) == ASTERISK) {
           multiLineCommentLevel++;
@@ -92,6 +108,10 @@ String removeCommentsAndTrimInternal(String sql) {
           continue;
         } else if (sql.length() > index + 1 && c == SLASH && sql.charAt(index + 1) == ASTERISK) {
           multiLineCommentLevel++;
+          if (index >= 1) {
+            whitespaceBeforeOrAfterMultiLineComment = Character.isWhitespace(sql.charAt(index - 1));
+          }
+          multiLineCommentStartIdx = index;
           index += 2;
           continue;
         } else {
@@ -267,4 +287,36 @@ private void appendIfNotNull(
       result.append(prefix).append(tag).append(suffix);
     }
   }
+
+  private boolean isReturning(String sql, int index) {
+    // RETURNING is a reserved keyword in PG, but requires a
+    // leading AS to be used as column label, to avoid ambiguity.
+    // We thus check for cases which do not have a leading AS.
+    // (https://www.postgresql.org/docs/current/sql-keywords-appendix.html)
+    return Pattern.compile("[\\s)]RETURNING[\\s(][\\s\\S]*", Pattern.CASE_INSENSITIVE)
+            .matcher(sql.substring(index))
+            .matches()
+        && !((index >= 3)
+            && Pattern.compile("[\\s]AS[\\s]RETURNING[\\s(][\\s\\S]*", Pattern.CASE_INSENSITIVE)
+                .matcher(sql.substring(index - 3))
+                .matches());
+  }
+
+  @InternalApi
+  @Override
+  boolean checkReturningClauseInternal(String rawSql) {
+    Preconditions.checkNotNull(rawSql);
+    int index = 0;
+    String sql = rawSql.replaceAll("\\s+", " ");
+    boolean hasReturningClause = false;
+    while (index < sql.length()) {
+      if (!hasReturningClause && isReturning(sql, index)) {
+        hasReturningClause = true;
+        index++;
+      } else {
+        index = skip(sql, index, null);
+      }
+    }
+    return hasReturningClause;
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java

@@ -45,6 +45,7 @@
 import com.google.cloud.spanner.TransactionContext;
 import com.google.cloud.spanner.TransactionManager;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
+import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.TransactionRetryListener.RetryResult;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -337,7 +338,10 @@ public ApiFuture<ResultSet> executeQueryAsync(
       final ParsedStatement statement,
       final AnalyzeMode analyzeMode,
       final QueryOption... options) {
-    Preconditions.checkArgument(statement.isQuery(), "Statement is not a query");
+    Preconditions.checkArgument(
+        (statement.getType() == StatementType.QUERY)
+            || (statement.getType() == StatementType.UPDATE && statement.hasReturningClause()),
+        "Statement must be a query or DML with returning clause");
     checkValidTransaction();
 
     ApiFuture<ResultSet> res;

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java

@@ -179,12 +179,18 @@ public ApiFuture<ResultSet> executeQueryAsync(
       final QueryOption... options) {
     Preconditions.checkNotNull(statement);
     Preconditions.checkArgument(
-        statement.isQuery() || (statement.isUpdate() && analyzeMode != AnalyzeMode.NONE),
+        statement.isQuery()
+            || (statement.isUpdate()
+                && (analyzeMode != AnalyzeMode.NONE || statement.hasReturningClause())),
         "The statement must be a query, or the statement must be DML and AnalyzeMode must be PLAN or PROFILE");
     checkAndMarkUsed();
 
-    if (statement.isUpdate() && analyzeMode != AnalyzeMode.NONE) {
-      return analyzeTransactionalUpdateAsync(statement, analyzeMode);
+    if (statement.isUpdate()) {
+      if (analyzeMode != AnalyzeMode.NONE) {
+        return analyzeTransactionalUpdateAsync(statement, analyzeMode);
+      }
+      // DML with returning clause.
+      return executeDmlReturningAsync(statement, options);
     }
 
     final ReadOnlyTransaction currentTransaction =
@@ -217,6 +223,30 @@ public ApiFuture<ResultSet> executeQueryAsync(
     return executeStatementAsync(statement, callable, SpannerGrpc.getExecuteStreamingSqlMethod());
   }
 
+  private ApiFuture<ResultSet> executeDmlReturningAsync(
+      final ParsedStatement update, QueryOption... options) {
+    Callable<ResultSet> callable =
+        () -> {
+          try {
+            writeTransaction = createWriteTransaction();
+            ResultSet resultSet =
+                writeTransaction.run(
+                    transaction ->
+                        DirectExecuteResultSet.ofResultSet(
+                            transaction.executeQuery(update.getStatement(), options)));
+            state = UnitOfWorkState.COMMITTED;
+            return resultSet;
+          } catch (Throwable t) {
+            state = UnitOfWorkState.COMMIT_FAILED;
+            throw t;
+          }
+        };
+    return executeStatementAsync(
+        update,
+        callable,
+        ImmutableList.of(SpannerGrpc.getExecuteSqlMethod(), SpannerGrpc.getCommitMethod()));
+  }
+
   @Override
   public Timestamp getReadTimestamp() {
     ConnectionPreconditions.checkState(