Add destination cluster info to response cookie

gateway-ha/src/main/java/io/trino/gateway/ha/config/HaGatewayConfiguration.java

@@ -41,6 +41,7 @@ public class HaGatewayConfiguration
     private OAuth2GatewayCookieConfiguration oauth2GatewayCookieConfiguration = new OAuth2GatewayCookieConfiguration();
     private GatewayCookieConfiguration gatewayCookieConfiguration = new GatewayCookieConfiguration();
     private List<String> statementPaths = ImmutableList.of(V1_STATEMENT_PATH);
+    private boolean includeClusterHostInResponse;
 
     private RequestAnalyzerConfig requestAnalyzerConfig = new RequestAnalyzerConfig();
 
@@ -97,6 +98,16 @@ public RoutingConfiguration getRouting()
         return routing;
     }
 
+    public void setIncludeClusterHostInResponse(boolean includeClusterHostInResponse)
+    {
+        this.includeClusterHostInResponse = includeClusterHostInResponse;
+    }
+
+    public boolean isIncludeClusterHostInResponse()
+    {
+        return includeClusterHostInResponse;
+    }
+
     public void setRouting(RoutingConfiguration routing)
     {
         this.routing = routing;

gateway-ha/src/main/java/io/trino/gateway/proxyserver/ProxyRequestHandler.java

@@ -40,6 +40,7 @@
 
 import java.io.IOException;
 import java.net.URI;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
@@ -85,6 +86,7 @@ public class ProxyRequestHandler
     private final boolean cookiesEnabled;
     private final boolean addXForwardedHeaders;
     private final List<String> statementPaths;
+    private final boolean includeClusterInfoInResponse;
 
     @Inject
     public ProxyRequestHandler(
@@ -100,6 +102,7 @@ public ProxyRequestHandler(
         asyncTimeout = haGatewayConfiguration.getRouting().getAsyncTimeout();
         addXForwardedHeaders = haGatewayConfiguration.getRouting().isAddXForwardedHeaders();
         statementPaths = haGatewayConfiguration.getStatementPaths();
+        this.includeClusterInfoInResponse = haGatewayConfiguration.isIncludeClusterHostInResponse();
     }
 
     @PreDestroy
@@ -169,13 +172,19 @@ private void performRequest(
 
         FluentFuture<ProxyResponse> future = executeHttp(request);
 
+        List<NewCookie> sessionCookies = new ArrayList<>();
         if (statementPaths.stream().anyMatch(request.getUri().getPath()::startsWith) && request.getMethod().equals(HttpMethod.POST)) {
             future = future.transform(response -> recordBackendForQueryId(request, response), executor);
+            if (includeClusterInfoInResponse) {
+                // add trinoClusterHost to sessionCookies
+                sessionCookies.add(new NewCookie.Builder("trinoClusterHost").value(remoteUri.getHost()).build());
+            }
         }
 
+        ImmutableList<NewCookie> finalSessionCookie = ImmutableList.copyOf(sessionCookies);
         setupAsyncResponse(
                 asyncResponse,
-                future.transform(response -> buildResponse(response, oauth2GatewayCookie), executor)
+                future.transform(response -> buildResponse(response, oauth2GatewayCookie, finalSessionCookie), executor)
                         .catching(ProxyException.class, e -> handleProxyException(request, e), directExecutor()));
     }
 
@@ -206,11 +215,13 @@ private static String getRemoteTarget(URI remoteUri)
         return format("%s://%s", remoteUri.getScheme(), remoteUri.getAuthority());
     }
 
-    private Response buildResponse(ProxyResponse response, ImmutableList<NewCookie> cookie)
+    private Response buildResponse(ProxyResponse response, ImmutableList<NewCookie> cookie, ImmutableList<NewCookie> sessionCookie)
     {
         Response.ResponseBuilder builder = Response.status(response.statusCode()).entity(response.body());
         response.headers().forEach((headerName, value) -> builder.header(headerName.toString(), value));
+
         cookie.forEach(builder::cookie);
+        sessionCookie.forEach(builder::cookie);
         return builder.build();
     }