diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 2f5863054903..d394b1ba7a7e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -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(); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java index fec91467c430..e8a4c83b9f4e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaData.java @@ -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; @@ -3922,8 +3923,7 @@ public boolean isWrapperFor(Class iface) throws SQLException { * for querying BigQuery's metadata. * @see BigQueryConnection#isFilterTablesOnDefaultDataset() */ - private Tuple determineEffectiveCatalogAndSchema( - String catalog, String schemaPattern) { + Tuple determineEffectiveCatalogAndSchema(String catalog, String schemaPattern) { String effectiveCatalog = catalog; String effectiveSchemaPattern = schemaPattern; @@ -3932,7 +3932,10 @@ private Tuple 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(); @@ -4403,20 +4406,23 @@ private List fetchMatchingDatasets( return allDatasets; } - private List getAccessibleCatalogNames() throws SQLException { + List getAccessibleCatalogNames() throws SQLException { Set 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 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); diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java index 928ed785a795..d32867b43524 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java @@ -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; @@ -924,4 +925,44 @@ static Map 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) { + return splitQualifiedDataset(trimmed, colonIdx, ':'); + } + + int dotIdx = trimmed.indexOf('.'); + if (dotIdx >= 0) { + return splitQualifiedDataset(trimmed, dotIdx, '.'); + } + + return DatasetId.of(trimmed); + } + + private static DatasetId splitQualifiedDataset(String trimmed, int delimiterIdx, char delimiter) { + if (trimmed.indexOf(delimiter, delimiterIdx + 1) >= 0) { + throw invalidDefaultDatasetException(); + } + String project = trimmed.substring(0, delimiterIdx).trim(); + String dataset = trimmed.substring(delimiterIdx + 1).trim(); + if (project.isEmpty() || dataset.isEmpty()) { + throw invalidDefaultDatasetException(); + } + return DatasetId.of(project, dataset); + } + + private static IllegalArgumentException invalidDefaultDatasetException() { + IllegalArgumentException ex = + new IllegalArgumentException( + "DefaultDataset format is invalid. Supported options are datasetId," + + " projectId:datasetId, or projectId.datasetId"); + LOG.severe(ex.getMessage(), ex); + return ex; + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java index 5b58b80f2765..a059efa31c0a 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryDatabaseMetaDataTest.java @@ -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; @@ -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 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 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 catalogs = dbMetadata.getAccessibleCatalogNames(); + + assertEquals( + Arrays.asList("extra-proj1", "extra-proj2", "pcnt-warehouse", "primary-project"), catalogs); + } } diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java index 0bc580391b12..676d603f89e0 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java @@ -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; @@ -279,4 +282,84 @@ 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_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_dotDelimited() { + 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_twoTierPcntNamespace() { + DatasetId datasetId = BigQueryJdbcUrlUtility.parseDefaultDataset("my-warehouse.my_namespace"); + assertNotNull(datasetId); + assertEquals("my_namespace", datasetId.getDataset()); + assertEquals("my-warehouse", 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() { + assertThrows( + IllegalArgumentException.class, + () -> BigQueryJdbcUrlUtility.parseDefaultDataset("my-project:")); + assertThrows( + IllegalArgumentException.class, + () -> BigQueryJdbcUrlUtility.parseDefaultDataset(":my_dataset")); + assertThrows( + IllegalArgumentException.class, + () -> BigQueryJdbcUrlUtility.parseDefaultDataset("tpc:my-project:")); + assertThrows( + IllegalArgumentException.class, + () -> BigQueryJdbcUrlUtility.parseDefaultDataset("my-project.")); + assertThrows( + IllegalArgumentException.class, + () -> BigQueryJdbcUrlUtility.parseDefaultDataset(".my_dataset")); + assertThrows( + IllegalArgumentException.class, + () -> BigQueryJdbcUrlUtility.parseDefaultDataset("part1.part2.part3")); + } }