feat(fxgrpcserver): Provided module

ankorstore · Feb 26, 2024 · fa60d25 · fa60d25
1 parent 5d81abb
commit fa60d25
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 25 deletions.
diff --git a/fxgrpcserver/define.go b/fxgrpcserver/define.go
@@ -2,27 +2,27 @@ package fxgrpcserver
 
 import "google.golang.org/grpc"
 
-type GrpcServiceDefinition interface {
+type GrpcServerServiceDefinition interface {
 	ReturnType() string
 	Description() *grpc.ServiceDesc
 }
 
-type grpcServiceDefinition struct {
+type grpcServerServiceDefinition struct {
 	returnType  string
 	description *grpc.ServiceDesc
 }
 
-func NewGrpcServiceDefinition(returnType string, description *grpc.ServiceDesc) GrpcServiceDefinition {
-	return &grpcServiceDefinition{
+func NewGrpcServiceDefinition(returnType string, description *grpc.ServiceDesc) GrpcServerServiceDefinition {
+	return &grpcServerServiceDefinition{
 		returnType:  returnType,
 		description: description,
 	}
 }
 
-func (d *grpcServiceDefinition) ReturnType() string {
+func (d *grpcServerServiceDefinition) ReturnType() string {
 	return d.returnType
 }
 
-func (d *grpcServiceDefinition) Description() *grpc.ServiceDesc {
+func (d *grpcServerServiceDefinition) Description() *grpc.ServiceDesc {
 	return d.description
 }
diff --git a/fxgrpcserver/module.go b/fxgrpcserver/module.go
@@ -78,6 +78,14 @@ func NewFxGrpcServer(p FxGrpcServerParam) (*grpc.Server, error) {
 	// server interceptors
 	unaryInterceptors, streamInterceptors := createInterceptors(p)
 
+	for _, unaryInterceptor := range p.Registry.ResolveGrpcServerUnaryInterceptors() {
+		unaryInterceptors = append(unaryInterceptors, unaryInterceptor.HandleUnary())
+	}
+
+	for _, streamInterceptor := range p.Registry.ResolveGrpcServerStreamInterceptors() {
+		streamInterceptors = append(streamInterceptors, streamInterceptor.HandleStream())
+	}
+
 	// server options
 	grpcServerOptions := []grpc.ServerOption{
 		grpc.ChainUnaryInterceptor(unaryInterceptors...),
@@ -95,12 +103,12 @@ func NewFxGrpcServer(p FxGrpcServerParam) (*grpc.Server, error) {
 		return nil, err
 	}
 
-	// healthcheck
+	// server healthcheck
 	if p.Config.GetBool("modules.grpc.server.healthcheck.enabled") {
 		grpcServer.RegisterService(&grpc_health_v1.Health_ServiceDesc, grpcserver.NewGrpcHealthCheckService(p.Checker))
 	}
 
-	// registrations
+	// server services
 	resolvedServices, err := p.Registry.ResolveGrpcServerServices()
 	if err != nil {
 		return nil, err

diff --git a/fxgrpcserver/register.go b/fxgrpcserver/register.go
@@ -5,6 +5,26 @@ import (
 	"google.golang.org/grpc"
 )
 
+func AsGrpcServerUnaryInterceptor(constructor any) fx.Option {
+	return fx.Provide(
+		fx.Annotate(
+			constructor,
+			fx.As(new(GrpcServerUnaryInterceptor)),
+			fx.ResultTags(`group:"grpc-server-unary-interceptors"`),
+		),
+	)
+}
+
+func AsGrpcServerStreamInterceptor(constructor any) fx.Option {
+	return fx.Provide(
+		fx.Annotate(
+			constructor,
+			fx.As(new(GrpcServerStreamInterceptor)),
+			fx.ResultTags(`group:"grpc-server-stream-interceptors"`),
+		),
+	)
+}
+
 func AsGrpcServerService(constructor any, description *grpc.ServiceDesc) fx.Option {
 	return fx.Options(
 		fx.Provide(
@@ -17,7 +37,7 @@ func AsGrpcServerService(constructor any, description *grpc.ServiceDesc) fx.Opti
 		fx.Supply(
 			fx.Annotate(
 				NewGrpcServiceDefinition(GetReturnType(constructor), description),
-				fx.As(new(GrpcServiceDefinition)),
+				fx.As(new(GrpcServerServiceDefinition)),
 				fx.ResultTags(`group:"grpc-server-service-definitions"`),
 			),
 		),

diff --git a/fxgrpcserver/registry.go b/fxgrpcserver/registry.go
@@ -7,33 +7,55 @@ import (
 	"go.uber.org/fx"
 )
 
+type GrpcServerUnaryInterceptor interface {
+	HandleUnary() grpc.UnaryServerInterceptor
+}
+
+type GrpcServerStreamInterceptor interface {
+	HandleStream() grpc.StreamServerInterceptor
+}
+
 type GrpcServerRegistry struct {
-	options     []grpc.ServerOption
-	services    []any
-	definitions []GrpcServiceDefinition
+	options            []grpc.ServerOption
+	unaryInterceptors  []GrpcServerUnaryInterceptor
+	streamInterceptors []GrpcServerStreamInterceptor
+	services           []any
+	definitions        []GrpcServerServiceDefinition
 }
 
 type FxGrpcServiceRegistryParam struct {
 	fx.In
-	Options     []grpc.ServerOption     `group:"grpc-server-options"`
-	Services    []any                   `group:"grpc-server-services"`
-	Definitions []GrpcServiceDefinition `group:"grpc-server-service-definitions"`
+	Options            []grpc.ServerOption           `group:"grpc-server-options"`
+	UnaryInterceptors  []GrpcServerUnaryInterceptor  `group:"grpc-server-unary-interceptors"`
+	StreamInterceptors []GrpcServerStreamInterceptor `group:"grpc-server-stream-interceptors"`
+	Services           []any                         `group:"grpc-server-services"`
+	Definitions        []GrpcServerServiceDefinition `group:"grpc-server-service-definitions"`
 }
 
 func NewFxGrpcServerRegistry(p FxGrpcServiceRegistryParam) *GrpcServerRegistry {
 	return &GrpcServerRegistry{
-		options:     p.Options,
-		services:    p.Services,
-		definitions: p.Definitions,
+		options:            p.Options,
+		unaryInterceptors:  p.UnaryInterceptors,
+		streamInterceptors: p.StreamInterceptors,
+		services:           p.Services,
+		definitions:        p.Definitions,
 	}
 }
 
 func (r *GrpcServerRegistry) ResolveGrpcServerOptions() []grpc.ServerOption {
 	return r.options
 }
 
-func (r *GrpcServerRegistry) ResolveGrpcServerServices() ([]*ResolvedGrpcService, error) {
-	var grpcServices []*ResolvedGrpcService
+func (r *GrpcServerRegistry) ResolveGrpcServerUnaryInterceptors() []GrpcServerUnaryInterceptor {
+	return r.unaryInterceptors
+}
+
+func (r *GrpcServerRegistry) ResolveGrpcServerStreamInterceptors() []GrpcServerStreamInterceptor {
+	return r.streamInterceptors
+}
+
+func (r *GrpcServerRegistry) ResolveGrpcServerServices() ([]*ResolvedGrpcServerService, error) {
+	var grpcServices []*ResolvedGrpcServerService
 
 	for _, definition := range r.definitions {
 		implementation, err := r.lookupRegisteredServiceImplementation(definition.ReturnType())

diff --git a/fxgrpcserver/resolve.go b/fxgrpcserver/resolve.go
@@ -2,22 +2,22 @@ package fxgrpcserver
 
 import "google.golang.org/grpc"
 
-type ResolvedGrpcService struct {
+type ResolvedGrpcServerService struct {
 	implementation any
 	description    *grpc.ServiceDesc
 }
 
-func NewResolvedGrpcService(implementation any, description *grpc.ServiceDesc) *ResolvedGrpcService {
-	return &ResolvedGrpcService{
+func NewResolvedGrpcService(implementation any, description *grpc.ServiceDesc) *ResolvedGrpcServerService {
+	return &ResolvedGrpcServerService{
 		implementation: implementation,
 		description:    description,
 	}
 }
 
-func (r *ResolvedGrpcService) Implementation() any {
+func (r *ResolvedGrpcServerService) Implementation() any {
 	return r.implementation
 }
 
-func (r *ResolvedGrpcService) Description() *grpc.ServiceDesc {
+func (r *ResolvedGrpcServerService) Description() *grpc.ServiceDesc {
 	return r.description
 }