feat(x/accounts): Add schema caching feature and corresponding test case

x/accounts/query_server.go

@@ -8,17 +8,27 @@ import (
 	v1 "cosmossdk.io/x/accounts/v1"
 )
 
-var _ v1.QueryServer = queryServer{}
+var _ v1.QueryServer = &queryServer{}
 
+// NewQueryServer initializes a new instance of QueryServer.
+// It precalculates and stores schemas for efficient schema retrieval.
 func NewQueryServer(k Keeper) v1.QueryServer {
-	return &queryServer{k}
+	// Pre-calculate schemas for efficient retrieval.
+	schemas := v1.MakeAccountsSchemas(k.accounts)
+	return &queryServer{
+		k:           k,
+		schemaCache: make(map[string]*v1.SchemaResponse, len(schemas)), // Initialize with precomputed length
+		schemas:     schemas,                                           // Store precalculated schemas.
+	}
 }
 
 type queryServer struct {
-	k Keeper
+	k           Keeper
+	schemaCache map[string]*v1.SchemaResponse
+	schemas     map[string]*v1.SchemaResponse // Stores precalculated schemas.
 }
 
-func (q queryServer) AccountQuery(ctx context.Context, request *v1.AccountQueryRequest) (*v1.AccountQueryResponse, error) {
+func (q *queryServer) AccountQuery(ctx context.Context, request *v1.AccountQueryRequest) (*v1.AccountQueryResponse, error) {
 	// get target addr
 	targetAddr, err := q.k.addressCodec.StringToBytes(request.Target)
 	if err != nil {
@@ -47,18 +57,19 @@ func (q queryServer) AccountQuery(ctx context.Context, request *v1.AccountQueryR
 	}, nil
 }
 
-func (q queryServer) Schema(_ context.Context, request *v1.SchemaRequest) (*v1.SchemaResponse, error) {
-	// TODO: maybe we should cache this, considering accounts types are not
-	// added on the fly as the chain is running.
-	schemas := v1.MakeAccountsSchemas(q.k.accounts)
-	schema, ok := schemas[request.AccountType]
+// Schema retrieves the schema for a given account type.
+// It checks the cache first and computes the schema if not found.
+func (q *queryServer) Schema(_ context.Context, request *v1.SchemaRequest) (*v1.SchemaResponse, error) {
+	// Fetch schema from precalculated schemas.
+	schema, ok := q.schemas[request.AccountType]
 	if !ok {
 		return nil, fmt.Errorf("%w: %s", errAccountTypeNotFound, request.AccountType)
 	}
+
 	return schema, nil
 }
 
-func (q queryServer) AccountType(ctx context.Context, request *v1.AccountTypeRequest) (*v1.AccountTypeResponse, error) {
+func (q *queryServer) AccountType(ctx context.Context, request *v1.AccountTypeRequest) (*v1.AccountTypeResponse, error) {
 	addr, err := q.k.addressCodec.StringToBytes(request.Address)
 	if err != nil {
 		return nil, err
@@ -72,7 +83,7 @@ func (q queryServer) AccountType(ctx context.Context, request *v1.AccountTypeReq
 	}, nil
 }
 
-func (q queryServer) AccountNumber(ctx context.Context, request *v1.AccountNumberRequest) (*v1.AccountNumberResponse, error) {
+func (q *queryServer) AccountNumber(ctx context.Context, request *v1.AccountNumberRequest) (*v1.AccountNumberResponse, error) {
 	addr, err := q.k.addressCodec.StringToBytes(request.Address)
 	if err != nil {
 		return nil, err

x/accounts/query_server_test.go

@@ -58,4 +58,20 @@ func TestQueryServer(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "test", typ.AccountType)
 	})
+
+	t.Run("schema caching", func(t *testing.T) {
+		// Request schema once
+		schemaReq := &v1.SchemaRequest{AccountType: "test"}
+		schemaResp1, err := qs.Schema(ctx, schemaReq)
+		require.NoError(t, err)
+		require.NotNil(t, schemaResp1)
+
+		// Request schema again
+		schemaResp2, err := qs.Schema(ctx, schemaReq)
+		require.NoError(t, err)
+		require.NotNil(t, schemaResp2)
+
+		// Check if both responses are the same (cached)
+		require.Equal(t, schemaResp1, schemaResp2)
+	})
 }