From 9f267c1063f9d3d4632aa3d70db1a89d513e5392 Mon Sep 17 00:00:00 2001
From: Slawomir Pajak <slawomir.pajak@starburstdata.com>
Date: Mon, 7 Oct 2024 11:24:41 +0200
Subject: [PATCH] Fix BigQuery case-insensitive mapping when cache is disabled

When cache is disabled, all values put into cache are not persisted, causing cache-insensitive mapping failing
This PR decouples mapping from cache handling
---
 plugin/trino-bigquery/pom.xml                 |  2 ++
 .../trino/plugin/bigquery/BigQueryClient.java | 10 ++++--
 .../TestBigQueryCaseInsensitiveMapping.java   | 33 +++++++++++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java

diff --git a/plugin/trino-bigquery/pom.xml b/plugin/trino-bigquery/pom.xml
index 93261cde6c69..370d48475676 100644
--- a/plugin/trino-bigquery/pom.xml
+++ b/plugin/trino-bigquery/pom.xml
@@ -534,6 +534,7 @@
                                 <exclude>**/TestBigQueryMetadata.java</exclude>
                                 <exclude>**/TestBigQuerySplitManager.java</exclude>
                                 <exclude>**/TestBigQueryInstanceCleaner.java</exclude>
+                                <exclude>**/TestBigQueryCaseInsensitiveMapping.java</exclude>
                                 <exclude>**/TestBigQueryCaseInsensitiveMappingWithCache.java</exclude>
                                 <exclude>**/TestBigQuery*FailureRecoveryTest.java</exclude>
                                 <exclude>**/TestBigQueryWithProxyTest.java</exclude>
@@ -579,6 +580,7 @@
                         <artifactId>maven-surefire-plugin</artifactId>
                         <configuration>
                             <includes>
+                                <include>**/TestBigQueryCaseInsensitiveMapping.java</include>
                                 <include>**/TestBigQueryCaseInsensitiveMappingWithCache.java</include>
                                 <!-- Also included here to make sure the case-insesitive project also gets cleaned up -->
                                 <include>**/TestBigQueryInstanceCleaner.java</include>
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
index 7e70dbf86707..1781c2d6f0b6 100644
--- 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
@@ -51,6 +51,7 @@
 import io.trino.spi.connector.TableNotFoundException;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -171,12 +172,15 @@ public Optional<RemoteDatabaseObject> toRemoteDataset(String projectId, String d
         }
 
         // Get all information from BigQuery and update cache from all fetched information
+        Map<DatasetId, RemoteDatabaseObject> mapping = new HashMap<>(remoteDatasetCaseInsensitiveCache.getAllPresent(remoteDatasetCaseInsensitiveCache.asMap().keySet()));
         for (DatasetId datasetId : datasetIds.get()) {
             DatasetId newCacheKey = datasetIdToLowerCase(datasetId);
             RemoteDatabaseObject newValue = RemoteDatabaseObject.of(datasetId.getDataset());
+            mapping.merge(newCacheKey, newValue, (currentValue, collision) -> currentValue.registerCollision(collision.getOnlyRemoteName()));
             updateCache(remoteDatasetCaseInsensitiveCache, newCacheKey, newValue);
         }
-        return Optional.ofNullable(remoteDatasetCaseInsensitiveCache.getIfPresent(cacheKey));
+
+        return Optional.ofNullable(mapping.get(cacheKey));
     }
 
     public Optional<RemoteDatabaseObject> toRemoteTable(ConnectorSession session, String projectId, String remoteDatasetName, String tableName)
@@ -207,13 +211,15 @@ private Optional<RemoteDatabaseObject> toRemoteTable(String projectId, String re
         }
 
         // Get all information from BigQuery and update cache from all fetched information
+        Map<TableId, RemoteDatabaseObject> mapping = new HashMap<>(remoteTableCaseInsensitiveCache.getAllPresent(remoteTableCaseInsensitiveCache.asMap().keySet()));
         for (TableId table : tableIds.get()) {
             TableId newCacheKey = tableIdToLowerCase(table);
             RemoteDatabaseObject newValue = RemoteDatabaseObject.of(table.getTable());
+            mapping.merge(newCacheKey, newValue, (currentValue, collision) -> currentValue.registerCollision(collision.getOnlyRemoteName()));
             updateCache(remoteTableCaseInsensitiveCache, newCacheKey, newValue);
         }
 
-        return Optional.ofNullable(remoteTableCaseInsensitiveCache.getIfPresent(cacheKey));
+        return Optional.ofNullable(mapping.get(cacheKey));
     }
 
     private static <T> void updateCache(Cache<T, RemoteDatabaseObject> caseInsensitiveCache, T newCacheKey, RemoteDatabaseObject newValue)
diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java
new file mode 100644
index 000000000000..5d5afeb7d5d9
--- /dev/null
+++ b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryCaseInsensitiveMapping.java
@@ -0,0 +1,33 @@
+/*
+ * 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.common.collect.ImmutableMap;
+import io.trino.testing.QueryRunner;
+
+final class TestBigQueryCaseInsensitiveMapping
+        extends BaseBigQueryCaseInsensitiveMapping
+{
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return BigQueryQueryRunner.builder()
+                .setConnectorProperties(ImmutableMap.<String, String>builder()
+                        .put("bigquery.case-insensitive-name-matching", "true")
+                        .put("bigquery.case-insensitive-name-matching.cache-ttl", "0m")
+                        .buildOrThrow())
+                .build();
+    }
+}