From a398d66aa57ba14b61c5a208d9afc4f62c832ec6 Mon Sep 17 00:00:00 2001 From: Yuya Ebihara Date: Tue, 17 May 2022 16:09:36 +0900 Subject: [PATCH] Support reading wildcard tables in BigQuery --- .../plugin/bigquery/BigQueryErrorCode.java | 2 + .../plugin/bigquery/BigQueryMetadata.java | 5 ++ .../plugin/bigquery/BigQuerySplitManager.java | 5 ++ .../trino/plugin/bigquery/BigQueryUtil.java | 7 +++ .../bigquery/TestBigQueryConnectorTest.java | 55 +++++++++++++++++++ 5 files changed, 74 insertions(+) diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryErrorCode.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryErrorCode.java index 9abfeeaeba7d..a4ac606fa3e1 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryErrorCode.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryErrorCode.java @@ -18,6 +18,7 @@ import io.trino.spi.ErrorType; import static io.trino.spi.ErrorType.EXTERNAL; +import static io.trino.spi.ErrorType.USER_ERROR; public enum BigQueryErrorCode implements ErrorCodeSupplier @@ -27,6 +28,7 @@ public enum BigQueryErrorCode BIGQUERY_FAILED_TO_EXECUTE_QUERY(2, EXTERNAL), BIGQUERY_AMBIGUOUS_OBJECT_NAME(3, EXTERNAL), BIGQUERY_LISTING_DATASET_ERROR(4, EXTERNAL), + BIGQUERY_UNSUPPORTED_OPERATION(5, USER_ERROR), /**/; private final ErrorCode errorCode; diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java index 7789d42a601e..9efb930d6ec3 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java @@ -73,10 +73,12 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_LISTING_DATASET_ERROR; +import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_UNSUPPORTED_OPERATION; import static io.trino.plugin.bigquery.BigQueryPseudoColumn.PARTITION_DATE; import static io.trino.plugin.bigquery.BigQueryPseudoColumn.PARTITION_TIME; import static io.trino.plugin.bigquery.BigQueryTableHandle.BigQueryPartitionType.INGESTION; import static io.trino.plugin.bigquery.BigQueryType.toField; +import static io.trino.plugin.bigquery.BigQueryUtil.isWildcardTable; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; import static java.util.function.Function.identity; @@ -406,6 +408,9 @@ public void dropTable(ConnectorSession session, ConnectorTableHandle tableHandle { BigQueryClient client = bigQueryClientFactory.create(session); BigQueryTableHandle bigQueryTable = (BigQueryTableHandle) tableHandle; + if (isWildcardTable(TableDefinition.Type.valueOf(bigQueryTable.getType()), bigQueryTable.getRemoteTableName().getTableName())) { + throw new TrinoException(BIGQUERY_UNSUPPORTED_OPERATION, "This connector does not support dropping wildcard tables"); + } TableId tableId = bigQueryTable.getRemoteTableName().toTableId(); client.dropTable(tableId); } diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQuerySplitManager.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQuerySplitManager.java index a479fe66b3ae..deb01cea1594 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQuerySplitManager.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQuerySplitManager.java @@ -49,6 +49,7 @@ import static io.trino.plugin.bigquery.BigQuerySessionProperties.createDisposition; import static io.trino.plugin.bigquery.BigQuerySessionProperties.isQueryResultsCacheEnabled; import static io.trino.plugin.bigquery.BigQuerySessionProperties.isSkipViewMaterialization; +import static io.trino.plugin.bigquery.BigQueryUtil.isWildcardTable; import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED; import static java.util.Objects.requireNonNull; import static java.util.stream.Collectors.toList; @@ -117,6 +118,10 @@ private List readFromBigQuery(ConnectorSession session, TableDefi .map(column -> ((BigQueryColumnHandle) column).getName()) .collect(toImmutableList()); + if (isWildcardTable(type, remoteTableId.getTable())) { + // Storage API doesn't support reading wildcard tables + return ImmutableList.of(BigQuerySplit.forViewStream(columns, filter)); + } if (type == MATERIALIZED_VIEW) { // Storage API doesn't support reading materialized views return ImmutableList.of(BigQuerySplit.forViewStream(columns, filter)); diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java index 69faf8ebdb25..e98bcc0dba01 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryUtil.java @@ -15,6 +15,7 @@ import com.google.cloud.bigquery.BigQueryError; import com.google.cloud.bigquery.BigQueryException; +import com.google.cloud.bigquery.TableDefinition; import com.google.common.collect.ImmutableSet; import io.grpc.Status; import io.grpc.StatusRuntimeException; @@ -23,6 +24,7 @@ import java.util.Optional; import java.util.Set; +import static com.google.cloud.bigquery.TableDefinition.Type.TABLE; import static com.google.cloud.http.BaseHttpServiceException.UNKNOWN_CODE; import static com.google.common.base.Throwables.getCausalChain; @@ -66,4 +68,9 @@ public static String toBigQueryColumnName(String columnName) } return columnName; } + + public static boolean isWildcardTable(TableDefinition.Type type, String tableName) + { + return type == TABLE && tableName.contains("*"); + } } diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java index 8f15601f37b3..51047711468e 100644 --- a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java +++ b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryConnectorTest.java @@ -651,6 +651,61 @@ public void testQueryCache() } } + @Test + public void testWildcardTable() + { + String suffix = randomTableSuffix(); + String firstTable = format("test_wildcard_%s_1", suffix); + String secondTable = format("test_wildcard_%s_2", suffix); + String wildcardTable = format("test_wildcard_%s_*", suffix); + try { + onBigQuery("CREATE TABLE test." + firstTable + " AS SELECT 1 AS value"); + onBigQuery("CREATE TABLE test." + secondTable + " AS SELECT 2 AS value"); + + assertQuery("DESCRIBE test.\"" + wildcardTable + "\"", "VALUES ('value', 'bigint', '', '')"); + assertQuery("SELECT * FROM test.\"" + wildcardTable + "\"", "VALUES (1), (2)"); + + // Unsupported operations + assertQueryFails("DROP TABLE test.\"" + wildcardTable + "\"", "This connector does not support dropping wildcard tables"); + assertQueryFails("INSERT INTO test.\"" + wildcardTable + "\" VALUES (1)", "This connector does not support inserts"); + assertQueryFails("ALTER TABLE test.\"" + wildcardTable + "\" ADD COLUMN new_column INT", "This connector does not support adding columns"); + assertQueryFails("ALTER TABLE test.\"" + wildcardTable + "\" RENAME TO test.new_wildcard_table", "This connector does not support renaming tables"); + } + finally { + onBigQuery("DROP TABLE IF EXISTS test." + firstTable); + onBigQuery("DROP TABLE IF EXISTS test." + secondTable); + } + } + + @Test + public void testWildcardTableWithDifferentColumnDefinition() + { + String suffix = randomTableSuffix(); + String firstTable = format("test_invalid_wildcard_%s_1", suffix); + String secondTable = format("test_invalid_wildcard_%s_2", suffix); + String wildcardTable = format("test_invalid_wildcard_%s_*", suffix); + try { + onBigQuery("CREATE TABLE test." + firstTable + " AS SELECT 1 AS value"); + onBigQuery("CREATE TABLE test." + secondTable + " AS SELECT 'string' AS value"); + + assertQuery("DESCRIBE test.\"" + wildcardTable + "\"", "VALUES ('value', 'varchar', '', '')"); + + assertThatThrownBy(() -> query("SELECT * FROM test.\"" + wildcardTable + "\"")) + .hasMessageContaining("Cannot read field of type INT64 as STRING Field: value"); + } + finally { + onBigQuery("DROP TABLE IF EXISTS test." + firstTable); + onBigQuery("DROP TABLE IF EXISTS test." + secondTable); + } + } + + @Test + public void testMissingWildcardTable() + { + assertThatThrownBy(() -> query("SELECT * FROM test.\"test_missing_wildcard_table_*\"")) + .hasMessageEndingWith("does not match any table."); + } + private void onBigQuery(@Language("SQL") String sql) { bigQuerySqlExecutor.execute(sql);