Add cluster service and SelectCluster method

dataproxy/service/cluster_service.go

@@ -0,0 +1,32 @@
+package service
+
+import (
+	"context"
+
+	"connectrpc.com/connect"
+
+	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster"
+	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster/clusterconnect"
+)
+
+type ClusterService struct {
+	clusterconnect.UnimplementedClusterServiceHandler
+	dataplaneDomain string
+}
+
+func NewClusterService(dataplaneDomain string) *ClusterService {
+	return &ClusterService{
+		dataplaneDomain: dataplaneDomain,
+	}
+}
+
+var _ clusterconnect.ClusterServiceHandler = (*ClusterService)(nil)
+
+func (s *ClusterService) SelectCluster(
+	ctx context.Context,
+	req *connect.Request[cluster.SelectClusterRequest],
+) (*connect.Response[cluster.SelectClusterResponse], error) {
+	return connect.NewResponse(&cluster.SelectClusterResponse{
+		ClusterEndpoint: s.dataplaneDomain,
+	}), nil
+}

dataproxy/setup.go

@@ -8,6 +8,7 @@ import (
 	"github.com/flyteorg/flyte/v2/flytestdlib/app"
 	"github.com/flyteorg/flyte/v2/dataproxy/config"
 	"github.com/flyteorg/flyte/v2/dataproxy/service"
+	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster/clusterconnect"
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/dataproxy/dataproxyconnect"
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/task/taskconnect"
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/trigger/triggerconnect"
@@ -32,6 +33,11 @@ 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)
+	clusterPath, clusterHandler := clusterconnect.NewClusterServiceHandler(clusterSvc)
+	sc.Mux.Handle(clusterPath, clusterHandler)
+	logger.Infof(ctx, "Mounted ClusterService at %s", clusterPath)
+
 	sc.AddReadyCheck(func(r *http.Request) error {
 		baseContainer := sc.DataStore.GetBaseContainerFQN(r.Context())
 		if baseContainer == "" {

flyteidl2/cluster/payload.proto

@@ -0,0 +1,41 @@
+syntax = "proto3";
+
+package flyteidl2.cluster;
+
+import "buf/validate/validate.proto";
+import "flyteidl2/app/app_definition.proto";
+import "flyteidl2/common/identifier.proto";
+import "flyteidl2/task/task_definition.proto";
+
+option go_package = "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster";
+
+message SelectClusterRequest {
+  oneof resource {
+    option (buf.validate.oneof).required = true;
+    flyteidl2.common.OrgIdentifier org_id = 1;
+    flyteidl2.common.ProjectIdentifier project_id = 2;
+    flyteidl2.task.TaskIdentifier task_id = 3;
+    flyteidl2.common.ActionIdentifier action_id = 4;
+    flyteidl2.common.ActionAttemptIdentifier action_attempt_id = 5;
+    flyteidl2.app.Identifier app_id = 6;
+  }
+
+  enum Operation {
+    OPERATION_UNSPECIFIED = 0;
+    OPERATION_CREATE_UPLOAD_LOCATION = 1;
+    OPERATION_UPLOAD_INPUTS = 2;
+    OPERATION_GET_ACTION_DATA = 3;
+    OPERATION_QUERY_RANGE_METRICS = 4;
+    OPERATION_CREATE_DOWNLOAD_LINK = 5;
+    OPERATION_TAIL_LOGS = 6;
+    OPERATION_GET_ACTION_ATTEMPT_METRICS = 7;
+  }
+
+  Operation operation = 8 [(buf.validate.field).enum = {
+    not_in: [0]
+  }];
+}
+
+message SelectClusterResponse {
+  string cluster_endpoint = 1;
+}

flyteidl2/cluster/service.proto

@@ -0,0 +1,12 @@
+syntax = "proto3";
+
+package flyteidl2.cluster;
+
+import "flyteidl2/cluster/payload.proto";
+import "google/api/annotations.proto";
+
+option go_package = "github.com/flyteorg/flyte/v2/gen/go/flyteidl2/cluster";
+
+service ClusterService {
+  rpc SelectCluster(SelectClusterRequest) returns (SelectClusterResponse) {}
+}