refactor(serevr/v2/cometbft): update RegisterQueryHandlers and GRPC…

… queries (#22403)
cosmos · Nov 5, 2024 · 55f7cfc · 55f7cfc
1 parent e9436a6
commit 55f7cfc
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/runtime/v2/manager.go b/runtime/v2/manager.go
@@ -642,7 +642,7 @@ func registerServices[T transaction.Tx](s appmodulev2.AppModule, app *App[T], re
 
 	// if module implements register msg handlers
 	if module, ok := s.(appmodulev2.HasMsgHandlers); ok {
-		wrapper := stfRouterWrapper{stfRouter: app.msgRouterBuilder}
+		wrapper := newStfRouterWrapper(app.msgRouterBuilder)
 		module.RegisterMsgHandlers(&wrapper)
 		if wrapper.error != nil {
 			return fmt.Errorf("unable to register handlers: %w", wrapper.error)
@@ -651,7 +651,7 @@ func registerServices[T transaction.Tx](s appmodulev2.AppModule, app *App[T], re
 
 	// if module implements register query handlers
 	if module, ok := s.(appmodulev2.HasQueryHandlers); ok {
-		wrapper := stfRouterWrapper{stfRouter: app.queryRouterBuilder}
+		wrapper := newStfRouterWrapper(app.queryRouterBuilder)
 		module.RegisterQueryHandlers(&wrapper)
 
 		for path, handler := range wrapper.handlers {
@@ -842,6 +842,13 @@ type stfRouterWrapper struct {
 	handlers map[string]appmodulev2.Handler
 }
 
+func newStfRouterWrapper(stfRouterBuilder *stf.MsgRouterBuilder) stfRouterWrapper {
+	wrapper := stfRouterWrapper{stfRouter: stfRouterBuilder}
+	wrapper.error = nil
+	wrapper.handlers = map[string]appmodulev2.Handler{}
+	return wrapper
+}
+
 func (s *stfRouterWrapper) RegisterHandler(handler appmodulev2.Handler) {
 	req := handler.MakeMsg()
 	requestName := gogoproto.MessageName(req)
@@ -854,7 +861,7 @@ func (s *stfRouterWrapper) RegisterHandler(handler appmodulev2.Handler) {
 	s.error = errors.Join(s.error, err)
 
 	// also make the decoder
-	if s.error == nil {
+	if s.handlers == nil {
 		s.handlers = map[string]appmodulev2.Handler{}
 	}
 	s.handlers[requestName] = handler

diff --git a/server/v2/cometbft/abci.go b/server/v2/cometbft/abci.go
@@ -259,12 +259,15 @@ func (c *Consensus[T]) maybeRunGRPCQuery(ctx context.Context, req *abci.QueryReq
 		return nil, false, err
 	}
 
+	var handlerFullName string
 	md, isGRPC := desc.(protoreflect.MethodDescriptor)
 	if !isGRPC {
-		return nil, false, nil
+		handlerFullName = string(desc.FullName())
+	} else {
+		handlerFullName = string(md.Input().FullName())
 	}
 
-	handler, found := c.queryHandlersMap[string(md.Input().FullName())]
+	handler, found := c.queryHandlersMap[handlerFullName]
 	if !found {
 		return nil, true, fmt.Errorf("no query handler found for %s", req.Path)
 	}
@@ -282,7 +285,7 @@ func (c *Consensus[T]) maybeRunGRPCQuery(ctx context.Context, req *abci.QueryReq
 	}
 
 	resp, err = queryResponse(res, req.Height)
-	return resp, isGRPC, err
+	return resp, true, err
 }
 
 // InitChain implements types.Application.