feat: support DML with Returning clause in Connection API

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(