Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -334,24 +334,7 @@ public class BigQueryConnection extends BigQueryNoOpsConnection {
this.reqGoogleDriveScope,
httpTransportFactory,
this.connectionClassName);
String defaultDatasetString = ds.getDefaultDataset();
if (defaultDatasetString == null || defaultDatasetString.trim().isEmpty()) {
this.defaultDataset = null;
} else {
String[] parts = defaultDatasetString.split("\\.");
if (parts.length == 2) {
this.defaultDataset = DatasetId.of(parts[0], parts[1]);
} else if (parts.length == 1) {
this.defaultDataset = DatasetId.of(parts[0]);
} else {
IllegalArgumentException ex =
new IllegalArgumentException(
"DefaultDataset format is invalid. Supported options are datasetId or"
+ " projectId.datasetId");
LOG.severe(ex.getMessage(), ex);
throw ex;
}
}
this.defaultDataset = BigQueryJdbcUrlUtility.parseDefaultDataset(ds.getDefaultDataset());
this.location = ds.getLocation();
this.enableHighThroughputAPI = ds.getEnableHighThroughputAPI();
this.highThroughputMinTableSize = ds.getHighThroughputMinTableSize();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
import com.google.cloud.bigquery.jdbc.BigQueryTypeRegistry.ColumnTypeInfo;
import com.google.cloud.bigquery.jdbc.utils.BigQueryJdbcVersionUtility;
import com.google.common.base.Splitter;
import io.opentelemetry.context.Context;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
Expand Down Expand Up @@ -3922,8 +3923,7 @@ public boolean isWrapperFor(Class<?> iface) throws SQLException {
* for querying BigQuery's metadata.
* @see BigQueryConnection#isFilterTablesOnDefaultDataset()
*/
private Tuple<String, String> determineEffectiveCatalogAndSchema(
String catalog, String schemaPattern) {
Tuple<String, String> determineEffectiveCatalogAndSchema(String catalog, String schemaPattern) {
String effectiveCatalog = catalog;
String effectiveSchemaPattern = schemaPattern;

Expand All @@ -3932,7 +3932,10 @@ private Tuple<String, String> determineEffectiveCatalogAndSchema(
&& this.connection.getDefaultDataset().getDataset() != null
&& !this.connection.getDefaultDataset().getDataset().isEmpty()) {

String defaultProjectFromConnection = this.connection.getCatalog();
String defaultProjectFromConnection =
(this.connection.getDefaultDataset().getProject() != null)
? this.connection.getDefaultDataset().getProject()
: this.connection.getCatalog();
// We only use the dataset part of the DefaultDataset for schema filtering
String defaultSchemaFromConnection = this.connection.getDefaultDataset().getDataset();

Expand Down Expand Up @@ -4403,20 +4406,23 @@ private List<Dataset> fetchMatchingDatasets(
return allDatasets;
}

private List<String> getAccessibleCatalogNames() throws SQLException {
List<String> getAccessibleCatalogNames() throws SQLException {
Set<String> accessibleCatalogs = new HashSet<>();
String primaryCatalog = this.connection.getCatalog();
if (primaryCatalog != null && !primaryCatalog.isEmpty()) {
accessibleCatalogs.add(primaryCatalog);
}

if (this.connection.getDefaultDataset() != null
&& this.connection.getDefaultDataset().getProject() != null
&& !this.connection.getDefaultDataset().getProject().isEmpty()) {
accessibleCatalogs.add(this.connection.getDefaultDataset().getProject());
}

String additionalProjectsStr = this.connection.getAdditionalProjects();
if (additionalProjectsStr != null && !additionalProjectsStr.trim().isEmpty()) {
List<String> additionalProjects =
com.google.common.base.Splitter.on(',')
.trimResults()
.omitEmptyStrings()
.splitToList(additionalProjectsStr);
Splitter.on(',').trimResults().omitEmptyStrings().splitToList(additionalProjectsStr);
for (String project : additionalProjects) {
if (project != null && !project.isEmpty()) {
accessibleCatalogs.add(project);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import com.google.api.client.util.escape.CharEscapers;
import com.google.cloud.bigquery.BigQueryOptions;
import com.google.cloud.bigquery.DatasetId;
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -924,4 +925,25 @@ static Map<String, String> parsePropertiesMapFromValue(
}
return propertiesMap;
}

static DatasetId parseDefaultDataset(String defaultDataset) {
if (defaultDataset == null || defaultDataset.trim().isEmpty()) {
return null;
}

String trimmed = defaultDataset.trim();
int colonIdx = trimmed.lastIndexOf(':');
if (colonIdx >= 0) {
String project = trimmed.substring(0, colonIdx).trim();
String dataset = trimmed.substring(colonIdx + 1).trim();
if (project.isEmpty() || dataset.isEmpty()) {
throw new BigQueryJdbcRuntimeException(
"DefaultDataset format is invalid. Supported options are datasetId, "
+ "catalog.namespace, projectId:datasetId, or projectId:catalog.namespace");
}
return DatasetId.of(project, dataset);
}

return DatasetId.of(trimmed);
}
Comment thread
keshavdandeva marked this conversation as resolved.
}
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import static org.mockito.Mockito.*;

import com.google.api.gax.paging.Page;
import com.google.cloud.Tuple;
import com.google.cloud.bigquery.*;
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
import com.google.cloud.bigquery.jdbc.BigQueryTypeRegistry.ColumnTypeInfo;
Expand Down Expand Up @@ -3697,4 +3698,46 @@ public void testGetCrossReference_noKeys() throws SQLException {
assertFalse(rs.next());
}
}

@Test
public void testDetermineEffectiveCatalogAndSchema_withDefaultDatasetProject() {
when(bigQueryConnection.isFilterTablesOnDefaultDataset()).thenReturn(true);
when(bigQueryConnection.getCatalog()).thenReturn("primary-project");
when(bigQueryConnection.getDefaultDataset())
.thenReturn(DatasetId.of("custom-project", "warehouse.namespace"));

dbMetadata = new BigQueryDatabaseMetaData(bigQueryConnection);
Tuple<String, String> result = dbMetadata.determineEffectiveCatalogAndSchema(null, null);

assertEquals("custom-project", result.x());
assertEquals("warehouse.namespace", result.y());
}

@Test
public void testDetermineEffectiveCatalogAndSchema_withDefaultDatasetNoProject() {
when(bigQueryConnection.isFilterTablesOnDefaultDataset()).thenReturn(true);
when(bigQueryConnection.getCatalog()).thenReturn("primary-project");
when(bigQueryConnection.getDefaultDataset()).thenReturn(DatasetId.of("my_dataset"));

dbMetadata = new BigQueryDatabaseMetaData(bigQueryConnection);
Tuple<String, String> result = dbMetadata.determineEffectiveCatalogAndSchema(null, null);

assertEquals("primary-project", result.x());
assertEquals("my_dataset", result.y());
}

@Test
public void testGetAccessibleCatalogNames_includesDefaultDatasetProject() throws SQLException {
when(bigQueryConnection.getCatalog()).thenReturn("primary-project");
when(bigQueryConnection.getDefaultDataset())
.thenReturn(DatasetId.of("pcnt-warehouse", "pcnt_ns"));
when(bigQueryConnection.getAdditionalProjects()).thenReturn("extra-proj1,extra-proj2");
when(bigQueryConnection.isEnableProjectDiscovery()).thenReturn(false);

dbMetadata = new BigQueryDatabaseMetaData(bigQueryConnection);
List<String> catalogs = dbMetadata.getAccessibleCatalogNames();

assertEquals(
Arrays.asList("extra-proj1", "extra-proj2", "pcnt-warehouse", "primary-project"), catalogs);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,13 @@

import static com.google.common.truth.Truth.assertThat;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;

import com.google.cloud.bigquery.DatasetId;
import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException;
import java.util.Collections;
import java.util.Map;
Expand Down Expand Up @@ -279,4 +282,70 @@ public void testParseEnableProjectDiscovery() {
String result2 = BigQueryJdbcUrlUtility.parseUriProperty(url2, "EnableProjectDiscovery");
assertThat(result2).isEqualTo("false");
}

@Test
public void testParseDefaultDataset_singleIdentifier() {
DatasetId datasetId = BigQueryJdbcUrlUtility.parseDefaultDataset("my_dataset");
assertNotNull(datasetId);
assertEquals("my_dataset", datasetId.getDataset());
assertNull(datasetId.getProject());
}

@Test
public void testParseDefaultDataset_twoTierPcntNamespace() {
DatasetId datasetId = BigQueryJdbcUrlUtility.parseDefaultDataset("my-warehouse.my_namespace");
assertNotNull(datasetId);
assertEquals("my-warehouse.my_namespace", datasetId.getDataset());
assertNull(datasetId.getProject());
}

@Test
public void testParseDefaultDataset_colonDelimited() {
DatasetId datasetId = BigQueryJdbcUrlUtility.parseDefaultDataset("my-project:my_dataset");
assertNotNull(datasetId);
assertEquals("my_dataset", datasetId.getDataset());
assertEquals("my-project", datasetId.getProject());
}

@Test
public void testParseDefaultDataset_threeTierPcntNamespace() {
DatasetId datasetId =
BigQueryJdbcUrlUtility.parseDefaultDataset("my-project:my-warehouse.my_namespace");
assertNotNull(datasetId);
assertEquals("my-warehouse.my_namespace", datasetId.getDataset());
assertEquals("my-project", datasetId.getProject());
}

@Test
public void testParseDefaultDataset_tpcProject() {
DatasetId datasetId = BigQueryJdbcUrlUtility.parseDefaultDataset("tpc:my-project:my_dataset");
assertNotNull(datasetId);
assertEquals("my_dataset", datasetId.getDataset());
assertEquals("tpc:my-project", datasetId.getProject());
}

@Test
public void testParseDefaultDataset_tpcProject_threeTierPcntNamespace() {
DatasetId datasetId =
BigQueryJdbcUrlUtility.parseDefaultDataset("tpc:my-project:my-warehouse.my_namespace");
assertNotNull(datasetId);
assertEquals("my-warehouse.my_namespace", datasetId.getDataset());
assertEquals("tpc:my-project", datasetId.getProject());
}

@Test
public void testParseDefaultDataset_invalidFormats() {
Comment thread
keshavdandeva marked this conversation as resolved.
assertThrows(
BigQueryJdbcRuntimeException.class,
() -> BigQueryJdbcUrlUtility.parseDefaultDataset("my-project:"));
assertThrows(
BigQueryJdbcRuntimeException.class,
() -> BigQueryJdbcUrlUtility.parseDefaultDataset(":my_dataset"));
assertThrows(
BigQueryJdbcRuntimeException.class,
() -> BigQueryJdbcUrlUtility.parseDefaultDataset("tpc:my-project:"));
assertThrows(
BigQueryJdbcRuntimeException.class,
() -> BigQueryJdbcUrlUtility.parseDefaultDataset(" : "));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -344,7 +344,7 @@ public void testDefaultDataset() throws SQLException {
@Test
public void testDefaultDatasetWithProject() throws SQLException {
String connection_uri =
ITBigQueryJDBCTest.connection_uri + "DEFAULTDATASET=" + PROJECT_ID + ".testDataset";
ITBigQueryJDBCTest.connection_uri + "DEFAULTDATASET=" + PROJECT_ID + ":testDataset";

Driver driver = BigQueryDriver.getRegisteredDriver();
assertTrue(driver.acceptsURL(connection_uri));
Expand Down
Loading