feat!: add CommitStats to Connection API (#608)

* feat!: add support for CommitStats Adds support for returning CommitStats from read/write transactions.

Author: olavloite

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

@@ -476,6 +476,23 @@
  <method>com.google.cloud.spanner.CommitResponse getCommitResponse()</method>
  </difference>
 
+ <!-- Support for Commit Stats in Connection API -->
+ <difference>
+ <differenceType>7012</differenceType>
+ <className>com/google/cloud/spanner/connection/Connection</className>
+ <method>com.google.cloud.spanner.CommitResponse getCommitResponse()</method>
+ </difference>
+ <difference>
+ <differenceType>7012</differenceType>
+ <className>com/google/cloud/spanner/connection/Connection</className>
+ <method>boolean isReturnCommitStats()</method>
+ </difference>
+ <difference>
+ <differenceType>7012</differenceType>
+ <className>com/google/cloud/spanner/connection/Connection</className>
+ <method>void setReturnCommitStats(boolean)</method>
+ </difference>
+
  <!-- PITR -->
  <difference>
  <differenceType>7013</differenceType>

google-cloud-spanner/src/main/java/com/google/cloud/spanner/CommitResponse.java

@@ -41,6 +41,11 @@ public Timestamp getCommitTimestamp() {
  return Timestamp.fromProto(proto.getCommitTimestamp());
  }
 
+ /** @return true if the {@link CommitResponse} includes {@link CommitStats} */
+ public boolean hasCommitStats() {
+ return proto.hasCommitStats();
+ }
+
  /**
  * Commit statistics are returned by a read/write transaction if specifically requested by passing
  * in {@link Options#commitStats()} to the transaction.

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

@@ -22,6 +22,7 @@
 import com.google.cloud.spanner.AbortedDueToConcurrentModificationException;
 import com.google.cloud.spanner.AbortedException;
 import com.google.cloud.spanner.AsyncResultSet;
+import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Options.QueryOption;
@@ -438,6 +439,15 @@ public interface Connection extends AutoCloseable {
  */
  String getOptimizerVersion();
 
+ /**
+ * Sets whether this connection should request commit statistics from Cloud Spanner for read/write
+ * transactions and DML statements in autocommit mode.
+ */
+ void setReturnCommitStats(boolean returnCommitStats);
+
+ /** @return true if this connection requests commit statistics from Cloud Spanner */
+ boolean isReturnCommitStats();
+
  /**
  * Commits the current transaction of this connection. All mutations that have been buffered
  * during the current transaction will be written to the database.
@@ -623,15 +633,26 @@ public interface Connection extends AutoCloseable {
 
  /**
  * @return the commit timestamp of the last {@link TransactionMode#READ_WRITE_TRANSACTION}
- * transaction. This method will throw a {@link SpannerException} if there is no last {@link
- * TransactionMode#READ_WRITE_TRANSACTION} transaction (i.e. the last transaction was a {@link
- * TransactionMode#READ_ONLY_TRANSACTION}), or if the last {@link
- * TransactionMode#READ_WRITE_TRANSACTION} transaction rolled back. It will also throw a
- * {@link SpannerException} if the last {@link TransactionMode#READ_WRITE_TRANSACTION}
- * transaction was empty when committed.
+ * transaction. This method throws a {@link SpannerException} if there is no last {@link
+ * TransactionMode#READ_WRITE_TRANSACTION} transaction. That is, if the last transaction was a
+ * {@link TransactionMode#READ_ONLY_TRANSACTION}), or if the last {@link
+ * TransactionMode#READ_WRITE_TRANSACTION} transaction rolled back. It also throws a {@link
+ * SpannerException} if the last {@link TransactionMode#READ_WRITE_TRANSACTION} transaction
+ * was empty when committed.
  */
  Timestamp getCommitTimestamp();
 
+ /**
+ * @return the {@link CommitResponse} of the last {@link TransactionMode#READ_WRITE_TRANSACTION}
+ * transaction. This method throws a {@link SpannerException} if there is no last {@link
+ * TransactionMode#READ_WRITE_TRANSACTION} transaction. That is, if the last transaction was a
+ * {@link TransactionMode#READ_ONLY_TRANSACTION}), or if the last {@link
+ * TransactionMode#READ_WRITE_TRANSACTION} transaction rolled back. It also throws a {@link
+ * SpannerException} if the last {@link TransactionMode#READ_WRITE_TRANSACTION} transaction
+ * was empty when committed.
+ */
+ CommitResponse getCommitResponse();
+
  /**
  * Starts a new DDL batch on this connection. A DDL batch allows several DDL statements to be
  * grouped into a batch that can be executed as a group. DDL statements that are issued during the

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

@@ -22,6 +22,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.AsyncResultSet;
+import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Mutation;
@@ -179,6 +180,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
  private DatabaseClient dbClient;
  private boolean autocommit;
  private boolean readOnly;
+ private boolean returnCommitStats;
 
  private UnitOfWork currentUnitOfWork = null;
  /**
@@ -213,6 +215,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
  this.readOnly = options.isReadOnly();
  this.autocommit = options.isAutocommit();
  this.queryOptions = this.queryOptions.toBuilder().mergeFrom(options.getQueryOptions()).build();
+ this.returnCommitStats = options.isReturnCommitStats();
  this.ddlClient = createDdlClient();
  setDefaultTransactionOptions();
  }
@@ -237,6 +240,7 @@ static UnitOfWorkType of(TransactionMode transactionMode) {
  this.dbClient = dbClient;
  setReadOnly(options.isReadOnly());
  setAutocommit(options.isAutocommit());
+ setReturnCommitStats(options.isReturnCommitStats());
  setDefaultTransactionOptions();
  }
 
@@ -580,6 +584,31 @@ Timestamp getCommitTimestampOrNull() {
  : this.currentUnitOfWork.getCommitTimestampOrNull();
  }
 
+ @Override
+ public CommitResponse getCommitResponse() {
+ ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+ ConnectionPreconditions.checkState(
+ this.currentUnitOfWork != null, "There is no transaction on this connection");
+ return this.currentUnitOfWork.getCommitResponse();
+ }
+
+ CommitResponse getCommitResponseOrNull() {
+ ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+ return this.currentUnitOfWork == null ? null : this.currentUnitOfWork.getCommitResponseOrNull();
+ }
+
+ @Override
+ public void setReturnCommitStats(boolean returnCommitStats) {
+ ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+ this.returnCommitStats = returnCommitStats;
+ }
+
+ @Override
+ public boolean isReturnCommitStats() {
+ ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+ return this.returnCommitStats;
+ }
+
  /** Resets this connection to its default transaction options. */
  private void setDefaultTransactionOptions() {
  if (transactionStack.isEmpty()) {
@@ -954,6 +983,7 @@ private UnitOfWork createNewUnitOfWork() {
  .setReadOnly(isReadOnly())
  .setReadOnlyStaleness(readOnlyStaleness)
  .setAutocommitDmlMode(autocommitDmlMode)
+ .setReturnCommitStats(returnCommitStats)
  .setStatementTimeout(statementTimeout)
  .withStatementExecutor(statementExecutor)
  .build();
@@ -970,6 +1000,7 @@ private UnitOfWork createNewUnitOfWork() {
  return ReadWriteTransaction.newBuilder()
  .setDatabaseClient(dbClient)
  .setRetryAbortsInternally(retryAbortsInternally)
+ .setReturnCommitStats(returnCommitStats)
  .setTransactionRetryListeners(transactionRetryListeners)
  .setStatementTimeout(statementTimeout)
  .withStatementExecutor(statementExecutor)

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

@@ -155,6 +155,7 @@ public String[] getValidValues() {
  private static final String DEFAULT_NUM_CHANNELS = null;
  private static final String DEFAULT_USER_AGENT = null;
  private static final String DEFAULT_OPTIMIZER_VERSION = "";
+ private static final boolean DEFAULT_RETURN_COMMIT_STATS = false;
  private static final boolean DEFAULT_LENIENT = false;
 
  private static final String PLAIN_TEXT_PROTOCOL = "http:";
@@ -229,6 +230,7 @@ public String[] getValidValues() {
  ConnectionProperty.createStringProperty(
  OPTIMIZER_VERSION_PROPERTY_NAME,
  "Sets the default query optimizer version to use for this connection."),
+ ConnectionProperty.createBooleanProperty("returnCommitStats", "", false),
  ConnectionProperty.createBooleanProperty(
  LENIENT_PROPERTY_NAME,
  "Silently ignore unknown properties in the connection string/properties (true/false)",
@@ -456,6 +458,7 @@ public static Builder newBuilder() {
  private final Integer maxSessions;
  private final String userAgent;
  private final QueryOptions queryOptions;
+ private final boolean returnCommitStats;
 
  private final boolean autocommit;
  private final boolean readOnly;
@@ -485,6 +488,7 @@ private ConnectionOptions(Builder builder) {
  QueryOptions.Builder queryOptionsBuilder = QueryOptions.newBuilder();
  queryOptionsBuilder.setOptimizerVersion(parseOptimizerVersion(this.uri));
  this.queryOptions = queryOptionsBuilder.build();
+ this.returnCommitStats = parseReturnCommitStats(this.uri);
 
  this.host =
  matcher.group(Builder.HOST_GROUP) == null
@@ -634,6 +638,12 @@ static String parseOptimizerVersion(String uri) {
  return value != null ? value : DEFAULT_OPTIMIZER_VERSION;
  }
 
+ @VisibleForTesting
+ static boolean parseReturnCommitStats(String uri) {
+ String value = parseUriProperty(uri, "returnCommitStats");
+ return value != null ? Boolean.valueOf(value) : false;
+ }
+
  @VisibleForTesting
  static boolean parseLenient(String uri) {
  String value = parseUriProperty(uri, LENIENT_PROPERTY_NAME);
@@ -823,6 +833,11 @@ QueryOptions getQueryOptions() {
  return queryOptions;
  }
 
+ /** Whether connections created by this {@link ConnectionOptions} return commit stats. */
+ public boolean isReturnCommitStats() {
+ return returnCommitStats;
+ }
+
  /** Interceptors that should be executed after each statement */
  List<StatementExecutionInterceptor> getStatementExecutionInterceptors() {
  return statementExecutionInterceptors;

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

@@ -56,6 +56,8 @@ interface ConnectionStatementExecutor {
 
  StatementResult statementShowCommitTimestamp();
 
+ StatementResult statementShowCommitResponse();
+
  StatementResult statementSetReadOnlyStaleness(TimestampBound staleness);
 
  StatementResult statementShowReadOnlyStaleness();
@@ -64,6 +66,10 @@ interface ConnectionStatementExecutor {
 
  StatementResult statementShowOptimizerVersion();
 
+ StatementResult statementSetReturnCommitStats(Boolean returnCommitStats);
+
+ StatementResult statementShowReturnCommitStats();
+
  StatementResult statementBeginTransaction();
 
  StatementResult statementCommit();

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

@@ -27,26 +27,37 @@
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_READONLY;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_READ_ONLY_STALENESS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_RETRY_ABORTS_INTERNALLY;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_RETURN_COMMIT_STATS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_STATEMENT_TIMEOUT;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SET_TRANSACTION_MODE;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_AUTOCOMMIT;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_AUTOCOMMIT_DML_MODE;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_COMMIT_RESPONSE;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_COMMIT_TIMESTAMP;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_OPTIMIZER_VERSION;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_READONLY;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_READ_ONLY_STALENESS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_READ_TIMESTAMP;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_RETRY_ABORTS_INTERNALLY;
+import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_RETURN_COMMIT_STATS;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.SHOW_STATEMENT_TIMEOUT;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.START_BATCH_DDL;
 import static com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType.START_BATCH_DML;
 import static com.google.cloud.spanner.connection.StatementResultImpl.noResult;
 import static com.google.cloud.spanner.connection.StatementResultImpl.resultSet;
 
+import com.google.cloud.spanner.CommitResponse;
+import com.google.cloud.spanner.CommitStats;
+import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.ResultSets;
+import com.google.cloud.spanner.Struct;
 import com.google.cloud.spanner.TimestampBound;
+import com.google.cloud.spanner.Type;
+import com.google.cloud.spanner.Type.StructField;
 import com.google.cloud.spanner.connection.ReadOnlyStalenessUtil.DurationValueGetter;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.Duration;
+import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -170,6 +181,28 @@ public StatementResult statementShowCommitTimestamp() {
  "COMMIT_TIMESTAMP", getConnection().getCommitTimestampOrNull(), SHOW_COMMIT_TIMESTAMP);
  }
 
+ @Override
+ public StatementResult statementShowCommitResponse() {
+ CommitResponse response = getConnection().getCommitResponseOrNull();
+ CommitStats stats = null;
+ if (response != null && response.hasCommitStats()) {
+ stats = response.getCommitStats();
+ }
+ ResultSet resultSet =
+ ResultSets.forRows(
+ Type.struct(
+ StructField.of("COMMIT_TIMESTAMP", Type.timestamp()),
+ StructField.of("MUTATION_COUNT", Type.int64())),
+ Arrays.asList(
+ Struct.newBuilder()
+ .set("COMMIT_TIMESTAMP")
+ .to(response == null ? null : response.getCommitTimestamp())
+ .set("MUTATION_COUNT")
+ .to(stats == null ? null : stats.getMutationCount())
+ .build()));
+ return StatementResultImpl.of(resultSet, SHOW_COMMIT_RESPONSE);
+ }
+
  @Override
  public StatementResult statementSetReadOnlyStaleness(TimestampBound staleness) {
  getConnection().setReadOnlyStaleness(staleness);
@@ -197,6 +230,18 @@ public StatementResult statementShowOptimizerVersion() {
  "OPTIMIZER_VERSION", getConnection().getOptimizerVersion(), SHOW_OPTIMIZER_VERSION);
  }
 
+ @Override
+ public StatementResult statementSetReturnCommitStats(Boolean returnCommitStats) {
+ getConnection().setReturnCommitStats(returnCommitStats);
+ return noResult(SET_RETURN_COMMIT_STATS);
+ }
+
+ @Override
+ public StatementResult statementShowReturnCommitStats() {
+ return resultSet(
+ "RETURN_COMMIT_STATS", getConnection().isReturnCommitStats(), SHOW_RETURN_COMMIT_STATS);
+ }
+
  @Override
  public StatementResult statementBeginTransaction() {
  getConnection().beginTransaction();

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

@@ -20,6 +20,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Mutation;
@@ -168,6 +169,17 @@ public Timestamp getCommitTimestampOrNull() {
  return null;
  }
 
+ @Override
+ public CommitResponse getCommitResponse() {
+ throw SpannerExceptionFactory.newSpannerException(
+ ErrorCode.FAILED_PRECONDITION, "There is no commit response available for DDL batches.");
+ }
+
+ @Override
+ public CommitResponse getCommitResponseOrNull() {
+ return null;
+ }
+
  @Override
  public ApiFuture<Void> executeDdlAsync(ParsedStatement ddl) {
  ConnectionPreconditions.checkState(

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

@@ -21,6 +21,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.SettableApiFuture;
 import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.Options.QueryOption;
@@ -119,6 +120,17 @@ public Timestamp getCommitTimestampOrNull() {
  return null;
  }
 
+ @Override
+ public CommitResponse getCommitResponse() {
+ throw SpannerExceptionFactory.newSpannerException(
+ ErrorCode.FAILED_PRECONDITION, "There is no commit response available for DML batches.");
+ }
+
+ @Override
+ public CommitResponse getCommitResponseOrNull() {
+ return null;
+ }
+
  @Override
  public ApiFuture<Void> executeDdlAsync(ParsedStatement ddl) {
  throw SpannerExceptionFactory.newSpannerException(