rag/internal/cmd/rpc.go

package cmd

import (
	grpc2 "framework_v2/internal/middleware/grpc"
	"framework_v2/internal/providers/jwks"
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/auth"
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"

	"net"
)

var rpcCommand = &cobra.Command{
	Use: "rpc",

	Run: func(cmd *cobra.Command, args []string) {
		jwks.InitJwksRefresh()
		StartSpiderService()
	},
}

func StartSpiderService() {
	if config.ListenAddr.GRPC == "" {
		config.ListenAddr.GRPC = "0.0.0.0:8081"
	}

	lis, err := net.Listen("tcp", config.ListenAddr.GRPC)
	if err != nil {
		panic("failed to listen: " + err.Error())
	}
	logger.Info("Server listening at " + config.ListenAddr.GRPC)

	var opts = []grpc.ServerOption{
		grpc.ChainUnaryInterceptor(
			logging.UnaryServerInterceptor(grpc2.ZapLogInterceptor()),
			auth.UnaryServerInterceptor(grpc2.JwtAuth),
		),
		grpc.ChainStreamInterceptor(
			logging.StreamServerInterceptor(grpc2.ZapLogInterceptor()),
			auth.StreamServerInterceptor(grpc2.JwtAuth),
		),
	}
	grpcServer := grpc.NewServer(opts...)

	reflection.Register(grpcServer)

	// 如需实现，需要手动实现这里
	//library.RegisterLibraryServiceServer(grpcServer, &services.LibraryService{})
	//document.RegisterDocumentServiceServer(grpcServer, &services.DocumentService{})

	err = grpcServer.Serve(lis)
	if err != nil {
		panic(err)
		return
	}

}