feat: add support for CommitStats

src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java

@@ -17,6 +17,8 @@
 package com.google.cloud.spanner.jdbc;
 
 import com.google.cloud.spanner.AbortedException;
+import com.google.cloud.spanner.CommitResponse;
+import com.google.cloud.spanner.CommitStats;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.ResultSet;
 import java.sql.Connection;
@@ -42,6 +44,24 @@ public interface CloudSpannerJdbcConnection extends Connection {
    */
   Timestamp getCommitTimestamp() throws SQLException;
 
+  /**
+   * @return the {@link CommitResponse} of the last read/write transaction. If the last transaction
+   *     was not a read/write transaction, or a read/write transaction that did not return a {@link
+   *     CommitResponse} because the transaction was not committed, the method will throw a {@link
+   *     SQLException}. The {@link CommitResponse} will include {@link CommitStats} if {@link
+   *     #isReturnCommitStats()} is true.
+   */
+  CommitResponse getCommitResponse() throws SQLException;
+
+  /**
+   * 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) throws SQLException;
+
+  /** @return true if this connection requests commit statistics from Cloud Spanner. */
+  boolean isReturnCommitStats() throws SQLException;
+
   /**
    * @return the read {@link Timestamp} of the last read-only transaction. If the last transaction
    *     was not a read-only transaction, or a read-only transaction that did not return a read

src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java

@@ -17,6 +17,7 @@
 package com.google.cloud.spanner.jdbc;
 
 import com.google.api.client.util.Preconditions;
+import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.connection.ConnectionOptions;
@@ -318,6 +319,27 @@ public Timestamp getCommitTimestamp() throws SQLException {
     }
   }
 
+  @Override
+  public CommitResponse getCommitResponse() throws SQLException {
+    checkClosed();
+    try {
+      return getSpannerConnection().getCommitResponse();
+    } catch (SpannerException e) {
+      throw JdbcSqlExceptionFactory.of(e);
+    }
+  }
+
+  @Override
+  public void setReturnCommitStats(boolean returnCommitStats) throws SQLException {
+    checkClosed();
+    getSpannerConnection().setReturnCommitStats(returnCommitStats);
+  }
+
+  public boolean isReturnCommitStats() throws SQLException {
+    checkClosed();
+    return getSpannerConnection().isReturnCommitStats();
+  }
+
   @Override
   public Timestamp getReadTimestamp() throws SQLException {
     checkClosed();

src/test/java/com/google/cloud/spanner/jdbc/JdbcCommitStatsTest.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.jdbc;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.spanner.CommitResponse;
+import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class JdbcCommitStatsTest extends AbstractMockServerTest {
+  @Test
+  public void testDefaultReturnCommitStats() throws SQLException {
+    try (java.sql.Connection connection = createJdbcConnection()) {
+      try (java.sql.ResultSet rs =
+          connection.createStatement().executeQuery("SHOW VARIABLE RETURN_COMMIT_STATS")) {
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getBoolean("RETURN_COMMIT_STATS")).isFalse();
+        assertThat(rs.next()).isFalse();
+      }
+    }
+  }
+
+  @Test
+  public void testReturnCommitStatsInConnectionUrl() throws SQLException {
+    try (java.sql.Connection connection =
+        DriverManager.getConnection(
+            String.format("jdbc:%s;returnCommitStats=true", getBaseUrl()))) {
+      try (java.sql.ResultSet rs =
+          connection.createStatement().executeQuery("SHOW VARIABLE RETURN_COMMIT_STATS")) {
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getBoolean("RETURN_COMMIT_STATS")).isTrue();
+        assertThat(rs.next()).isFalse();
+      }
+    }
+  }
+
+  @Test
+  public void testSetReturnCommitStats() throws SQLException {
+    try (java.sql.Connection connection = createJdbcConnection()) {
+      connection.createStatement().execute("SET RETURN_COMMIT_STATS=true");
+      try (java.sql.ResultSet rs =
+          connection.createStatement().executeQuery("SHOW VARIABLE RETURN_COMMIT_STATS")) {
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getBoolean("RETURN_COMMIT_STATS")).isTrue();
+        assertThat(rs.next()).isFalse();
+      }
+      connection.createStatement().execute("SET RETURN_COMMIT_STATS=false");
+      try (java.sql.ResultSet rs =
+          connection.createStatement().executeQuery("SHOW VARIABLE RETURN_COMMIT_STATS")) {
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getBoolean("RETURN_COMMIT_STATS")).isFalse();
+        assertThat(rs.next()).isFalse();
+      }
+    }
+  }
+
+  @Test
+  public void testSetAndUseReturnCommitStats() throws SQLException {
+    try (CloudSpannerJdbcConnection connection =
+        createJdbcConnection().unwrap(CloudSpannerJdbcConnection.class)) {
+      connection.setReturnCommitStats(true);
+      connection.bufferedWrite(Mutation.newInsertBuilder("FOO").set("ID").to(1L).build());
+      connection.commit();
+      CommitResponse response = connection.getCommitResponse();
+      assertThat(response).isNotNull();
+      assertThat(response.getCommitStats()).isNotNull();
+      assertThat(response.getCommitStats().getMutationCount()).isAtLeast(1L);
+      assertThat(response.getCommitStats().getOverloadDelay()).isNotNull();
+    }
+  }
+
+  @Test
+  public void testSetAndUseReturnCommitStatsUsingSql() throws SQLException {
+    try (java.sql.Connection connection = createJdbcConnection()) {
+      connection.createStatement().execute("SET RETURN_COMMIT_STATS=true");
+      // Use a Mutation as the mock server only returns a non-zero mutation count for mutations, and
+      // not for DML statements.
+      connection
+          .unwrap(CloudSpannerJdbcConnection.class)
+          .bufferedWrite(Mutation.newInsertBuilder("FOO").set("ID").to(1L).build());
+      connection.commit();
+      try (ResultSet rs =
+          connection.createStatement().executeQuery("SHOW VARIABLE COMMIT_RESPONSE")) {
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getTimestamp("COMMIT_TIMESTAMP")).isNotNull();
+        assertThat(rs.getLong("MUTATION_COUNT")).isAtLeast(1L);
+        assertThat(rs.getString("OVERLOAD_DELAY")).isNotNull();
+        assertThat(rs.next()).isFalse();
+      }
+    }
+  }
+}