Ignore disappeared datasets during listing tables in BigQuery

trinodb · Nov 15, 2021 · a0d3dc2 · a0d3dc2
1 parent 884bc09
commit a0d3dc2
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 9 deletions.
diff --git a/plugin/trino-bigquery/pom.xml b/plugin/trino-bigquery/pom.xml
@@ -334,6 +334,7 @@
                                 <!-- If you are adding entry here also add an entry to cloud-tests or cloud-tests-case-insensitive-mapping profile below -->
                                 <exclude>**/TestBigQueryIntegrationSmokeTest.java</exclude>
                                 <exclude>**/TestBigQueryTypeMapping.java</exclude>
+                                <exclude>**/TestBigQueryMetadata.java</exclude>
                                 <exclude>**/TestBigQueryInstanceCleaner.java</exclude>
                                 <exclude>**/TestBigQueryCaseInsensitiveMapping.java</exclude>
                             </excludes>
@@ -357,6 +358,7 @@
                             <includes>
                                 <include>**/TestBigQueryIntegrationSmokeTest.java</include>
                                 <include>**/TestBigQueryTypeMapping.java</include>
+                                <include>**/TestBigQueryMetadata.java</include>
                                 <include>**/TestBigQueryInstanceCleaner.java</include>
                             </includes>
                         </configuration>

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
@@ -25,7 +25,9 @@ public enum BigQueryErrorCode
     BIGQUERY_VIEW_DESTINATION_TABLE_CREATION_FAILED(0, EXTERNAL),
     BIGQUERY_DATETIME_PARSING_ERROR(1, EXTERNAL),
     BIGQUERY_FAILED_TO_EXECUTE_QUERY(2, EXTERNAL),
-    BIGQUERY_AMBIGUOUS_OBJECT_NAME(3, EXTERNAL);
+    BIGQUERY_AMBIGUOUS_OBJECT_NAME(3, EXTERNAL),
+    BIGQUERY_LISTING_DATASET_ERROR(4, EXTERNAL),
+    /**/;
 
     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
@@ -30,6 +30,7 @@
 import com.google.common.collect.Streams;
 import io.airlift.log.Logger;
 import io.trino.plugin.bigquery.BigQueryClient.RemoteDatabaseObject;
+import io.trino.spi.TrinoException;
 import io.trino.spi.connector.Assignment;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
@@ -70,6 +71,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 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.BigQueryType.toField;
 import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
@@ -152,14 +154,25 @@ public List<SchemaTableName> listTables(ConnectorSession session, Optional<Strin
 
         ImmutableList.Builder<SchemaTableName> tableNames = ImmutableList.builder();
         for (String remoteSchemaName : remoteSchemaNames) {
-            Iterable<Table> tables = bigQueryClient.listTables(DatasetId.of(projectId, remoteSchemaName), TABLE, VIEW);
-            for (Table table : tables) {
-                // filter ambiguous tables
-                bigQueryClient.toRemoteTable(projectId, remoteSchemaName, table.getTableId().getTable().toLowerCase(ENGLISH), tables)
-                        .filter(RemoteDatabaseObject::isAmbiguous)
-                        .ifPresentOrElse(
-                                remoteTable -> log.debug("Filtered out [%s.%s] from list of tables due to ambiguous name", remoteSchemaName, table.getTableId().getTable()),
-                                () -> tableNames.add(new SchemaTableName(table.getTableId().getDataset(), table.getTableId().getTable())));
+            try {
+                Iterable<Table> tables = bigQueryClient.listTables(DatasetId.of(projectId, remoteSchemaName), TABLE, VIEW);
+                for (Table table : tables) {
+                    // filter ambiguous tables
+                    bigQueryClient.toRemoteTable(projectId, remoteSchemaName, table.getTableId().getTable().toLowerCase(ENGLISH), tables)
+                            .filter(RemoteDatabaseObject::isAmbiguous)
+                            .ifPresentOrElse(
+                                    remoteTable -> log.debug("Filtered out [%s.%s] from list of tables due to ambiguous name", remoteSchemaName, table.getTableId().getTable()),
+                                    () -> tableNames.add(new SchemaTableName(table.getTableId().getDataset(), table.getTableId().getTable())));
+                }
+            }
+            catch (BigQueryException e) {
+                if (e.getCode() == 404 && e.getMessage().contains("Not found: Dataset")) {
+                    // Dataset not found error is ignored because listTables is used for metadata queries (SELECT FROM information_schema)
+                    log.debug("Dataset disappeared during listing operation: %s", remoteSchemaName);
+                }
+                else {
+                    throw new TrinoException(BIGQUERY_LISTING_DATASET_ERROR, "Exception happened during listing BigQuery dataset: " + remoteSchemaName, e);
+                }
             }
         }
         return tableNames.build();

diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/BigQueryQueryRunner.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/BigQueryQueryRunner.java
@@ -157,6 +157,11 @@ public List<String> getTableNames(String dataset)
             return tableNames.build();
         }
 
+        public BigQuery getBigQuery()
+        {
+            return bigQuery;
+        }
+
         private static BigQuery createBigQueryClient()
         {
             try {

diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryMetadata.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryMetadata.java
@@ -0,0 +1,34 @@
+/*
+ * 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
+ *
+ *     http://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.
+ */
+package io.trino.plugin.bigquery;
+
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.BigQueryException;
+import org.testng.annotations.Test;
+
+import static io.trino.plugin.bigquery.BigQueryQueryRunner.BigQuerySqlExecutor;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+
+public class TestBigQueryMetadata
+{
+    @Test
+    public void testDatasetNotFoundMessage()
+    {
+        // Update the catch block condition in 'BigQueryMetadata.listTables()' method if this test failed
+        BigQuery bigQuery = new BigQuerySqlExecutor().getBigQuery();
+        assertThatExceptionOfType(BigQueryException.class)
+                .isThrownBy(() -> bigQuery.listTables("test_dataset_not_found"))
+                .matches(e -> e.getCode() == 404 && e.getMessage().contains("Not found: Dataset"));
+    }
+}