Skip to content

Commit ab9669a

Browse files
keshavdandevakeshavdandeva
authored
feat(bigquery-jdbc): add EnableProjectDiscovery connection property for metadata methods (#13344)
b/499078725 This PR implements the `EnableProjectDiscovery` connection property into the BigQuery JDBC driver. Enabling this property (default `false`) allows JDBC database metadata methods (like `getCatalogs()` and `getSchemas()`) to discover and query datasets across all Google Cloud projects accessible to the client credentials, rather than being confined to the single default `ProjectId` specified in the connection URL. #### Changes - **Connection Parameter Parsing**: Added parsing for `EnableProjectDiscovery` in `BigQueryJdbcUrlUtility` and configured it in `DataSource`. - **SDK Integration**: Implemented `BigQueryConnection.getDiscoveredProjects()` to utilize the `BigQuery.listProjects()` core SDK method - **Resilience and Caching**: - Connection-scoped caching is implemented via `discoveredProjectsCache`. - **Metadata Integration**: Integrated discovered projects inside `BigQueryDatabaseMetaData.getAccessibleCatalogNames()` so that catalog and schema metadata fetches automatically query across all accessible projects when `EnableProjectDiscovery` is enabled.
1 parent 87dda5d commit ab9669a

7 files changed

Lines changed: 316 additions & 17 deletions

File tree

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.google.api.gax.core.CredentialsProvider;
2020
import com.google.api.gax.core.FixedCredentialsProvider;
2121
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
22+
import com.google.api.gax.paging.Page;
2223
import com.google.api.gax.retrying.RetrySettings;
2324
import com.google.api.gax.rpc.FixedHeaderProvider;
2425
import com.google.api.gax.rpc.HeaderProvider;
@@ -32,6 +33,7 @@
3233
import com.google.cloud.bigquery.DatasetId;
3334
import com.google.cloud.bigquery.Job;
3435
import com.google.cloud.bigquery.JobInfo;
36+
import com.google.cloud.bigquery.Project;
3537
import com.google.cloud.bigquery.QueryJobConfiguration;
3638
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
3739
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
@@ -42,6 +44,7 @@
4244
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
4345
import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
4446
import com.google.cloud.http.HttpTransportOptions;
47+
import com.google.common.collect.ImmutableList;
4548
import com.google.common.collect.ImmutableSortedSet;
4649
import java.io.IOException;
4750
import java.io.InputStream;
@@ -122,6 +125,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
122125
BigQueryJdbcUrlUtility.SWA_APPEND_ROW_COUNT_PROPERTY_NAME,
123126
BigQueryJdbcUrlUtility.SWA_ACTIVATION_ROW_COUNT_PROPERTY_NAME,
124127
BigQueryJdbcUrlUtility.FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME,
128+
BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME,
125129
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME,
126130
BigQueryJdbcUrlUtility.SSL_TRUST_STORE_PROPERTY_NAME,
127131
BigQueryJdbcUrlUtility.MAX_BYTES_BILLED_PROPERTY_NAME,
@@ -171,6 +175,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
171175
int highThroughputMinTableSize;
172176
int highThroughputActivationRatio;
173177
boolean enableSession;
178+
boolean enableProjectDiscovery;
179+
private List<String> discoveredProjectsCache;
174180
boolean unsupportedHTAPIFallback;
175181
boolean useQueryCache;
176182
String queryDialect;
@@ -346,6 +352,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
346352
this.additionalProjects = ds.getAdditionalProjects();
347353

348354
this.filterTablesOnDefaultDataset = ds.getFilterTablesOnDefaultDataset();
355+
this.enableProjectDiscovery = ds.getEnableProjectDiscovery();
349356
this.requestGoogleDriveScope = ds.getRequestGoogleDriveScope();
350357
this.metadataFetchThreadCount = ds.getMetadataFetchThreadCount();
351358
this.requestReason = ds.getRequestReason();
@@ -1320,6 +1327,29 @@ private boolean checkIsReadOnlyTokenUsed(Map<String, String> authProps) {
13201327
return false;
13211328
}
13221329

1330+
public boolean isEnableProjectDiscovery() {
1331+
return this.enableProjectDiscovery;
1332+
}
1333+
1334+
public synchronized List<String> getDiscoveredProjects() throws SQLException {
1335+
if (this.discoveredProjectsCache != null) {
1336+
return this.discoveredProjectsCache;
1337+
}
1338+
1339+
try {
1340+
BigQuery bigQuery = getBigQuery();
1341+
List<String> projects = new ArrayList<>();
1342+
Page<Project> projectPage = bigQuery.listProjects();
1343+
for (Project project : projectPage.iterateAll()) {
1344+
projects.add(project.getProjectId());
1345+
}
1346+
this.discoveredProjectsCache = ImmutableList.copyOf(projects);
1347+
} catch (Exception e) {
1348+
throw new BigQueryJdbcException("Failed to list all accessible projects.", e);
1349+
}
1350+
return this.discoveredProjectsCache;
1351+
}
1352+
13231353
@Override
13241354
public <T> T unwrap(Class<T> iface) throws SQLException {
13251355
if (iface.isInstance(this)) {

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java‎

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1996,14 +1996,14 @@ Comparator<FieldValueList> defineGetTablesComparator(FieldList resultSchemaField
19961996
}
19971997

19981998
@Override
1999-
public ResultSet getSchemas() {
1999+
public ResultSet getSchemas() throws SQLException {
20002000
LOG.info("getSchemas() called");
20012001

20022002
return getSchemas(null, null);
20032003
}
20042004

20052005
@Override
2006-
public ResultSet getCatalogs() {
2006+
public ResultSet getCatalogs() throws SQLException {
20072007
LOG.info("getCatalogs() called");
20082008

20092009
final List<String> accessibleCatalogs = getAccessibleCatalogNames();
@@ -3618,7 +3618,7 @@ public RowIdLifetime getRowIdLifetime() {
36183618
}
36193619

36203620
@Override
3621-
public ResultSet getSchemas(String catalog, String schemaPattern) {
3621+
public ResultSet getSchemas(String catalog, String schemaPattern) throws SQLException {
36223622
if ((catalog != null && catalog.isEmpty())
36233623
|| (schemaPattern != null && schemaPattern.isEmpty())) {
36243624
LOG.warning("Returning empty ResultSet as catalog or schemaPattern is an empty string.");
@@ -3641,20 +3641,20 @@ public ResultSet getSchemas(String catalog, String schemaPattern) {
36413641
final FieldList localResultSchemaFields = resultSchemaFields;
36423642
List<String> projectsToScanList = new ArrayList<>();
36433643

3644-
if (catalogParam != null) {
3645-
projectsToScanList.add(catalogParam);
3646-
} else {
3647-
projectsToScanList.addAll(getAccessibleCatalogNames());
3648-
}
3644+
try {
3645+
if (catalogParam != null) {
3646+
projectsToScanList.add(catalogParam);
3647+
} else {
3648+
projectsToScanList.addAll(getAccessibleCatalogNames());
3649+
}
36493650

3650-
if (projectsToScanList.isEmpty()) {
3651-
LOG.info(
3652-
"No valid projects to scan (primary, specified, or additional). Returning empty"
3653-
+ " resultset.");
3654-
return;
3655-
}
3651+
if (projectsToScanList.isEmpty()) {
3652+
LOG.info(
3653+
"No valid projects to scan (primary, specified, or additional). Returning empty"
3654+
+ " resultset.");
3655+
return;
3656+
}
36563657

3657-
try {
36583658
for (String currentProjectToScan : projectsToScanList) {
36593659
if (Thread.currentThread().isInterrupted()) {
36603660
LOG.warning(
@@ -3707,6 +3707,13 @@ public ResultSet getSchemas(String catalog, String schemaPattern) {
37073707

37083708
} catch (Throwable t) {
37093709
LOG.severe("Unexpected error in schema fetcher runnable: " + t.getMessage());
3710+
Exception ex = (t instanceof Exception) ? (Exception) t : new Exception(t);
3711+
try {
3712+
queue.put(BigQueryFieldValueListWrapper.ofError(ex));
3713+
} catch (InterruptedException ie) {
3714+
LOG.warning("Failed to put exception to queue due to interruption.");
3715+
Thread.currentThread().interrupt();
3716+
}
37103717
} finally {
37113718
signalEndOfData(queue, localResultSchemaFields);
37123719
LOG.info("Schema fetcher thread finished.");
@@ -5178,7 +5185,7 @@ private String getCurrentCatalogName() {
51785185
return this.connection.getCatalog();
51795186
}
51805187

5181-
private List<String> getAccessibleCatalogNames() {
5188+
private List<String> getAccessibleCatalogNames() throws SQLException {
51825189
Set<String> accessibleCatalogs = new HashSet<>();
51835190
String primaryCatalog = getCurrentCatalogName();
51845191
if (primaryCatalog != null && !primaryCatalog.isEmpty()) {
@@ -5199,6 +5206,10 @@ private List<String> getAccessibleCatalogNames() {
51995206
}
52005207
}
52015208

5209+
if (this.connection.isEnableProjectDiscovery()) {
5210+
accessibleCatalogs.addAll(this.connection.getDiscoveredProjects());
5211+
}
5212+
52025213
List<String> sortedCatalogs = new ArrayList<>(accessibleCatalogs);
52035214
Collections.sort(sortedCatalogs);
52045215
return sortedCatalogs;

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,8 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
168168
static final String FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME =
169169
"FilterTablesOnDefaultDataset";
170170
static final boolean DEFAULT_FILTER_TABLES_ON_DEFAULT_DATASET_VALUE = false;
171+
static final String ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME = "EnableProjectDiscovery";
172+
static final boolean DEFAULT_ENABLE_PROJECT_DISCOVERY_VALUE = false;
171173
static final String REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME = "RequestGoogleDriveScope";
172174
static final String SSL_TRUST_STORE_PROPERTY_NAME = "SSLTrustStore";
173175
static final String SSL_TRUST_STORE_PWD_PROPERTY_NAME = "SSLTrustStorePwd";
@@ -577,6 +579,13 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
577579
.setDefaultValue(
578580
String.valueOf(DEFAULT_FILTER_TABLES_ON_DEFAULT_DATASET_VALUE))
579581
.build(),
582+
BigQueryConnectionProperty.newBuilder()
583+
.setName(ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME)
584+
.setDescription(
585+
"Enables or disables automatic discovery of all accessible Google Cloud projects. "
586+
+ "When disabled, only the default ProjectId and AdditionalProjects are listed as catalogs.")
587+
.setDefaultValue(String.valueOf(DEFAULT_ENABLE_PROJECT_DISCOVERY_VALUE))
588+
.build(),
580589
BigQueryConnectionProperty.newBuilder()
581590
.setName(REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME)
582591
.setDescription(

‎java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/DataSource.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ public class DataSource implements javax.sql.DataSource {
8484
private Boolean enableWriteAPI;
8585
private String additionalProjects;
8686
private Boolean filterTablesOnDefaultDataset;
87+
private Boolean enableProjectDiscovery;
8788
private Integer requestGoogleDriveScope;
8889
private Integer metadataFetchThreadCount;
8990
private String sslTrustStorePath;
@@ -242,6 +243,12 @@ public class DataSource implements javax.sql.DataSource {
242243
BigQueryJdbcUrlUtility.convertIntToBoolean(
243244
val,
244245
BigQueryJdbcUrlUtility.FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME)))
246+
.put(
247+
BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME,
248+
(ds, val) ->
249+
ds.setEnableProjectDiscovery(
250+
BigQueryJdbcUrlUtility.convertIntToBoolean(
251+
val, BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME)))
245252
.put(
246253
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME,
247254
(ds, val) -> ds.setRequestGoogleDriveScope(Integer.parseInt(val)))
@@ -555,6 +562,11 @@ Properties createProperties() {
555562
BigQueryJdbcUrlUtility.FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME,
556563
String.valueOf(this.filterTablesOnDefaultDataset));
557564
}
565+
if (this.enableProjectDiscovery != null) {
566+
connectionProperties.setProperty(
567+
BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME,
568+
String.valueOf(this.enableProjectDiscovery));
569+
}
558570
if (this.requestGoogleDriveScope != null) {
559571
connectionProperties.setProperty(
560572
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME,
@@ -1060,6 +1072,16 @@ public void setFilterTablesOnDefaultDataset(Boolean filterTablesOnDefaultDataset
10601072
this.filterTablesOnDefaultDataset = filterTablesOnDefaultDataset;
10611073
}
10621074

1075+
public Boolean getEnableProjectDiscovery() {
1076+
return enableProjectDiscovery != null
1077+
? enableProjectDiscovery
1078+
: BigQueryJdbcUrlUtility.DEFAULT_ENABLE_PROJECT_DISCOVERY_VALUE;
1079+
}
1080+
1081+
public void setEnableProjectDiscovery(Boolean enableProjectDiscovery) {
1082+
this.enableProjectDiscovery = enableProjectDiscovery;
1083+
}
1084+
10631085
public Integer getRequestGoogleDriveScope() {
10641086
return requestGoogleDriveScope != null
10651087
? requestGoogleDriveScope

‎java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryConnectionTest.java‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,18 +23,27 @@
2323
import static org.junit.jupiter.api.Assertions.assertSame;
2424
import static org.junit.jupiter.api.Assertions.assertThrows;
2525
import static org.junit.jupiter.api.Assertions.assertTrue;
26+
import static org.mockito.Mockito.mock;
27+
import static org.mockito.Mockito.times;
28+
import static org.mockito.Mockito.verify;
29+
import static org.mockito.Mockito.when;
2630

2731
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
32+
import com.google.api.gax.paging.Page;
2833
import com.google.api.gax.rpc.HeaderProvider;
2934
import com.google.api.gax.rpc.TransportChannelProvider;
3035
import com.google.cloud.bigquery.BigQuery;
36+
import com.google.cloud.bigquery.BigQueryException;
37+
import com.google.cloud.bigquery.Project;
3138
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
3239
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
3340
import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
3441
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
3542
import java.io.IOException;
3643
import java.io.InputStream;
3744
import java.sql.SQLException;
45+
import java.util.Arrays;
46+
import java.util.List;
3847
import java.util.Optional;
3948
import java.util.Properties;
4049
import java.util.logging.Level;
@@ -519,4 +528,79 @@ public void testWrapperMethods() throws Exception {
519528
assertTrue(e.getMessage().contains("Cannot unwrap to java.sql.Statement"));
520529
}
521530
}
531+
532+
@Test
533+
public void testGetDiscoveredProjects_Success() throws Exception {
534+
try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
535+
BigQuery mockBigQuery = mock(BigQuery.class);
536+
connection.bigQuery = mockBigQuery;
537+
538+
Page<Project> mockPage = mock(Page.class);
539+
Project project1 = mock(Project.class);
540+
when(project1.getProjectId()).thenReturn("discovered-p1");
541+
Project project2 = mock(Project.class);
542+
when(project2.getProjectId()).thenReturn("discovered-p2");
543+
544+
when(mockPage.iterateAll()).thenReturn(Arrays.asList(project1, project2));
545+
when(mockBigQuery.listProjects()).thenReturn(mockPage);
546+
547+
List<String> discovered = connection.getDiscoveredProjects();
548+
assertEquals(Arrays.asList("discovered-p1", "discovered-p2"), discovered);
549+
550+
// Verify caching: second call should not invoke listProjects again
551+
List<String> discoveredCached = connection.getDiscoveredProjects();
552+
assertSame(discovered, discoveredCached);
553+
verify(mockBigQuery, times(1)).listProjects();
554+
}
555+
}
556+
557+
@Test
558+
public void testGetDiscoveredProjects_BigQueryExceptionThrown() throws Exception {
559+
try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
560+
BigQuery mockBigQuery = mock(BigQuery.class);
561+
connection.bigQuery = mockBigQuery;
562+
563+
BigQueryException exception = new BigQueryException(403, "Access Denied");
564+
when(mockBigQuery.listProjects()).thenThrow(exception);
565+
566+
// Verify that it throws BigQueryJdbcException
567+
BigQueryJdbcException ex =
568+
assertThrows(
569+
BigQueryJdbcException.class,
570+
() -> {
571+
connection.getDiscoveredProjects();
572+
});
573+
assertTrue(ex.getMessage().contains("Failed to list all accessible projects."));
574+
assertEquals(exception, ex.getCause());
575+
576+
// Subsequent call should retry since no cache is set
577+
assertThrows(
578+
BigQueryJdbcException.class,
579+
() -> {
580+
connection.getDiscoveredProjects();
581+
});
582+
verify(mockBigQuery, times(2)).listProjects();
583+
}
584+
}
585+
586+
@Test
587+
public void testGetDiscoveredProjects_OtherExceptionThrown() throws Exception {
588+
try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
589+
BigQuery mockBigQuery = mock(BigQuery.class);
590+
connection.bigQuery = mockBigQuery;
591+
592+
RuntimeException exception = new RuntimeException("Generic Network Failure");
593+
when(mockBigQuery.listProjects()).thenThrow(exception);
594+
595+
// Verify that it throws BigQueryJdbcException
596+
BigQueryJdbcException ex =
597+
assertThrows(
598+
BigQueryJdbcException.class,
599+
() -> {
600+
connection.getDiscoveredProjects();
601+
});
602+
assertTrue(ex.getMessage().contains("Failed to list all accessible projects."));
603+
assertEquals(exception, ex.getCause());
604+
}
605+
}
522606
}

0 commit comments

Comments
 (0)