Avoid null cache writes for query routing mappings

Guard query-id cache writes in routing manager against null/empty values
to prevent runtime failures when query history rows have been removed
by retention.

- Use direct static import for isNullOrEmpty.
- Change log level from warn to debug for empty queryId checks.
- Standardize null/empty handling across all cache setters.
- Add regression tests for null/empty external URL and query ID.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java

@@ -17,7 +17,6 @@
 import com.github.benmanes.caffeine.cache.LoadingCache;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
-import com.google.common.base.Strings;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.clustermonitor.ClusterStats;
 import io.trino.gateway.ha.clustermonitor.TrinoStatus;
@@ -40,6 +39,8 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
+import static com.google.common.base.Strings.isNullOrEmpty;
+
 /**
  * This class performs health check, stats counts for each backend and provides a backend given
  * request object. Default implementation comes here.
@@ -76,12 +77,30 @@ public BaseRoutingManager(GatewayBackendManager gatewayBackendManager, QueryHist
     @Override
     public void setBackendForQueryId(String queryId, String backend)
     {
+        if (isNullOrEmpty(queryId)) {
+            log.debug("Skipping backend cache update for empty queryId");
+            return;
+        }
+        if (isNullOrEmpty(backend)) {
+            queryIdBackendCache.invalidate(queryId);
+            log.debug("Invalidated backend cache for queryId [%s] due to empty backend", queryId);
+            return;
+        }
         queryIdBackendCache.put(queryId, backend);
     }
 
     @Override
     public void setRoutingGroupForQueryId(String queryId, String routingGroup)
     {
+        if (isNullOrEmpty(queryId)) {
+            log.debug("Skipping routing group cache update for empty queryId");
+            return;
+        }
+        if (isNullOrEmpty(routingGroup)) {
+            queryIdRoutingGroupCache.invalidate(queryId);
+            log.debug("Invalidated routing group cache for queryId [%s] due to empty routing group", queryId);
+            return;
+        }
         queryIdRoutingGroupCache.put(queryId, routingGroup);
     }
 
@@ -177,6 +196,15 @@ public void updateClusterStats(List<ClusterStats> stats)
     @Override
     public void setExternalUrlForQueryId(String queryId, String externalUrl)
     {
+        if (isNullOrEmpty(queryId)) {
+            log.debug("Skipping externalUrl cache update for empty queryId");
+            return;
+        }
+        if (isNullOrEmpty(externalUrl)) {
+            queryIdExternalUrlCache.invalidate(queryId);
+            log.debug("Invalidated externalUrl cache for queryId [%s] due to empty externalUrl", queryId);
+            return;
+        }
         queryIdExternalUrlCache.put(queryId, externalUrl);
     }
 
@@ -185,7 +213,7 @@ String findBackendForUnknownQueryId(String queryId)
     {
         String backend;
         backend = queryHistoryManager.getBackendForQueryId(queryId);
-        if (Strings.isNullOrEmpty(backend)) {
+        if (isNullOrEmpty(backend)) {
             log.debug("Unable to find backend mapping for [%s]. Searching for suitable backend", queryId);
             backend = searchAllBackendForQuery(queryId);
         }

gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingManagerExternalUrlCache.java

@@ -21,6 +21,7 @@
 import org.mockito.Mockito;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.mockito.Mockito.when;
 
 @TestInstance(Lifecycle.PER_CLASS)
@@ -106,11 +107,56 @@ void testMultipleQueryIdsWithDifferentExternalUrls()
     void testEmptyStringExternalUrl()
     {
         String queryId = "empty-url-test";
-        String emptyUrl = "";
-        routingManager.setExternalUrlForQueryId(queryId, emptyUrl);
+        String expectedFallbackUrl = "https://fallback-after-empty.example.com";
+
+        when(queryHistoryManager.getExternalUrlForQueryId(queryId)).thenReturn(expectedFallbackUrl);
+
+        routingManager.setExternalUrlForQueryId(queryId, "");
+        String retrievedUrl = routingManager.findExternalUrlForQueryId(queryId);
+
+        assertThat(retrievedUrl).isEqualTo(expectedFallbackUrl);
+        Mockito.verify(queryHistoryManager).getExternalUrlForQueryId(queryId);
+    }
+
+    @Test
+    void testNullExternalUrlDoesNotThrowAndFallsBackToQueryHistory()
+    {
+        String queryId = "null-external-url-query";
+        String expectedFallbackUrl = "https://fallback-after-null.example.com";
+
+        when(queryHistoryManager.getExternalUrlForQueryId(queryId)).thenReturn(expectedFallbackUrl);
+
+        assertThatCode(() -> routingManager.setExternalUrlForQueryId(queryId, null))
+                .doesNotThrowAnyException();
+
         String retrievedUrl = routingManager.findExternalUrlForQueryId(queryId);
 
-        assertThat(retrievedUrl).isEqualTo(emptyUrl);
+        assertThat(retrievedUrl).isEqualTo(expectedFallbackUrl);
+        Mockito.verify(queryHistoryManager).getExternalUrlForQueryId(queryId);
+    }
+
+    @Test
+    void testEmptyExternalUrlDoesNotThrowAndFallsBackToQueryHistory()
+    {
+        String queryId = "empty-external-url-query";
+        String expectedFallbackUrl = "https://fallback-after-empty.example.com";
+
+        when(queryHistoryManager.getExternalUrlForQueryId(queryId)).thenReturn(expectedFallbackUrl);
+
+        assertThatCode(() -> routingManager.setExternalUrlForQueryId(queryId, ""))
+                .doesNotThrowAnyException();
+
+        String retrievedUrl = routingManager.findExternalUrlForQueryId(queryId);
+
+        assertThat(retrievedUrl).isEqualTo(expectedFallbackUrl);
+        Mockito.verify(queryHistoryManager).getExternalUrlForQueryId(queryId);
+    }
+
+    @Test
+    void testNullQueryIdDoesNotThrowForExternalUrlCacheSet()
+    {
+        assertThatCode(() -> routingManager.setExternalUrlForQueryId(null, "https://external-url.example.com"))
+                .doesNotThrowAnyException();
     }
 
     @Test