Derived Loader Fixes (#1340)

* cli: add derived loader

* Fix wrong generated type for derived loaders

* Run prettier

* Refactoring

* Generate derived field getter

* Remove usage of any, add changeset

* chore(dependencies): updated changesets for modified dependencies

* change any to concrete types

* chore(dependencies): updated changesets for modified dependencies

* Create big-hornets-relate.md

---------

Co-authored-by: Gustavo Inacio <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Saihajpreet Singh <[email protected]>

Author: 4 people

.changeset/big-hornets-relate.md

@@ -0,0 +1,5 @@
+---
+"@graphprotocol/graph-ts": minor
+---
+
+export `loadRelated` host function 

.changeset/rotten-lemons-pull.md

@@ -0,0 +1,6 @@
+---
+'@graphprotocol/graph-cli': minor
+---
+
+Add support for codegen for derived field loaders, This adds getters for derived fields defined in
+the schema for entities.

packages/cli/src/codegen/schema.test.ts

@@ -2,15 +2,7 @@ import * as graphql from 'graphql/language';
 import prettier from 'prettier';
 import Schema from '../schema';
 import SchemaCodeGenerator from './schema';
-import {
-  ArrayType,
-  Class,
-  Method,
-  NamedType,
-  NullableType,
-  Param,
-  StaticMethod,
-} from './typescript';
+import { Class, Method, NamedType, NullableType, Param, StaticMethod } from './typescript';
 
 const formatTS = (code: string) => prettier.format(code, { parser: 'typescript', semi: false });
 
@@ -273,14 +265,9 @@ describe('Schema code generator', () => {
           {
             name: 'get wallets',
             params: [],
-            returnType: new NullableType(new ArrayType(new NamedType('string'))),
+            returnType: new NamedType('WalletLoader'),
             body: `
-              let value = this.get('wallets')
-              if (!value || value.kind == ValueKind.NULL) {
-                return null
-              } else {
-                return value.toStringArray()
-              }
+              return new WalletLoader("Account", this.get('id')!.toString(), "wallets") 
             `,
           },
         ],

packages/cli/src/codegen/schema.ts

@@ -1,4 +1,5 @@
 /* eslint-disable unicorn/no-array-for-each */
+import debug from 'debug';
 import Schema from '../schema';
 import * as typesCodegen from './types';
 import * as tsCodegen from './typescript';
@@ -104,13 +105,36 @@ export default class SchemaCodeGenerator {
       .filter(Boolean) as Array<tsCodegen.Class>;
   }
 
+  generateDerivedLoaders() {
+    const fields = (
+      (
+        this.schema.ast.definitions.filter(def =>
+          this._isEntityTypeDefinition(def),
+        ) as ObjectTypeDefinitionNode[]
+      )
+        .flatMap((def: ObjectTypeDefinitionNode) => def.fields)
+        .filter(def => this._isDerivedField(def))
+        .filter(def => def?.type !== undefined) as FieldDefinitionNode[]
+    ).map(def => this._getTypeNameForField(def.type));
+
+    return [...new Set(fields)].map(typeName => {
+      return this._generateDerivedLoader(typeName);
+    });
+  }
+
   _isEntityTypeDefinition(def: DefinitionNode): def is ObjectTypeDefinitionNode {
     return (
       def.kind === 'ObjectTypeDefinition' &&
       def.directives?.find(directive => directive.name.value === 'entity') !== undefined
     );
   }
 
+  _isDerivedField(field: FieldDefinitionNode | undefined): boolean {
+    return (
+      field?.directives?.find((directive: any) => directive.name.value === 'derivedFrom') !==
+      undefined
+    );
+  }
   _isInterfaceDefinition(def: DefinitionNode): def is InterfaceTypeDefinitionNode {
     return def.kind === 'InterfaceTypeDefinition';
   }
@@ -138,6 +162,61 @@ export default class SchemaCodeGenerator {
     return klass;
   }
 
+  _generateDerivedLoader(typeName: string): any {
+    // <field>Loader
+    const klass = tsCodegen.klass(`${typeName}Loader`, { export: true, extends: 'Entity' });
+
+    klass.addMember(tsCodegen.klassMember('_entity', 'string'));
+    klass.addMember(tsCodegen.klassMember('_field', 'string'));
+    klass.addMember(tsCodegen.klassMember('_id', 'string'));
+    // Generate and add a constructor
+    klass.addMethod(
+      tsCodegen.method(
+        'constructor',
+        [
+          tsCodegen.param('entity', 'string'),
+          tsCodegen.param('id', 'string'),
+          tsCodegen.param('field', 'string'),
+        ],
+        undefined,
+        `
+      super();
+      this._entity = entity;
+      this._id = id;
+      this._field = field;
+`,
+      ),
+    );
+
+    // Generate load() method for the Loader
+    klass.addMethod(
+      tsCodegen.method(
+        'load',
+        [],
+        `${typeName}[]`,
+        `
+  let value = store.loadRelated(this._entity, this._id, this._field);
+  return changetype<${typeName}[]>(value);
+  `,
+      ),
+    );
+
+    return klass;
+  }
+
+  _getTypeNameForField(gqlType: TypeNode): string {
+    if (gqlType.kind === 'NonNullType') {
+      return this._getTypeNameForField(gqlType.type);
+    }
+    if (gqlType.kind === 'ListType') {
+      return this._getTypeNameForField(gqlType.type);
+    }
+    if (gqlType.kind === 'NamedType') {
+      return (gqlType as NamedTypeNode).name.value;
+    }
+
+    throw new Error(`Unknown type kind: ${gqlType}`);
+  }
   _generateConstructor(_entityName: string, fields: readonly FieldDefinitionNode[] | undefined) {
     const idField = IdField.fromFields(fields);
     return tsCodegen.method(
@@ -207,6 +286,13 @@ export default class SchemaCodeGenerator {
   }
 
   _generateEntityFieldGetter(_entityDef: ObjectTypeDefinitionNode, fieldDef: FieldDefinitionNode) {
+    const isDerivedField = this._isDerivedField(fieldDef);
+    const codegenDebug = debug('codegen');
+    if (isDerivedField) {
+      codegenDebug(`Generating derived field getter for ${fieldDef.name.value}`);
+      return this._generateDerivedFieldGetter(_entityDef, fieldDef);
+    }
+
     const name = fieldDef.name.value;
     const gqlType = fieldDef.type;
     const fieldValueType = this._valueTypeFromGraphQl(gqlType);
@@ -240,7 +326,59 @@ export default class SchemaCodeGenerator {
       `,
     );
   }
+  _generateDerivedFieldGetter(entityDef: ObjectTypeDefinitionNode, fieldDef: FieldDefinitionNode) {
+    const entityName = entityDef.name.value;
+    const name = fieldDef.name.value;
+    const gqlType = fieldDef.type;
+    const returnType = this._returnTypeForDervied(gqlType);
+    return tsCodegen.method(
+      `get ${name}`,
+      [],
+      returnType,
+      `
+        return new ${returnType}('${entityName}', this.get('id')!.toString(), '${name}')
+      `,
+    );
+  }
+
+  _returnTypeForDervied(gqlType: TypeNode): tsCodegen.NamedType {
+    if (gqlType.kind === 'NonNullType') {
+      return this._returnTypeForDervied(gqlType.type);
+    }
+    if (gqlType.kind === 'ListType') {
+      return this._returnTypeForDervied(gqlType.type);
+    }
+    const type = tsCodegen.namedType(gqlType.name.value + 'Loader');
+    return type;
+  }
+
+  _generatedEntityDerivedFieldGetter(
+    _entityDef: ObjectTypeDefinitionNode,
+    fieldDef: FieldDefinitionNode,
+  ) {
+    const name = fieldDef.name.value;
+    const gqlType = fieldDef.type;
+    const fieldValueType = this._valueTypeFromGraphQl(gqlType);
+    const returnType = this._typeFromGraphQl(gqlType);
+    const isNullable = returnType instanceof tsCodegen.NullableType;
 
+    const getNonNullable = `return ${typesCodegen.valueToAsc('value!', fieldValueType)}`;
+    const getNullable = `if (!value || value.kind == ValueKind.NULL) {
+                          return null
+                        } else {
+                          return ${typesCodegen.valueToAsc('value', fieldValueType)}
+                        }`;
+
+    return tsCodegen.method(
+      `get ${name}`,
+      [],
+      returnType,
+      `
+       let value = this.get('${name}')
+       ${isNullable ? getNullable : getNonNullable}
+      `,
+    );
+  }
   _generateEntityFieldSetter(_entityDef: ObjectTypeDefinitionNode, fieldDef: FieldDefinitionNode) {
     const name = fieldDef.name.value;
     const isDerivedField = !!fieldDef.directives?.find(

packages/cli/src/type-generator.ts

@@ -161,6 +161,7 @@ export default class TypeGenerator {
             GENERATED_FILE_NOTE,
             ...codeGenerator.generateModuleImports(),
             ...codeGenerator.generateTypes(),
+            ...codeGenerator.generateDerivedLoaders(),
           ].join('\n'),
           {
             parser: 'typescript',

packages/ts/index.ts

@@ -26,6 +26,7 @@ export declare namespace store {
   /** If the entity was not created in the block, this function will return null. */
   // Matches the host function https://github.com/graphprotocol/graph-node/blob/9f4a1821146b18f6f49165305e9a8c0795120fad/runtime/wasm/src/module/mod.rs#L1091-L1099
   function get_in_block(entity: string, id: string): Entity | null;
+  function loadRelated(entity: string, id: string, field: string): Array<Entity>;
   function set(entity: string, id: string, data: Entity): void;
   function remove(entity: string, id: string): void;
 }