Optimize Optional usage and improve code consistency in gateway-ha module

Author: Akanksha-kedia

gateway-ha/src/main/java/io/trino/gateway/ha/handler/RoutingTargetHandler.java

@@ -84,7 +84,7 @@ public RoutingTargetResponse resolveRouting(HttpServletRequest request)
             return new RoutingTargetResponse(
                     new RoutingDestination(routingGroup, cluster, buildUriWithNewCluster(cluster, request)),
                     request);
-        }).orElse(getRoutingTargetResponse(request));
+        }).orElseGet(() -> getRoutingTargetResponse(request));
         logRewrite(routingTargetResponse.routingDestination().clusterHost(), request);
         return routingTargetResponse;
     }

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

@@ -25,8 +25,8 @@
 import io.trino.gateway.ha.persistence.dao.SelectorsDao;
 import jakarta.annotation.Nullable;
 
-import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
@@ -90,13 +90,10 @@ public ResourceGroupsDetail updateResourceGroup(ResourceGroupsDetail resourceGro
             @Nullable String routingGroupDatabase)
     {
         ResourceGroupsDao dao = getResourceGroupsDao(routingGroupDatabase);
-        ResourceGroups model = dao.findFirstById(resourceGroup.getResourceGroupId());
-        if (model == null) {
-            dao.create(resourceGroup);
-        }
-        else {
-            dao.update(resourceGroup);
-        }
+        Optional.ofNullable(dao.findFirstById(resourceGroup.getResourceGroupId()))
+                .ifPresentOrElse(
+                        _ -> dao.update(resourceGroup),
+                        () -> dao.create(resourceGroup));
         return resourceGroup;
     }
 
@@ -153,13 +150,10 @@ public SelectorsDetail updateSelector(SelectorsDetail selector, SelectorsDetail
             @Nullable String routingGroupDatabase)
     {
         SelectorsDao dao = getSelectorsDao(routingGroupDatabase);
-        Selectors model = dao.findFirst(selector);
-        if (model == null) {
-            dao.insert(updatedSelector);
-        }
-        else {
-            dao.update(selector, updatedSelector);
-        }
+        Optional.ofNullable(dao.findFirst(selector))
+                .ifPresentOrElse(
+                        _ -> dao.update(selector, updatedSelector),
+                        () -> dao.insert(updatedSelector));
         return updatedSelector;
     }
 
@@ -213,14 +207,10 @@ public GlobalPropertiesDetail updateGlobalProperty(GlobalPropertiesDetail global
             @Nullable String routingGroupDatabase)
     {
         ResourceGroupsGlobalPropertiesDao dao = getDao(routingGroupDatabase);
-        ResourceGroupsGlobalProperties model = dao.findFirstByName(globalProperty.getName());
-
-        if (model == null) {
-            dao.insert(globalProperty.getName(), globalProperty.getValue());
-        }
-        else {
-            dao.update(globalProperty.getName(), globalProperty.getValue());
-        }
+        Optional.ofNullable(dao.findFirstByName(globalProperty.getName()))
+                .ifPresentOrElse(
+                        _ -> dao.update(globalProperty.getName(), globalProperty.getValue()),
+                        () -> dao.insert(globalProperty.getName(), globalProperty.getValue()));
         return globalProperty;
     }
 
@@ -263,8 +253,9 @@ public List<ExactSelectorsDetail> readExactMatchSourceSelector()
     public ExactSelectorsDetail getExactMatchSourceSelector(
             ExactSelectorsDetail exactSelectorDetail)
     {
-        ExactMatchSourceSelectors exactSelector = exactMatchSourceSelectorsDao.findFirst(exactSelectorDetail);
-        return upcastExactSelectors(exactSelector);
+        return Optional.ofNullable(exactMatchSourceSelectorsDao.findFirst(exactSelectorDetail))
+                .map(HaResourceGroupsManager::upcastExactSelectors)
+                .orElse(null);
     }
 
     private SelectorsDao getSelectorsDao(@Nullable String routingGroupDatabase)
@@ -284,15 +275,14 @@ private ResourceGroupsGlobalPropertiesDao getDao(@Nullable String routingGroupDa
 
     private static List<GlobalPropertiesDetail> upcast(List<ResourceGroupsGlobalProperties> globalPropertiesList)
     {
-        List<GlobalPropertiesDetail> globalProperties = new ArrayList<>();
-        for (ResourceGroupsGlobalProperties dao : globalPropertiesList) {
-            GlobalPropertiesDetail globalPropertiesDetail = new GlobalPropertiesDetail();
-            globalPropertiesDetail.setName(dao.name());
-            globalPropertiesDetail.setValue(dao.value());
-
-            globalProperties.add(globalPropertiesDetail);
-        }
-        return globalProperties;
+        return globalPropertiesList.stream()
+                .map(dao -> {
+                    GlobalPropertiesDetail detail = new GlobalPropertiesDetail();
+                    detail.setName(dao.name());
+                    detail.setValue(dao.value());
+                    return detail;
+                })
+                .collect(toImmutableList());
     }
 
     private static ExactSelectorsDetail upcastExactSelectors(ExactMatchSourceSelectors exactMatchSourceSelector)

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

@@ -206,9 +206,9 @@ private void updateLocalStats(LocalStats stats, String user)
         // so that they can be used for next queries to route
         // We assume that if a user has queued queries then newly arriving queries
         // for that user would also be queued
-        int count = stats.userQueuedCount() == null ? 0 : stats.userQueuedCount().getOrDefault(user, 0);
-        if (count > 0) {
-            stats.userQueuedCount().put(user, count + 1);
+        Map<String, Integer> queuedCounts = stats.userQueuedCount();
+        if (queuedCounts != null && queuedCounts.getOrDefault(user, 0) > 0) {
+            queuedCounts.merge(user, 1, Integer::sum);
             return;
         }
         // Else the we assume that the query would be running
@@ -233,7 +233,7 @@ public String provideAdhocCluster(String user)
     public String provideClusterForRoutingGroup(String routingGroup, String user)
     {
         return getBackendForRoutingGroup(routingGroup, user)
-                .orElse(provideAdhocCluster(user));
+                .orElseGet(() -> provideAdhocCluster(user));
     }
 
     @Override

gateway-ha/src/main/java/io/trino/gateway/ha/security/LbFilter.java

@@ -61,7 +61,7 @@ public void filter(final ContainerRequestContext requestContext)
             String idToken = Optional
                     .ofNullable(requestContext.getCookies().get(SessionCookie.OAUTH_ID_TOKEN))
                     .map(Cookie::getValue)
-                    .orElse(getToken(requestContext.getHeaders().getFirst(HttpHeaders.AUTHORIZATION)));
+                    .orElseGet(() -> getToken(requestContext.getHeaders().getFirst(HttpHeaders.AUTHORIZATION)));
 
             LbPrincipal principal = idTokenAuthenticator.authenticate(idToken)
                     .orElseThrow(() -> new AuthenticationException("Authentication error"));

gateway-ha/src/main/java/io/trino/gateway/ha/security/LbOAuthManager.java

@@ -127,7 +127,7 @@ public Response exchangeCodeForToken(String code, String nonce, String redirectL
             return buildUnauthorizedResponse();
         }
         return Response.status(FOUND)
-                .location(oauthConfig.getRedirectWebUrl().orElse(URI.create(redirectLocation)))
+                .location(oauthConfig.getRedirectWebUrl().orElseGet(() -> URI.create(redirectLocation)))
                 .cookie(SessionCookie.getTokenCookie(idToken))
                 .cookie(OidcCookie.delete())
                 .build();