Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.api.gax.paging.Page;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.FixedHeaderProvider;
import com.google.api.gax.rpc.HeaderProvider;
Expand All @@ -31,6 +32,7 @@
import com.google.cloud.bigquery.DatasetId;
import com.google.cloud.bigquery.Job;
import com.google.cloud.bigquery.JobInfo;
import com.google.cloud.bigquery.Project;
import com.google.cloud.bigquery.QueryJobConfiguration;
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
Expand All @@ -41,6 +43,7 @@
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
import com.google.cloud.bigquery.storage.v1.BigQueryWriteSettings;
import com.google.cloud.http.HttpTransportOptions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSortedSet;
import java.io.IOException;
import java.io.InputStream;
Expand Down Expand Up @@ -121,6 +124,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
BigQueryJdbcUrlUtility.SWA_APPEND_ROW_COUNT_PROPERTY_NAME,
BigQueryJdbcUrlUtility.SWA_ACTIVATION_ROW_COUNT_PROPERTY_NAME,
BigQueryJdbcUrlUtility.FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME,
BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME,
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME,
BigQueryJdbcUrlUtility.SSL_TRUST_STORE_PROPERTY_NAME,
BigQueryJdbcUrlUtility.MAX_BYTES_BILLED_PROPERTY_NAME,
Expand Down Expand Up @@ -170,6 +174,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
int highThroughputMinTableSize;
int highThroughputActivationRatio;
boolean enableSession;
boolean enableProjectDiscovery;
private List<String> discoveredProjectsCache;
boolean unsupportedHTAPIFallback;
boolean useQueryCache;
String queryDialect;
Expand Down Expand Up @@ -338,6 +344,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
this.additionalProjects = ds.getAdditionalProjects();

this.filterTablesOnDefaultDataset = ds.getFilterTablesOnDefaultDataset();
this.enableProjectDiscovery = ds.getEnableProjectDiscovery();
this.requestGoogleDriveScope = ds.getRequestGoogleDriveScope();
this.metadataFetchThreadCount = ds.getMetadataFetchThreadCount();
this.requestReason = ds.getRequestReason();
Expand Down Expand Up @@ -1312,6 +1319,32 @@ private boolean checkIsReadOnlyTokenUsed(Map<String, String> authProps) {
return false;
}

public boolean isEnableProjectDiscovery() {
return this.enableProjectDiscovery;
}

public synchronized List<String> getDiscoveredProjects() throws SQLException {
if (this.discoveredProjectsCache != null) {
return this.discoveredProjectsCache;
}

try {
BigQuery bigQuery = getBigQuery();
List<String> projects = new ArrayList<>();
Page<Project> projectPage = bigQuery.listProjects();
for (Project project : projectPage.iterateAll()) {
projects.add(project.getProjectId());
}
this.discoveredProjectsCache = ImmutableList.copyOf(projects);
} catch (BigQueryException e) {
throw new BigQueryJdbcException(
"Failed to list all accessible projects due to BigQuery error.", e);
} catch (Exception e) {
throw new BigQueryJdbcException("Failed to list all accessible projects.", e);
}
Comment thread
keshavdandeva marked this conversation as resolved.
Comment thread
keshavdandeva marked this conversation as resolved.
return this.discoveredProjectsCache;
}

@Override
public <T> T unwrap(Class<T> iface) throws SQLException {
if (iface.isInstance(this)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1996,14 +1996,14 @@ Comparator<FieldValueList> defineGetTablesComparator(FieldList resultSchemaField
}

@Override
public ResultSet getSchemas() {
public ResultSet getSchemas() throws SQLException {
LOG.info("getSchemas() called");

return getSchemas(null, null);
}

@Override
public ResultSet getCatalogs() {
public ResultSet getCatalogs() throws SQLException {
LOG.info("getCatalogs() called");

final List<String> accessibleCatalogs = getAccessibleCatalogNames();
Expand Down Expand Up @@ -3618,7 +3618,7 @@ public RowIdLifetime getRowIdLifetime() {
}

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

if (catalogParam != null) {
projectsToScanList.add(catalogParam);
} else {
projectsToScanList.addAll(getAccessibleCatalogNames());
}
try {
if (catalogParam != null) {
projectsToScanList.add(catalogParam);
} else {
projectsToScanList.addAll(getAccessibleCatalogNames());
}

if (projectsToScanList.isEmpty()) {
LOG.info(
"No valid projects to scan (primary, specified, or additional). Returning empty"
+ " resultset.");
return;
}
if (projectsToScanList.isEmpty()) {
LOG.info(
"No valid projects to scan (primary, specified, or additional). Returning empty"
+ " resultset.");
return;
}

try {
for (String currentProjectToScan : projectsToScanList) {
if (Thread.currentThread().isInterrupted()) {
LOG.warning(
Expand Down Expand Up @@ -3707,6 +3707,13 @@ public ResultSet getSchemas(String catalog, String schemaPattern) {

} catch (Throwable t) {
LOG.severe("Unexpected error in schema fetcher runnable: " + t.getMessage());
Exception ex = (t instanceof Exception) ? (Exception) t : new Exception(t);
try {
queue.put(BigQueryFieldValueListWrapper.ofError(ex));
} catch (InterruptedException ie) {
LOG.warning("Failed to put exception to queue due to interruption.");
Thread.currentThread().interrupt();
}
} finally {
signalEndOfData(queue, localResultSchemaFields);
LOG.info("Schema fetcher thread finished.");
Expand Down Expand Up @@ -5178,7 +5185,7 @@ private String getCurrentCatalogName() {
return this.connection.getCatalog();
}

private List<String> getAccessibleCatalogNames() {
private List<String> getAccessibleCatalogNames() throws SQLException {
Set<String> accessibleCatalogs = new HashSet<>();
String primaryCatalog = getCurrentCatalogName();
if (primaryCatalog != null && !primaryCatalog.isEmpty()) {
Expand All @@ -5199,6 +5206,10 @@ private List<String> getAccessibleCatalogNames() {
}
}

if (this.connection.isEnableProjectDiscovery()) {
accessibleCatalogs.addAll(this.connection.getDiscoveredProjects());
}

List<String> sortedCatalogs = new ArrayList<>(accessibleCatalogs);
Collections.sort(sortedCatalogs);
return sortedCatalogs;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,8 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
static final String FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME =
"FilterTablesOnDefaultDataset";
static final boolean DEFAULT_FILTER_TABLES_ON_DEFAULT_DATASET_VALUE = false;
static final String ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME = "EnableProjectDiscovery";
static final boolean DEFAULT_ENABLE_PROJECT_DISCOVERY_VALUE = false;
static final String REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME = "RequestGoogleDriveScope";
static final String SSL_TRUST_STORE_PROPERTY_NAME = "SSLTrustStore";
static final String SSL_TRUST_STORE_PWD_PROPERTY_NAME = "SSLTrustStorePwd";
Expand Down Expand Up @@ -577,6 +579,13 @@ protected boolean removeEldestEntry(Map.Entry<String, Map<String, String>> eldes
.setDefaultValue(
String.valueOf(DEFAULT_FILTER_TABLES_ON_DEFAULT_DATASET_VALUE))
.build(),
BigQueryConnectionProperty.newBuilder()
.setName(ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME)
.setDescription(
"Enables or disables automatic discovery of all accessible Google Cloud projects. "
+ "When disabled, only the default ProjectId and AdditionalProjects are listed as catalogs.")
.setDefaultValue(String.valueOf(DEFAULT_ENABLE_PROJECT_DISCOVERY_VALUE))
.build(),
BigQueryConnectionProperty.newBuilder()
.setName(REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME)
.setDescription(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ public class DataSource implements javax.sql.DataSource {
private Boolean enableWriteAPI;
private String additionalProjects;
private Boolean filterTablesOnDefaultDataset;
private Boolean enableProjectDiscovery;
private Integer requestGoogleDriveScope;
private Integer metadataFetchThreadCount;
private String sslTrustStorePath;
Expand Down Expand Up @@ -242,6 +243,12 @@ public class DataSource implements javax.sql.DataSource {
BigQueryJdbcUrlUtility.convertIntToBoolean(
val,
BigQueryJdbcUrlUtility.FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME)))
.put(
BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME,
(ds, val) ->
ds.setEnableProjectDiscovery(
BigQueryJdbcUrlUtility.convertIntToBoolean(
val, BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME)))
.put(
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME,
(ds, val) -> ds.setRequestGoogleDriveScope(Integer.parseInt(val)))
Expand Down Expand Up @@ -555,6 +562,11 @@ Properties createProperties() {
BigQueryJdbcUrlUtility.FILTER_TABLES_ON_DEFAULT_DATASET_PROPERTY_NAME,
String.valueOf(this.filterTablesOnDefaultDataset));
}
if (this.enableProjectDiscovery != null) {
connectionProperties.setProperty(
BigQueryJdbcUrlUtility.ENABLE_PROJECT_DISCOVERY_PROPERTY_NAME,
String.valueOf(this.enableProjectDiscovery));
}
if (this.requestGoogleDriveScope != null) {
connectionProperties.setProperty(
BigQueryJdbcUrlUtility.REQUEST_GOOGLE_DRIVE_SCOPE_PROPERTY_NAME,
Expand Down Expand Up @@ -1060,6 +1072,16 @@ public void setFilterTablesOnDefaultDataset(Boolean filterTablesOnDefaultDataset
this.filterTablesOnDefaultDataset = filterTablesOnDefaultDataset;
}

public Boolean getEnableProjectDiscovery() {
return enableProjectDiscovery != null
? enableProjectDiscovery
: BigQueryJdbcUrlUtility.DEFAULT_ENABLE_PROJECT_DISCOVERY_VALUE;
}

public void setEnableProjectDiscovery(Boolean enableProjectDiscovery) {
this.enableProjectDiscovery = enableProjectDiscovery;
}

public Integer getRequestGoogleDriveScope() {
return requestGoogleDriveScope != null
? requestGoogleDriveScope
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,18 +23,27 @@
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
import com.google.api.gax.paging.Page;
import com.google.api.gax.rpc.HeaderProvider;
import com.google.api.gax.rpc.TransportChannelProvider;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQueryException;
import com.google.cloud.bigquery.Project;
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.Properties;
import java.util.logging.Level;
Expand Down Expand Up @@ -519,4 +528,81 @@ public void testWrapperMethods() throws Exception {
assertTrue(e.getMessage().contains("Cannot unwrap to java.sql.Statement"));
}
}

@Test
public void testGetDiscoveredProjects_Success() throws Exception {
try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
BigQuery mockBigQuery = mock(BigQuery.class);
connection.bigQuery = mockBigQuery;

Page<Project> mockPage = mock(Page.class);
Project project1 = mock(Project.class);
when(project1.getProjectId()).thenReturn("discovered-p1");
Project project2 = mock(Project.class);
when(project2.getProjectId()).thenReturn("discovered-p2");

when(mockPage.iterateAll()).thenReturn(Arrays.asList(project1, project2));
when(mockBigQuery.listProjects()).thenReturn(mockPage);

List<String> discovered = connection.getDiscoveredProjects();
assertEquals(Arrays.asList("discovered-p1", "discovered-p2"), discovered);

// Verify caching: second call should not invoke listProjects again
List<String> discoveredCached = connection.getDiscoveredProjects();
assertSame(discovered, discoveredCached);
verify(mockBigQuery, times(1)).listProjects();
}
}

@Test
public void testGetDiscoveredProjects_BigQueryExceptionThrown() throws Exception {
try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
BigQuery mockBigQuery = mock(BigQuery.class);
connection.bigQuery = mockBigQuery;

BigQueryException exception = new BigQueryException(403, "Access Denied");
when(mockBigQuery.listProjects()).thenThrow(exception);

// Verify that it throws BigQueryJdbcException
BigQueryJdbcException ex =
assertThrows(
BigQueryJdbcException.class,
() -> {
connection.getDiscoveredProjects();
});
assertTrue(
ex.getMessage()
.contains("Failed to list all accessible projects due to BigQuery error."));
assertEquals(exception, ex.getCause());

// Subsequent call should retry since no cache is set
assertThrows(
BigQueryJdbcException.class,
() -> {
connection.getDiscoveredProjects();
});
verify(mockBigQuery, times(2)).listProjects();
}
}

Comment thread
keshavdandeva marked this conversation as resolved.
@Test
public void testGetDiscoveredProjects_OtherExceptionThrown() throws Exception {
try (BigQueryConnection connection = new BigQueryConnection(BASE_URL)) {
BigQuery mockBigQuery = mock(BigQuery.class);
connection.bigQuery = mockBigQuery;

RuntimeException exception = new RuntimeException("Generic Network Failure");
when(mockBigQuery.listProjects()).thenThrow(exception);

// Verify that it throws BigQueryJdbcException
BigQueryJdbcException ex =
assertThrows(
BigQueryJdbcException.class,
() -> {
connection.getDiscoveredProjects();
});
assertTrue(ex.getMessage().contains("Failed to list all accessible projects."));
assertEquals(exception, ex.getCause());
}
}
}
Loading
Loading