Add Additional gRPC service handler hook

Signed-off-by: Katrina Rogan <[email protected]>

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

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()

plugins/registry.go

@@ -14,6 +14,7 @@ const (
 	PluginIDUnaryServiceMiddleware PluginID = "UnaryServiceMiddleware"
 	PluginIDPreRedirectHook        PluginID = "PreRedirectHook"
 	PluginIDLogoutHook             PluginID = "LogoutHook"
+	PluginIDAdditionalGRPCService  PluginID = "AdditionalGRPCService"
 )
 
 type AtomicRegistry struct {