Set setFailIfPoolExhausted in SessionPoolOptions for SpannerAccessor

nielm · Jul 5, 2024 · efee92b · efee92b
1 parent 3914ad0
commit efee92b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/...ogle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java b/...ogle-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessor.java
@@ -30,6 +30,7 @@
 import com.google.cloud.spanner.DatabaseAdminClient;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.SessionPoolOptions;
 import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
@@ -229,7 +230,9 @@ static SpannerOptions buildSpannerOptions(SpannerConfig spannerConfig) {
     if (credentials != null && credentials.get() != null) {
       builder.setCredentials(credentials.get());
     }
-
+    SessionPoolOptions sessionPoolOptions =
+        SessionPoolOptions.newBuilder().setFailIfPoolExhausted().build();
+    builder.setSessionPoolOption(sessionPoolOptions);
     return builder.build();
   }
 

diff --git a/...-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessorTest.java b/...-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/SpannerAccessorTest.java
@@ -18,6 +18,7 @@
 package org.apache.beam.sdk.io.gcp.spanner;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -163,5 +164,6 @@ public void testBuildSpannerOptionsWithCredential() {
     assertEquals("project", options.getProjectId());
     assertEquals("test-role", options.getDatabaseRole());
     assertEquals(testCredential, options.getCredentials());
+    assertNotNull(options.getSessionPoolOptions());
   }
 }