Prechádzať zdrojové kódy

chore: refactor gateway (#2303)

Kevin Wan 2 rokov pred
rodič
commit
93584c6ca6
2 zmenil súbory, kde vykonal 48 pridanie a 5 odobranie
  1. 47 0
      gateway/internal/eventhandler.go
  2. 1 5
      gateway/server.go

+ 47 - 0
gateway/internal/eventhandler.go

@@ -0,0 +1,47 @@
+package internal
+
+import (
+	"io"
+
+	"github.com/golang/protobuf/jsonpb"
+	"github.com/golang/protobuf/proto"
+	"github.com/jhump/protoreflect/desc"
+	"github.com/zeromicro/go-zero/core/logx"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+type EventHandler struct {
+	Status    *status.Status
+	writer    io.Writer
+	marshaler jsonpb.Marshaler
+}
+
+func NewEventHandler(writer io.Writer, resolver jsonpb.AnyResolver) *EventHandler {
+	return &EventHandler{
+		writer: writer,
+		marshaler: jsonpb.Marshaler{
+			EmitDefaults: true,
+			AnyResolver:  resolver,
+		},
+	}
+}
+
+func (h *EventHandler) OnResolveMethod(_ *desc.MethodDescriptor) {
+}
+
+func (h *EventHandler) OnSendHeaders(_ metadata.MD) {
+}
+
+func (h *EventHandler) OnReceiveHeaders(_ metadata.MD) {
+}
+
+func (h *EventHandler) OnReceiveResponse(message proto.Message) {
+	if err := h.marshaler.Marshal(h.writer, message); err != nil {
+		logx.Error(err)
+	}
+}
+
+func (h *EventHandler) OnReceiveTrailers(status *status.Status, _ metadata.MD) {
+	h.Status = status
+}

+ 1 - 5
gateway/server.go

@@ -120,11 +120,6 @@ func (s *Server) build() error {
 func (s *Server) buildHandler(source grpcurl.DescriptorSource, resolver jsonpb.AnyResolver,
 	cli zrpc.Client, rpcPath string) func(http.ResponseWriter, *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
-		handler := &grpcurl.DefaultEventHandler{
-			Out: w,
-			Formatter: grpcurl.NewJSONFormatter(true,
-				grpcurl.AnyResolverFromDescriptorSource(source)),
-		}
 		parser, err := internal.NewRequestParser(r, resolver)
 		if err != nil {
 			httpx.Error(w, err)
@@ -136,6 +131,7 @@ func (s *Server) buildHandler(source grpcurl.DescriptorSource, resolver jsonpb.A
 		defer can()
 
 		w.Header().Set(httpx.ContentType, httpx.JsonContentType)
+		handler := internal.NewEventHandler(w, resolver)
 		if err := grpcurl.InvokeRPC(ctx, source, cli.Conn(), rpcPath, s.prepareMetadata(r.Header),
 			handler, parser.Next); err != nil {
 			httpx.Error(w, err)