diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index 7ca564912c43..9766e8911a07 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -40,7 +40,7 @@ * * @see Google Cloud BigQuery */ -public interface BigQuery extends Service { +public interface BigQuery extends Service, AutoCloseable { /** * Fields of a BigQuery Dataset resource. @@ -1821,4 +1821,14 @@ Object queryWithTimeout( * represents the subset of granted permissions. */ List testIamPermissions(TableId table, List permissions, IAMOption... options); + + /** + * Closes any background resources and transport channels held by this service. + * + *

The default implementation does nothing. Implementations that manage background resources + * (such as gRPC channels or storage clients) should override this method to release them + * deterministically. + */ + @Override + default void close() {} } diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 44d951fdb9cf..2d24e7bb82c1 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -504,6 +504,7 @@ private void closeClient() { private transient ConcurrentHashMap bqReadClients; private transient boolean isGlobalClientUserProvided; + private transient volatile boolean closed = false; /** * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming @@ -527,6 +528,9 @@ BigQueryReadClient getBigQueryReadClient() { * @throws BigQueryException if initializing the storage read client fails */ BigQueryReadClient getBigQueryReadClient(String location) { + if (closed) { + throw new IllegalStateException("BigQuery service has been closed"); + } String cacheKey = location != null ? location.toLowerCase() : "global"; if (bqReadClients == null) { synchronized (this) { @@ -543,6 +547,9 @@ BigQueryReadClient getBigQueryReadClient(String location) { return client; } synchronized (this) { + if (closed) { + throw new IllegalStateException("BigQuery service has been closed"); + } client = bqReadClients.get(cacheKey); if (client == null && isGlobalClientUserProvided) { client = bqReadClients.get("global"); @@ -552,6 +559,10 @@ BigQueryReadClient getBigQueryReadClient(String location) { configureReadSettings(settingsBuilder, getOptions()); try { client = BigQueryReadClient.create(settingsBuilder.build()); + if (closed) { + client.close(); + throw new IllegalStateException("BigQuery service has been closed"); + } if (bqReadClients.size() < MAX_CACHED_READ_CLIENTS) { bqReadClients.put(cacheKey, client); } @@ -583,6 +594,30 @@ void setBigQueryReadClient(String location, BigQueryReadClient client) { } } + /** + * Closes any background resources and transport channels held by this {@link BigQueryImpl}, + * including the underlying {@link BigQueryReadClient} instances used for Arrow query streaming. + */ + @Override + public void close() { + synchronized (this) { + if (closed) { + return; + } + closed = true; + } + if (bqReadClients != null) { + for (BigQueryReadClient client : bqReadClients.values()) { + try { + client.close(); + } catch (Exception e) { + // Ignore exceptions during teardown + } + } + bqReadClients.clear(); + } + } + /** * Configures a {@link BigQueryReadSettings.Builder} with credentials, header provider, and * universe domain mapped from the given {@link BigQueryOptions}. diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 79028a7c8140..5a90822cb25b 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -4260,4 +4260,70 @@ void testTestIamPermissionsWhenNoPermissionsGranted() throws IOException { .testIamPermissionsSkipExceptionTranslation( resourceId, checkedPermissions, EMPTY_RPC_OPTIONS); } + + @Test + void testCloseClosesBigQueryReadClient() { + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient); + + bigquery.close(); + + verify(mockReadClient, times(1)).close(); + } + + @Test + void testCloseIsIdempotent() { + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient); + + bigquery.close(); + bigquery.close(); + + verify(mockReadClient, times(1)).close(); + } + + @Test + void testCloseWithoutReadClientDoesNotThrow() { + bigquery = options.getService(); + bigquery.close(); + } + + @Test + void testTryWithResources() { + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + try (BigQuery bq = options.getService()) { + ((BigQueryImpl) bq).setBigQueryReadClient(mockReadClient); + assertNotNull(bq); + } + verify(mockReadClient, times(1)).close(); + } + + @Test + void testGetBigQueryReadClientAfterCloseThrows() { + bigquery = options.getService(); + bigquery.close(); + assertThrows( + IllegalStateException.class, () -> ((BigQueryImpl) bigquery).getBigQueryReadClient()); + } + + @Test + void testCloseClosesAllRegionalBigQueryReadClients() { + BigQueryReadClient mockReadClientUs = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + BigQueryReadClient mockReadClientEu = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient("us", mockReadClientUs); + ((BigQueryImpl) bigquery).setBigQueryReadClient("eu", mockReadClientEu); + + bigquery.close(); + + verify(mockReadClientUs, times(1)).close(); + verify(mockReadClientEu, times(1)).close(); + } }