From 35ec9b89d2e380024a32712eb49be230b8a774ea Mon Sep 17 00:00:00 2001 From: Haytham Abuelfutuh Date: Fri, 17 Apr 2026 10:35:19 -0700 Subject: [PATCH] Use request host in SelectCluster Signed-off-by: Haytham Abuelfutuh --- dataproxy/service/cluster_service.go | 12 ++++++------ dataproxy/setup.go | 7 ++++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/dataproxy/service/cluster_service.go b/dataproxy/service/cluster_service.go index d0d9ddf078..0c1a684c0d 100644 --- a/dataproxy/service/cluster_service.go +++ b/dataproxy/service/cluster_service.go @@ -4,6 +4,7 @@ import ( "context" "connectrpc.com/connect" + "github.com/flyteorg/flyte/v2/flytestdlib/logger" "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster" "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster/clusterconnect" @@ -11,13 +12,10 @@ import ( type ClusterService struct { clusterconnect.UnimplementedClusterServiceHandler - dataplaneDomain string } -func NewClusterService(dataplaneDomain string) *ClusterService { - return &ClusterService{ - dataplaneDomain: dataplaneDomain, - } +func NewClusterService() *ClusterService { + return &ClusterService{} } var _ clusterconnect.ClusterServiceHandler = (*ClusterService)(nil) @@ -26,7 +24,9 @@ func (s *ClusterService) SelectCluster( ctx context.Context, req *connect.Request[cluster.SelectClusterRequest], ) (*connect.Response[cluster.SelectClusterResponse], error) { + requestHost := req.Header().Get("Host") + logger.Debugf(ctx, "Request Host: %s", requestHost) return connect.NewResponse(&cluster.SelectClusterResponse{ - ClusterEndpoint: s.dataplaneDomain, + ClusterEndpoint: requestHost, }), nil } diff --git a/dataproxy/setup.go b/dataproxy/setup.go index 4652d997fe..adbe891a17 100644 --- a/dataproxy/setup.go +++ b/dataproxy/setup.go @@ -3,12 +3,13 @@ package dataproxy import ( "context" "fmt" - "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/workflow/workflowconnect" "net/http" + "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/workflow/workflowconnect" + "github.com/flyteorg/flyte/v2/dataproxy/config" - "github.com/flyteorg/flyte/v2/dataproxy/service" "github.com/flyteorg/flyte/v2/dataproxy/logs" + "github.com/flyteorg/flyte/v2/dataproxy/service" "github.com/flyteorg/flyte/v2/flytestdlib/app" "github.com/flyteorg/flyte/v2/flytestdlib/logger" "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster/clusterconnect" @@ -42,7 +43,7 @@ func Setup(ctx context.Context, sc *app.SetupContext) error { sc.Mux.Handle(path, handler) logger.Infof(ctx, "Mounted DataProxyService at %s", path) - clusterSvc := service.NewClusterService(baseURL) + clusterSvc := service.NewClusterService() clusterPath, clusterHandler := clusterconnect.NewClusterServiceHandler(clusterSvc) sc.Mux.Handle(clusterPath, clusterHandler) logger.Infof(ctx, "Mounted ClusterService at %s", clusterPath)