Repository navigation
feat(bigquery-jdbc): add graph metadata integration for getTables and getTableTypes #14868
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: jdbc/graph-metadata-integration
Are you sure you want to change the base?
Changes from all commits
7a79749
c430c6b
a285b57
c8c8467
b30489f
9e062da
410f17f
577c492
6401731
b08aa2b
2eb99b7
fd85ea0
a831f43
ff49045
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -70,6 +70,7 @@ | |
| import java.util.Comparator; | ||
| import java.util.HashSet; | ||
| import java.util.List; | ||
| import java.util.Locale; | ||
| import java.util.Set; | ||
| import java.util.concurrent.BlockingQueue; | ||
| import java.util.concurrent.Callable; | ||
|
|
@@ -103,6 +104,12 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData { | |
| private static final int DEFAULT_QUEUE_CAPACITY = 5000; | ||
| private static final String GET_EXPORTED_KEYS_SQL = "DatabaseMetaData_GetExportedKeys.sql"; | ||
| private static String exportedKeysSqlContent; | ||
| private static final String GET_GRAPH_TABLES_REGION_SQL = | ||
| "DatabaseMetaData_GetGraphTables_Region.sql"; | ||
| private static String graphTablesRegionSqlContent; | ||
| private static final String GET_GRAPH_TABLES_DATASET_SQL = | ||
| "DatabaseMetaData_GetGraphTables_Dataset.sql"; | ||
| private static String graphTablesDatasetSqlContent; | ||
| // Declared package-private for testing. | ||
| static final String GOOGLE_SQL_QUOTED_IDENTIFIER = "`"; | ||
| // Does not include SQL:2003 Keywords as per JDBC spec. | ||
|
|
@@ -138,6 +145,15 @@ class BigQueryDatabaseMetaData implements DatabaseMetaData { | |
| static final String GOOGLE_SQL_CATALOG_SEPARATOR = "."; | ||
| static final int GOOGLE_SQL_MAX_COL_NAME_LEN = 300; | ||
| static final int GOOGLE_SQL_MAX_COLS_PER_TABLE = 10000; | ||
| private static final String[] DEFAULT_TABLE_TYPES = { | ||
| "EXTERNAL", "MATERIALIZED VIEW", "SNAPSHOT", "TABLE", "VIEW" | ||
| }; | ||
| private static final String[] TABLE_TYPES_WITH_GRAPH = { | ||
| "EXTERNAL", "GRAPH TABLE", "MATERIALIZED VIEW", "SNAPSHOT", "TABLE", "VIEW" | ||
| }; | ||
| static final String GRAPH_TABLE_TYPE = "GRAPH TABLE"; | ||
| private static final String DEFAULT_GRAPH_REGION = "us"; | ||
| private static final Pattern VALID_IDENTIFIER_PATTERN = Pattern.compile("^[A-Za-z0-9_\\-:.]+$"); | ||
|
|
||
| String URL; | ||
| BigQueryConnection connection; | ||
|
|
@@ -1308,19 +1324,37 @@ private ResultSet getTablesImpl( | |
| () -> { | ||
| final List<FieldValueList> collectedResults = | ||
| Collections.synchronizedList(new ArrayList<>()); | ||
| processTargetTablesConcurrently( | ||
| targetDatasets, | ||
| boolean onlyGraphTablesRequested = | ||
| requestedTypes != null | ||
| && requestedTypes.size() == 1 | ||
| && requestedTypes.contains(GRAPH_TABLE_TYPE); | ||
| boolean isExactGraphTableName = | ||
| tableNamePattern != null | ||
| && !tableNamePattern.contains("%") | ||
| && tableNamePattern.contains("."); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please add a comment explaining this assumption that if table name contains |
||
| if (!isExactGraphTableName && !onlyGraphTablesRequested) { | ||
| processTargetTablesConcurrently( | ||
| targetDatasets, | ||
| tableNamePattern, | ||
| true, | ||
| false, | ||
| false, | ||
| collectedResults, | ||
| resultSchemaFields, | ||
| (bqTable, results, fields) -> { | ||
| processTableInfo(bqTable, requestedTypes, results, fields); | ||
| }, | ||
| TableOption.fields( | ||
| TableField.TABLE_REFERENCE, TableField.TYPE, TableField.DESCRIPTION)); | ||
| } | ||
|
|
||
| appendGraphTableRows( | ||
| effectiveCatalog, | ||
| effectiveSchemaPattern, | ||
| tableNamePattern, | ||
| true, | ||
| false, | ||
| false, | ||
| requestedTypes, | ||
| collectedResults, | ||
| resultSchemaFields, | ||
| (bqTable, results, fields) -> { | ||
| processTableInfo(bqTable, requestedTypes, results, fields); | ||
| }, | ||
| TableOption.fields( | ||
| TableField.TABLE_REFERENCE, TableField.TYPE, TableField.DESCRIPTION)); | ||
| resultSchemaFields); | ||
|
|
||
| Comparator<FieldValueList> comparator = defineGetTablesComparator(resultSchemaFields); | ||
| sortResults(collectedResults, comparator, "getTables", LOG); | ||
|
|
@@ -1441,6 +1475,120 @@ Comparator<FieldValueList> defineGetTablesComparator(FieldList resultSchemaField | |
| Comparator.nullsFirst(String::compareTo)); | ||
| } | ||
|
|
||
| static String quoteIdentifier(String identifier) throws SQLException { | ||
| if (identifier == null || !VALID_IDENTIFIER_PATTERN.matcher(identifier).matches()) { | ||
| throw new BigQueryJdbcException( | ||
| String.format("Invalid BigQuery identifier: '%s'", identifier)); | ||
| } | ||
| return GOOGLE_SQL_QUOTED_IDENTIFIER + identifier + GOOGLE_SQL_QUOTED_IDENTIFIER; | ||
| } | ||
|
|
||
| boolean shouldSkipGraphTables(String tableNamePattern, Set<String> requestedTypes) { | ||
| if (this.connection == null || !this.connection.isEnableGraphMetadata()) { | ||
| return true; | ||
| } | ||
| if (requestedTypes != null && !requestedTypes.contains(GRAPH_TABLE_TYPE)) { | ||
| return true; | ||
| } | ||
| if (tableNamePattern != null | ||
| && !tableNamePattern.contains("%") | ||
| && !tableNamePattern.contains(".")) { | ||
|
Neenu1995 marked this conversation as resolved.
Resolved
|
||
| return true; | ||
| } | ||
| return false; | ||
| } | ||
|
|
||
| String resolveGraphRegion() throws SQLException { | ||
| String location = this.connection != null ? this.connection.getLocation() : null; | ||
| if (location == null || location.trim().isEmpty()) { | ||
| location = DEFAULT_GRAPH_REGION; | ||
|
Neenu1995 marked this conversation as resolved.
Resolved
|
||
| LOG.info(String.format("Deafulting to Graph Location %s", DEFAULT_GRAPH_REGION)); | ||
| } | ||
| String normalized = location.trim().toLowerCase(Locale.US); | ||
| return normalized.startsWith("region-") ? normalized : "region-" + normalized; | ||
| } | ||
|
|
||
| void appendGraphTableRows( | ||
| String effectiveCatalog, | ||
| String effectiveSchemaPattern, | ||
| String tableNamePattern, | ||
| Set<String> requestedTypes, | ||
| List<FieldValueList> collectedResults, | ||
| FieldList resultSchemaResults) | ||
| throws SQLException { | ||
| if (shouldSkipGraphTables(tableNamePattern, requestedTypes)) { | ||
| return; | ||
| } | ||
|
|
||
| LOG.info("Appending Graph table rows"); | ||
|
|
||
| List<String> targetCatalogs = | ||
| effectiveCatalog != null | ||
| ? Collections.singletonList(effectiveCatalog) | ||
| : getAccessibleCatalogNames(); | ||
| boolean isExactDataset = | ||
| effectiveSchemaPattern != null && !needsListing(effectiveSchemaPattern); | ||
| String effectiveTablePattern = (tableNamePattern == null) ? "%" : tableNamePattern; | ||
|
|
||
| for (String project : targetCatalogs) { | ||
| String formattedSql; | ||
| if (isExactDataset) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: I'd rather do it this way, easier to see what's changing |
||
| formattedSql = | ||
| replaceSqlParameters( | ||
| getGraphTablesDatasetSqlContent(), | ||
| quoteIdentifier(project), | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Per comment in SQL, rather than escaping identifiers here, better to do it at SQL file level |
||
| quoteIdentifier(effectiveSchemaPattern)); | ||
| } else { | ||
| formattedSql = | ||
| replaceSqlParameters( | ||
| getGraphTablesRegionSqlContent(), | ||
| quoteIdentifier(project), | ||
| quoteIdentifier(resolveGraphRegion())); | ||
| } | ||
|
|
||
| try (PreparedStatement statement = this.connection.prepareStatement(formattedSql)) { | ||
| if (isExactDataset) { | ||
| statement.setString(1, effectiveTablePattern); | ||
| } else { | ||
| statement.setString(1, effectiveSchemaPattern == null ? "%" : effectiveSchemaPattern); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we can have single SQL query, having an extra condition |
||
| statement.setString(2, effectiveTablePattern); | ||
| } | ||
|
|
||
| try (ResultSet rs = statement.executeQuery()) { | ||
| while (rs.next()) { | ||
| String tableCat = rs.getString("TABLE_CAT"); | ||
| String tableSchem = rs.getString("TABLE_SCHEM"); | ||
| String tableName = rs.getString("TABLE_NAME"); | ||
| String sourceTable = rs.getString("SOURCE_TABLE"); | ||
| String remarks = sourceTable != null ? "Source: " + sourceTable : null; | ||
|
|
||
| List<FieldValue> values = new ArrayList<>(resultSchemaResults.size()); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you add inline comments what maps to what? |
||
| values.add(createStringFieldValue(tableCat)); | ||
| values.add(createStringFieldValue(tableSchem)); | ||
| values.add(createStringFieldValue(tableName)); | ||
| values.add(createStringFieldValue(GRAPH_TABLE_TYPE)); | ||
| values.add(createStringFieldValue(remarks)); | ||
| values.add(createNullFieldValue()); | ||
| values.add(createNullFieldValue()); | ||
| values.add(createNullFieldValue()); | ||
| values.add(createNullFieldValue()); | ||
| values.add(createNullFieldValue()); | ||
| collectedResults.add(FieldValueList.of(values, resultSchemaResults)); | ||
|
Neenu1995 marked this conversation as resolved.
Resolved
|
||
| } | ||
| } | ||
| } catch (SQLException e) { | ||
| if (e.getMessage() != null && e.getMessage().contains("Not found")) { | ||
| LOG.info( | ||
| "Dataset or project not found while querying PROPERTY_GRAPHS in project '%s'." | ||
| + " Skipping.", | ||
| project); | ||
| } else { | ||
| throw e; | ||
| } | ||
| } | ||
|
Neenu1995 marked this conversation as resolved.
Resolved
|
||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public ResultSet getSchemas() throws SQLException { | ||
| LOG.info("getSchemas() called"); | ||
|
|
@@ -1488,7 +1636,8 @@ public ResultSet getTableTypes() throws SQLException { | |
| LOG.info("getTableTypes() called"); | ||
|
|
||
| final Schema tableTypesSchema = defineGetTableTypesSchema(); | ||
| final List<FieldValueList> tableTypeRows = prepareGetTableTypesRows(tableTypesSchema); | ||
| final List<FieldValueList> tableTypeRows = | ||
| prepareGetTableTypesRows(tableTypesSchema, this.connection.isEnableGraphMetadata()); | ||
|
|
||
| BlockingQueue<BigQueryFieldValueListWrapper> queue = | ||
| fetchAndPopulateQueue(() -> tableTypeRows, tableTypesSchema.getFields()); | ||
|
|
@@ -1504,9 +1653,13 @@ static Schema defineGetTableTypesSchema() { | |
| } | ||
|
|
||
| static List<FieldValueList> prepareGetTableTypesRows(Schema schema) { | ||
| final String[] tableTypes = {"EXTERNAL", "MATERIALIZED VIEW", "SNAPSHOT", "TABLE", "VIEW"}; | ||
| List<FieldValueList> rows = new ArrayList<>(tableTypes.length); | ||
| return prepareGetTableTypesRows(schema, false); | ||
| } | ||
|
|
||
| static List<FieldValueList> prepareGetTableTypesRows(Schema schema, boolean includeGraphTables) { | ||
| final String[] tableTypes = includeGraphTables ? TABLE_TYPES_WITH_GRAPH : DEFAULT_TABLE_TYPES; | ||
| FieldList schemaFields = schema.getFields(); | ||
| List<FieldValueList> rows = new ArrayList<>(tableTypes.length); | ||
|
|
||
| for (String typeName : tableTypes) { | ||
| FieldValue fieldValue = FieldValue.of(FieldValue.Attribute.PRIMITIVE, typeName); | ||
|
|
@@ -4853,6 +5006,20 @@ private static synchronized String getExportedKeysSqlContent() { | |
| return exportedKeysSqlContent; | ||
| } | ||
|
|
||
| static synchronized String getGraphTablesDatasetSqlContent() { | ||
| if (graphTablesDatasetSqlContent == null) { | ||
| graphTablesDatasetSqlContent = readSqlFromFile(GET_GRAPH_TABLES_DATASET_SQL); | ||
| } | ||
| return graphTablesDatasetSqlContent; | ||
| } | ||
|
|
||
| static synchronized String getGraphTablesRegionSqlContent() { | ||
| if (graphTablesRegionSqlContent == null) { | ||
| graphTablesRegionSqlContent = readSqlFromFile(GET_GRAPH_TABLES_REGION_SQL); | ||
| } | ||
| return graphTablesRegionSqlContent; | ||
| } | ||
|
|
||
| static String readSqlFromFile(String filename) { | ||
| try (InputStream in = BigQueryDatabaseMetaData.class.getResourceAsStream(filename)) { | ||
| if (in == null) { | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,29 @@ | ||
| /* | ||
| * Copyright 2026 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * https://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| SELECT | ||
| property_graph_catalog AS TABLE_CAT, | ||
| property_graph_schema AS TABLE_SCHEM, | ||
| CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name')) AS TABLE_NAME, | ||
| CONCAT( | ||
| JSON_VALUE(t, '$.dataSourceTable.projectId'), '.', | ||
| JSON_VALUE(t, '$.dataSourceTable.datasetId'), '.', | ||
| JSON_VALUE(t, '$.dataSourceTable.tableId')) AS SOURCE_TABLE | ||
| FROM %1$s.%2$s.INFORMATION_SCHEMA.PROPERTY_GRAPHS, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Add escape ``` characters for the table name |
||
| UNNEST(ARRAY_CONCAT( | ||
| IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.nodeTables'), []), | ||
| IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.edgeTables'), []))) AS t | ||
| WHERE LOWER(CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name'))) LIKE LOWER(?) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| /* | ||
| * Copyright 2026 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * https://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| SELECT | ||
| property_graph_catalog AS TABLE_CAT, | ||
| property_graph_schema AS TABLE_SCHEM, | ||
| CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name')) AS TABLE_NAME, | ||
| CONCAT( | ||
| JSON_VALUE(t, '$.dataSourceTable.projectId'), '.', | ||
| JSON_VALUE(t, '$.dataSourceTable.datasetId'), '.', | ||
| JSON_VALUE(t, '$.dataSourceTable.tableId')) AS SOURCE_TABLE | ||
| FROM %1$s.%2$s.INFORMATION_SCHEMA.PROPERTY_GRAPHS, | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Escape table name |
||
| UNNEST(ARRAY_CONCAT( | ||
| IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.nodeTables'), []), | ||
| IFNULL(JSON_QUERY_ARRAY(property_graph_metadata_json, '$.edgeTables'), []))) AS t | ||
| WHERE LOWER(property_graph_schema) LIKE LOWER(?) | ||
| AND LOWER(CONCAT(property_graph_name, '.', JSON_VALUE(t, '$.name'))) LIKE LOWER(?) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: can we update it to be
DEFAULT_TABLE_TYPES + {"GRAPH TABLE"}in some way?