diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml
index c750b2a7f40..f65bfd94332 100644
--- a/google-cloud-spanner/clirr-ignored-differences.xml
+++ b/google-cloud-spanner/clirr-ignored-differences.xml
@@ -6,6 +6,11 @@
com/google/cloud/spanner/connection/Connection
com.google.cloud.spanner.Dialect getDialect()
+
+ 7012
+ com/google/cloud/spanner/BatchReadOnlyTransaction
+ void cleanup()
+
8001
com/google/cloud/spanner/connection/StatementParser
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java
index c84bef77cf8..fe31a3b0101 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java
@@ -214,9 +214,14 @@ public ResultSet execute(Partition partition) throws SpannerException {
partition.getPartitionToken());
}
+ /**
+ * Closes the session as part of the cleanup. It is the responsibility of the caller to make a
+ * call to this method once the transaction completes execution across all the channels (which
+ * is understandably hard to identify). It is okay if the caller does not call the method
+ * because the backend will anyways clean up the unused session.
+ */
@Override
- public void close() {
- super.close();
+ public void cleanup() {
session.close();
}
}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchReadOnlyTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchReadOnlyTransaction.java
index 9a9613f3247..03b08a11730 100644
--- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchReadOnlyTransaction.java
+++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchReadOnlyTransaction.java
@@ -197,4 +197,12 @@ List partitionQuery(
* BatchTransactionId guarantees the subsequent read/query to be executed at the same timestamp.
*/
BatchTransactionId getBatchTransactionId();
+
+ /**
+ * Closes the session as part of the cleanup. It is the responsibility of the caller to make a
+ * call to this method once the transaction completes execution across all the channels (which is
+ * understandably hard to identify). It is okay if the caller does not call the method because the
+ * backend will anyways clean up the unused session.
+ */
+ default void cleanup() {}
}
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
index 20e8c6b8ea6..21c9d3eaa77 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
@@ -54,6 +54,7 @@
import com.google.common.util.concurrent.SettableFuture;
import com.google.protobuf.AbstractMessage;
import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.DeleteSessionRequest;
import com.google.spanner.v1.ExecuteBatchDmlRequest;
import com.google.spanner.v1.ExecuteSqlRequest;
import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
@@ -1630,16 +1631,20 @@ public void testBackendPartitionQueryOptions() {
try (ResultSet rs = transaction.execute(partitions.get(0))) {
// Just iterate over the results to execute the query.
while (rs.next()) {}
+ } finally {
+ transaction.cleanup();
}
- // Check that the last query was executed using a custom optimizer version and statistics
- // package.
+ // Check if the last query executed is a DeleteSessionRequest and the second last query
+ // executed is a ExecuteSqlRequest and was executed using a custom optimizer version and
+ // statistics package.
List requests = mockSpanner.getRequests();
- assertThat(requests).isNotEmpty();
- assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class);
- ExecuteSqlRequest request = (ExecuteSqlRequest) requests.get(requests.size() - 1);
- assertThat(request.getQueryOptions()).isNotNull();
- assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("1");
- assertThat(request.getQueryOptions().getOptimizerStatisticsPackage())
+ assert requests.size() >= 2 : "required to have at least 2 requests";
+ assertThat(requests.get(requests.size() - 1)).isInstanceOf(DeleteSessionRequest.class);
+ assertThat(requests.get(requests.size() - 2)).isInstanceOf(ExecuteSqlRequest.class);
+ ExecuteSqlRequest executeSqlRequest = (ExecuteSqlRequest) requests.get(requests.size() - 2);
+ assertThat(executeSqlRequest.getQueryOptions()).isNotNull();
+ assertThat(executeSqlRequest.getQueryOptions().getOptimizerVersion()).isEqualTo("1");
+ assertThat(executeSqlRequest.getQueryOptions().getOptimizerStatisticsPackage())
.isEqualTo("custom-package");
}
}