table type is now correctly shown in the information_schema

trinodb · Feb 28, 2020 · bcdfad9 · bcdfad9
1 parent 07d3da5
commit bcdfad9
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryClient.java b/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryClient.java
@@ -21,19 +21,24 @@
 import com.google.cloud.bigquery.JobInfo;
 import com.google.cloud.bigquery.QueryJobConfiguration;
 import com.google.cloud.bigquery.Table;
+import com.google.cloud.bigquery.TableDefinition;
 import com.google.cloud.bigquery.TableId;
 import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.TableResult;
 import com.google.cloud.http.BaseHttpServiceException;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterators;
 
 import java.util.Iterator;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.stream.StreamSupport;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.lang.String.format;
 import static java.util.Locale.ENGLISH;
 import static java.util.UUID.randomUUID;
@@ -90,10 +95,14 @@ Iterable<Dataset> listDatasets(String projectId)
         return () -> Iterators.transform(datasets, this::addDataSetMappingIfNeeded);
     }
 
-    Iterable<Table> listTables(DatasetId datasetId)
+    Iterable<Table> listTables(DatasetId datasetId, TableDefinition.Type... types)
     {
+        Set<TableDefinition.Type> allowedTypes = ImmutableSet.copyOf(types);
         DatasetId bigQueryDatasetId = datasetIds.getOrDefault(datasetId, datasetId);
-        return bigQuery.listTables(bigQueryDatasetId).iterateAll();
+        Iterable<Table> allTables = bigQuery.listTables(bigQueryDatasetId).iterateAll();
+        return StreamSupport.stream(allTables.spliterator(), false)
+                .filter(table -> allowedTypes.contains(table.getDefinition().getType()))
+                .collect(toImmutableList());
     }
 
     private Dataset addDataSetMappingIfNeeded(Dataset dataset)

diff --git a/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryMetadata.java b/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryMetadata.java
@@ -18,6 +18,7 @@
 import com.google.cloud.bigquery.Field;
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.Table;
+import com.google.cloud.bigquery.TableDefinition;
 import com.google.cloud.bigquery.TableId;
 import com.google.cloud.bigquery.TableInfo;
 import com.google.common.collect.ImmutableList;
@@ -48,6 +49,8 @@
 import java.util.Optional;
 import java.util.Set;
 
+import static com.google.cloud.bigquery.TableDefinition.Type.TABLE;
+import static com.google.cloud.bigquery.TableDefinition.Type.VIEW;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toMap;
@@ -83,6 +86,18 @@ public List<String> listSchemaNames(ConnectorSession session)
     public List<SchemaTableName> listTables(ConnectorSession session, Optional<String> schemaName)
     {
         log.debug("listTables(session=%s, schemaName=%s)", session, schemaName);
+        return listTablesWithTypes(session, schemaName, TABLE);
+    }
+
+    @Override
+    public List<SchemaTableName> listViews(ConnectorSession session, Optional<String> schemaName)
+    {
+        log.debug("listViews(session=%s, schemaName=%s)", session, schemaName);
+        return listTablesWithTypes(session, schemaName, VIEW);
+    }
+
+    private List<SchemaTableName> listTablesWithTypes(ConnectorSession session, Optional<String> schemaName, TableDefinition.Type... types)
+    {
         if (schemaName.isPresent() && schemaName.get().equalsIgnoreCase(INFORMATION_SCHEMA)) {
             return ImmutableList.of();
         }
@@ -91,7 +106,7 @@ public List<SchemaTableName> listTables(ConnectorSession session, Optional<Strin
 
         ImmutableList.Builder<SchemaTableName> tableNames = ImmutableList.builder();
         for (String datasetId : schemaNames) {
-            for (Table table : bigQueryClient.listTables(DatasetId.of(projectId, datasetId))) {
+            for (Table table : bigQueryClient.listTables(DatasetId.of(projectId, datasetId), types)) {
                 tableNames.add(new SchemaTableName(datasetId, table.getTableId().getTable()));
             }
         }