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 index bf0d397ce2e..4766e92f93d 100644 --- 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 listDatasets(String projectId) return () -> Iterators.transform(datasets, this::addDataSetMappingIfNeeded); } - Iterable listTables(DatasetId datasetId) + Iterable
listTables(DatasetId datasetId, TableDefinition.Type... types) { + Set allowedTypes = ImmutableSet.copyOf(types); DatasetId bigQueryDatasetId = datasetIds.getOrDefault(datasetId, datasetId); - return bigQuery.listTables(bigQueryDatasetId).iterateAll(); + Iterable
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 index 1afe1b4abdc..17e6b750ac8 100644 --- 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 listSchemaNames(ConnectorSession session) public List listTables(ConnectorSession session, Optional schemaName) { log.debug("listTables(session=%s, schemaName=%s)", session, schemaName); + return listTablesWithTypes(session, schemaName, TABLE); + } + + @Override + public List listViews(ConnectorSession session, Optional schemaName) + { + log.debug("listViews(session=%s, schemaName=%s)", session, schemaName); + return listTablesWithTypes(session, schemaName, VIEW); + } + + private List listTablesWithTypes(ConnectorSession session, Optional schemaName, TableDefinition.Type... types) + { if (schemaName.isPresent() && schemaName.get().equalsIgnoreCase(INFORMATION_SCHEMA)) { return ImmutableList.of(); } @@ -91,7 +106,7 @@ public List listTables(ConnectorSession session, Optional 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())); } }