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 {