Support reading wildcard tables in BigQuery

trinodb · May 23, 2022 · a398d66 · a398d66
1 parent be2b53c
commit a398d66
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 0 deletions.
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
@@ -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
@@ -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
@@ -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<BigQuerySplit> 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
@@ -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
@@ -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);