From 8cb23152beda11bb127c00cbe05b42cf6bf39cce Mon Sep 17 00:00:00 2001
From: Prathamesh Musale <prathamesh.musale0@gmail.com>
Date: Mon, 5 Aug 2024 16:08:44 +0530
Subject: [PATCH] Generate GQL schema types with arguments on plural fields

---
 packages/codegen/src/schema.ts | 45 +++++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 6 deletions(-)

diff --git a/packages/codegen/src/schema.ts b/packages/codegen/src/schema.ts
index f5800905..b03e716a 100644
--- a/packages/codegen/src/schema.ts
+++ b/packages/codegen/src/schema.ts
@@ -189,6 +189,8 @@ export class Schema {
   }
 
   _addSubgraphSchemaQueries (subgraphTypeDefs: ReadonlyArray<DefinitionNode>): void {
+    const subgraphTypeArgsMap = new Map<string, { [key: string]: any }>();
+
     for (const subgraphTypeDef of subgraphTypeDefs) {
       // Filtering out enums.
       if (subgraphTypeDef.kind !== 'ObjectTypeDefinition') {
@@ -303,21 +305,30 @@ export class Schema {
       let pluralQueryName = pluralize(queryName);
       pluralQueryName = (pluralQueryName === queryName) ? `${pluralQueryName}s` : pluralQueryName;
 
+      const queryArgs = {
+        where: `${subgraphType}_filter`,
+        orderBy: subgraphTypeOrderByEnum,
+        orderDirection: ORDER_DIRECTION,
+        first: { type: GraphQLInt, defaultValue: 100 },
+        skip: { type: GraphQLInt, defaultValue: 0 }
+      };
+
       queryObject[pluralQueryName] = {
         // Get type composer object for return type from the schema composer.
         type: this._composer.getAnyTC(subgraphType).NonNull.List.NonNull,
         args: {
           block: BLOCK_HEIGHT,
-          where: `${subgraphType}_filter`,
-          orderBy: subgraphTypeOrderByEnum,
-          orderDirection: ORDER_DIRECTION,
-          first: { type: GraphQLInt, defaultValue: 100 },
-          skip: { type: GraphQLInt, defaultValue: 0 }
+          ...queryArgs
         }
       };
-
       this._composer.Query.addFields(queryObject);
+
+      // Save the args for this type in a map (type -> args) for further usage.
+      subgraphTypeArgsMap.set(subgraphType, queryArgs);
     }
+
+    // Add args on plural fields for subgraph types.
+    this._addSubgraphPluralFieldArgs(subgraphTypeDefs, subgraphTypeArgsMap);
   }
 
   _getDetailsForSubgraphField (fieldType: ComposeOutputType<any>): {
@@ -381,6 +392,28 @@ export class Schema {
       });
   }
 
+  _addSubgraphPluralFieldArgs (subgraphTypeDefs: ReadonlyArray<DefinitionNode>, subgraphTypeArgsMap: Map<string, { [key: string]: any }>): void {
+    for (const subgraphTypeDef of subgraphTypeDefs) {
+      // Filtering out enums.
+      if (subgraphTypeDef.kind !== 'ObjectTypeDefinition') {
+        continue;
+      }
+
+      const subgraphType = subgraphTypeDef.name.value;
+      const subgraphTypeComposer = this._composer.getOTC(subgraphType);
+
+      // Process each field on the type.
+      Object.entries(subgraphTypeComposer.getFields()).forEach(([fieldName, field]) => {
+        const { isArray, entityType } = this._getDetailsForSubgraphField(field.type);
+
+        // Set args if it's a plural field of some entity type.
+        if (entityType && isArray) {
+          subgraphTypeComposer.setFieldArgs(fieldName, subgraphTypeArgsMap.get(entityType) || {});
+        }
+      });
+    }
+  }
+
   /**
    * Adds basic types to the schema and typemapping.
    */