From e9aa3428ce4deab38028e1cbe2ec127fb40121c8 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <katroganGH@gmail.com>
Date: Tue, 3 Oct 2023 11:10:05 -0700
Subject: [PATCH] Add Additional gRPC service handler hook

Signed-off-by: Katrina Rogan <katroganGH@gmail.com>
---
 pkg/common/service.go | 11 +++++++++++
 pkg/server/service.go |  7 +++++++
 plugins/registry.go   |  1 +
 3 files changed, 19 insertions(+)
 create mode 100644 pkg/common/service.go

diff --git a/pkg/common/service.go b/pkg/common/service.go
new file mode 100644
index 000000000..29c3fb6e2
--- /dev/null
+++ b/pkg/common/service.go
@@ -0,0 +1,11 @@
+package common
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+)
+
+// RegisterAdditionalGRPCService is the interface for the plugin hook for additional GRPC service handlers which
+// should be also served on the flyteadmin gRPC server.
+type RegisterAdditionalGRPCService = func(ctx context.Context, grpcServer *grpc.Server) error
diff --git a/pkg/server/service.go b/pkg/server/service.go
index b060c3ec7..d9f7847d3 100644
--- a/pkg/server/service.go
+++ b/pkg/server/service.go
@@ -128,6 +128,13 @@ func newGRPCServer(ctx context.Context, pluginRegistry *plugins.Registry, cfg *c
 	pluginRegistry.RegisterDefault(plugins.PluginIDDataProxy, dataProxySvc)
 	service.RegisterDataProxyServiceServer(grpcServer, plugins.Get[service.DataProxyServiceServer](pluginRegistry, plugins.PluginIDDataProxy))
 
+	additionalService := plugins.Get[common.RegisterAdditionalGRPCService](pluginRegistry, plugins.PluginIDAdditionalGRPCService)
+	if additionalService != nil {
+		if err := additionalService(ctx, grpcServer); err != nil {
+			return nil, err
+		}
+	}
+
 	service.RegisterSignalServiceServer(grpcServer, rpc.NewSignalServer(ctx, configuration, scope.NewSubScope("signal")))
 
 	healthServer := health.NewServer()
diff --git a/plugins/registry.go b/plugins/registry.go
index 90389008b..92644b136 100644
--- a/plugins/registry.go
+++ b/plugins/registry.go
@@ -14,6 +14,7 @@ const (
 	PluginIDUnaryServiceMiddleware PluginID = "UnaryServiceMiddleware"
 	PluginIDPreRedirectHook        PluginID = "PreRedirectHook"
 	PluginIDLogoutHook             PluginID = "LogoutHook"
+	PluginIDAdditionalGRPCService  PluginID = "AdditionalGRPCService"
 )
 
 type AtomicRegistry struct {