feat: support DML with Returning clause in Connection API (#1978)

* feat: incorporate dml with returning clause

* feat: changes

* feat: change handling of AsyncResultSet.

* fix: lint

* doc: add comments

* fix: lint

* test: add tests for executeBatchUpdate

* test: import fix

* test: remove redundant import

* test: add abort tests for dml returning

* test: add pg dml returning tests

* feat: change error statement

* doc: add doc for dml with returning clause usage

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

Co-authored-by: Knut Olav Løite <koloite@gmail.com>

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

Co-authored-by: Knut Olav Løite <koloite@gmail.com>

* fix: incorporate review comments

* test: add more test cases

* test: add todo

* test: add separate abort tests for dml returning

* fix: add try-with-resources block around ResultSet

* feat: enhancement by adding a pre-check

* feat: changes

* test: delete unnecessary test

* test: add few more tests to PG parser

* feat: method doc update

* test: nit fixes

* feat: handle another corner case

* test: add another test

* clirr fixes

* revert env for integration tests

* remove comments

* skip returning tests in emulator

* fix: linting

Co-authored-by: Knut Olav Løite <koloite@gmail.com>

google-cloud-spanner/clirr-ignored-differences.xml

@@ -202,4 +202,9 @@
     <className>com/google/cloud/spanner/DatabaseClient</className>
     <method>java.lang.String getDatabaseRole()</method>
   </difference>
+  <difference>
+    <differenceType>7013</differenceType>
+    <className>com/google/cloud/spanner/connection/AbstractStatementParser</className>
+    <method>boolean checkReturningClauseInternal(java.lang.String)</method>
+  </difference>
 </differences>

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}.
@@ -355,7 +375,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);
     }
@@ -460,6 +480,10 @@ private boolean statementStartsWith(String sql, Iterable<String> checkStatements
   static final char SLASH = '/';
   static final char ASTERISK = '*';
   static final char DOLLAR = '$';
+  static final char SPACE = ' ';
+  static final char CLOSE_PARENTHESIS = ')';
+  static final char COMMA = ',';
+  static final char UNDERSCORE = '_';
 
   /**
    * Removes comments from and trims the given sql statement using the dialect of this parser.
@@ -522,4 +546,19 @@ 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.
+   */
+  @InternalApi
+  protected 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/Connection.java

@@ -856,8 +856,8 @@ default RpcPriority getRPCPriority() {
    * state. The returned value depends on the type of statement:
    *
    * <ul>
-   *   <li>Queries will return a {@link ResultSet}
-   *   <li>DML statements will return an update count
+   *   <li>Queries and DML statements with returning clause will return a {@link ResultSet}.
+   *   <li>Simple DML statements will return an update count
    *   <li>DDL statements will return a {@link ResultType#NO_RESULT}
    *   <li>Connection and transaction statements (SET AUTOCOMMIT=TRUE|FALSE, SHOW AUTOCOMMIT, SET
    *       TRANSACTION READ ONLY, etc) will return either a {@link ResultSet} or {@link
@@ -874,9 +874,9 @@ default RpcPriority getRPCPriority() {
    * state asynchronously. The returned value depends on the type of statement:
    *
    * <ul>
-   *   <li>Queries will return an {@link AsyncResultSet}
-   *   <li>DML statements will return an {@link ApiFuture} with an update count that is done when
-   *       the DML statement has been applied successfully, or that throws an {@link
+   *   <li>Queries and DML statements with returning clause will return an {@link AsyncResultSet}.
+   *   <li>Simple DML statements will return an {@link ApiFuture} with an update count that is done
+   *       when the DML statement has been applied successfully, or that throws an {@link
    *       ExecutionException} if the DML statement failed.
    *   <li>DDL statements will return an {@link ApiFuture} containing a {@link Void} that is done
    *       when the DDL statement has been applied successfully, or that throws an {@link
@@ -894,31 +894,33 @@ default RpcPriority getRPCPriority() {
   AsyncStatementResult executeAsync(Statement statement);
 
   /**
-   * Executes the given statement as a query and returns the result as a {@link ResultSet}. This
-   * method blocks and waits for a response from Spanner. If the statement does not contain a valid
-   * query, the method will throw a {@link SpannerException}.
+   * Executes the given statement (a query or a DML statement with returning clause) and returns the
+   * result as a {@link ResultSet}. This method blocks and waits for a response from Spanner. If the
+   * statement does not contain a valid query or a DML statement with returning clause, the method
+   * will throw a {@link SpannerException}.
    *
-   * @param query The query statement to execute
+   * @param query The query statement or DML statement with returning clause to execute
    * @param options the options to configure the query
-   * @return a {@link ResultSet} with the results of the query
+   * @return a {@link ResultSet} with the results of the statement
    */
   ResultSet executeQuery(Statement query, QueryOption... options);
 
   /**
-   * Executes the given statement asynchronously as a query and returns the result as an {@link
-   * AsyncResultSet}. This method is guaranteed to be non-blocking. If the statement does not
-   * contain a valid query, the method will throw a {@link SpannerException}.
+   * Executes the given statement (a query or a DML statement with returning clause) asynchronously
+   * and returns the result as an {@link AsyncResultSet}. This method is guaranteed to be
+   * non-blocking. If the statement does not contain a valid query or a DML statement with returning
+   * clause, the method will throw a {@link SpannerException}.
    *
    * <p>See {@link AsyncResultSet#setCallback(java.util.concurrent.Executor,
    * com.google.cloud.spanner.AsyncResultSet.ReadyCallback)} for more information on how to consume
-   * the results of the query asynchronously.
+   * the results of the statement asynchronously.
    *
    * <p>It is also possible to consume the returned {@link AsyncResultSet} in the same way as a
    * normal {@link ResultSet}, i.e. in a while-loop calling {@link AsyncResultSet#next()}.
    *
-   * @param query The query statement to execute
+   * @param query The query statement or DML statement with returning clause to execute
    * @param options the options to configure the query
-   * @return an {@link AsyncResultSet} with the results of the query
+   * @return an {@link AsyncResultSet} with the results of the statement
    */
   AsyncResultSet executeQueryAsync(Statement query, QueryOption... options);
 
@@ -951,8 +953,8 @@ default RpcPriority getRPCPriority() {
   ResultSet analyzeQuery(Statement query, QueryAnalyzeMode queryMode);
 
   /**
-   * Executes the given statement as a DML statement. If the statement does not contain a valid DML
-   * statement, the method will throw a {@link SpannerException}.
+   * Executes the given statement as a simple DML statement. If the statement does not contain a
+   * valid DML statement, the method will throw a {@link SpannerException}.
    *
    * @param update The update statement to execute
    * @return the number of records that were inserted/updated/deleted by this statement
@@ -972,8 +974,9 @@ default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeM
   }
 
   /**
-   * Executes the given statement asynchronously as a DML statement. If the statement does not
-   * contain a valid DML statement, the method will throw a {@link SpannerException}.
+   * Executes the given statement asynchronously as a simple DML statement. If the statement does
+   * not contain a simple DML statement, the method will throw a {@link SpannerException}. A DML
+   * statement with returning clause will throw a {@link SpannerException}.
    *
    * <p>This method is guaranteed to be non-blocking.
    *
@@ -984,8 +987,9 @@ default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeM
   ApiFuture<Long> executeUpdateAsync(Statement update);
 
   /**
-   * Executes a list of DML statements in a single request. The statements will be executed in order
-   * and the semantics is the same as if each statement is executed by {@link
+   * Executes a list of DML statements (can be simple DML statements or DML statements with
+   * returning clause) in a single request. The statements will be executed in order and the
+   * semantics is the same as if each statement is executed by {@link
    * Connection#executeUpdate(Statement)} in a loop. This method returns an array of long integers,
    * each representing the number of rows modified by each statement.
    *
@@ -1006,8 +1010,9 @@ default ResultSetStats analyzeUpdate(Statement update, QueryAnalyzeMode analyzeM
   long[] executeBatchUpdate(Iterable<Statement> updates);
 
   /**
-   * Executes a list of DML statements in a single request. The statements will be executed in order
-   * and the semantics is the same as if each statement is executed by {@link
+   * Executes a list of DML statements (can be simple DML statements or DML statements with
+   * returning clause) in a single request. The statements will be executed in order and the
+   * semantics is the same as if each statement is executed by {@link
    * Connection#executeUpdate(Statement)} in a loop. This method returns an {@link ApiFuture} that
    * contains an array of long integers, each representing the number of rows modified by each
    * statement.

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,22 +935,36 @@ private ResultSet parseAndExecuteQuery(
         case QUERY:
           return internalExecuteQuery(parsedStatement, analyzeMode, options);
         case UPDATE:
+          if (parsedStatement.hasReturningClause()) {
+            // Cannot execute DML statement with returning clause 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:
       }
     }
     throw SpannerExceptionFactory.newSpannerException(
         ErrorCode.INVALID_ARGUMENT,
-        "Statement is not a query: " + parsedStatement.getSqlWithoutComments());
+        "Statement is not a query or DML with returning clause: "
+            + parsedStatement.getSqlWithoutComments());
   }
 
   private AsyncResultSet parseAndExecuteQueryAsync(
       Statement query, AnalyzeMode analyzeMode, QueryOption... options) {
     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,14 +977,28 @@ private AsyncResultSet parseAndExecuteQueryAsync(
         case QUERY:
           return internalExecuteQueryAsync(parsedStatement, analyzeMode, options);
         case UPDATE:
+          if (parsedStatement.hasReturningClause()) {
+            // Cannot execute DML statement with returning clause 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:
       }
     }
     throw SpannerExceptionFactory.newSpannerException(
         ErrorCode.INVALID_ARGUMENT,
-        "Statement is not a query: " + parsedStatement.getSqlWithoutComments());
+        "Statement is not a query or DML with returning clause: "
+            + parsedStatement.getSqlWithoutComments());
   }
 
   @Override
@@ -974,6 +1009,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 using executeUpdate: "
+                    + parsedStatement.getSqlWithoutComments()
+                    + ". Please use executeQuery instead.");
+          }
           return get(internalExecuteUpdateAsync(parsedStatement));
         case CLIENT_SIDE:
         case QUERY:
@@ -995,6 +1037,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 using executeUpdateAsync: "
+                    + parsedStatement.getSqlWithoutComments()
+                    + ". Please use executeQueryAsync instead.");
+          }
           return internalExecuteUpdateAsync(parsedStatement);
         case CLIENT_SIDE:
         case QUERY:
@@ -1141,8 +1190,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 +1204,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(