Add BigQuery table list cache (trinodb#4)

regadas · Aug 16, 2023 · 0a2c36a · 0a2c36a
1 parent c53ba9c
commit 0a2c36a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java
@@ -69,6 +69,7 @@
 import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_FAILED_TO_EXECUTE_QUERY;
 import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_INVALID_STATEMENT;
 import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_LISTING_DATASET_ERROR;
+import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_LISTING_TABLE_ERROR;
 import static io.trino.plugin.bigquery.BigQuerySessionProperties.createDisposition;
 import static io.trino.plugin.bigquery.BigQuerySessionProperties.isQueryResultsCacheEnabled;
 import static java.lang.String.format;
@@ -88,6 +89,7 @@ public class BigQueryClient
     private final ViewMaterializationCache materializationCache;
     private final boolean caseInsensitiveNameMatching;
     private final LoadingCache<String, List<Dataset>> remoteDatasetCache;
+    private final LoadingCache<DatasetId, List<Table>> remoteTableCache;
     private final Optional<String> configProjectId;
 
     public BigQueryClient(
@@ -106,6 +108,10 @@ public BigQueryClient(
                 .expireAfterWrite(metadataCacheTtl.toMillis(), MILLISECONDS)
                 .shareNothingWhenDisabled()
                 .build(CacheLoader.from(this::listDatasetsFromBigQuery));
+        this.remoteTableCache = EvictableCacheBuilder.newBuilder()
+                .expireAfterWrite(metadataCacheTtl.toMillis(), MILLISECONDS)
+                .shareNothingWhenDisabled()
+                .build(CacheLoader.from(this::listTablesFromBigQuery));
         this.configProjectId = requireNonNull(configProjectId, "projectId is null");
     }
 
@@ -235,6 +241,16 @@ private List<Dataset> listDatasetsFromBigQuery(String projectId)
     }
 
     public Iterable<Table> listTables(DatasetId remoteDatasetId)
+    {
+        try {
+            return remoteTableCache.get(remoteDatasetId);
+        }
+        catch (ExecutionException e) {
+            throw new TrinoException(BIGQUERY_LISTING_TABLE_ERROR, "Failed to retrieve tables from BigQuery", e);
+        }
+    }
+
+    private List<Table> listTablesFromBigQuery(DatasetId remoteDatasetId)
     {
         Iterable<Table> allTables = bigQuery.listTables(remoteDatasetId).iterateAll();
         return stream(allTables)

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
@@ -31,6 +31,7 @@ public enum BigQueryErrorCode
     BIGQUERY_UNSUPPORTED_OPERATION(5, USER_ERROR),
     BIGQUERY_INVALID_STATEMENT(6, USER_ERROR),
     BIGQUERY_PROXY_SSL_INITIALIZATION_FAILED(7, EXTERNAL),
+    BIGQUERY_LISTING_TABLE_ERROR(8, EXTERNAL),
     /**/;
 
     private final ErrorCode errorCode;