feat: Add loadInBlock function for entities as part of codegen (#1306)

* feat: Add loadInBlock function for entities as part of codegen

* fix tests

* match with host fn def

Author: saihaj

.changeset/slow-emus-share.md

@@ -0,0 +1,5 @@
+---
+'@graphprotocol/graph-cli': minor
+---
+
+Add loadInBlock function for entities as part of codegen

.changeset/warm-goats-agree.md

@@ -0,0 +1,5 @@
+---
+'@graphprotocol/graph-ts': minor
+---
+
+expose Host's `get_in_block` function

examples/ethereum-basic-event-handlers/generated/schema.ts

@@ -29,6 +29,12 @@ export class NewGravatar extends Entity {
     }
   }
 
+  static loadInBlock(id: string): NewGravatar | null {
+    return changetype<NewGravatar | null>(
+      store.get_in_block("NewGravatar", id)
+    );
+  }
+
   static load(id: string): NewGravatar | null {
     return changetype<NewGravatar | null>(store.get("NewGravatar", id));
   }
@@ -104,6 +110,12 @@ export class UpdatedGravatar extends Entity {
     }
   }
 
+  static loadInBlock(id: string): UpdatedGravatar | null {
+    return changetype<UpdatedGravatar | null>(
+      store.get_in_block("UpdatedGravatar", id)
+    );
+  }
+
   static load(id: string): UpdatedGravatar | null {
     return changetype<UpdatedGravatar | null>(store.get("UpdatedGravatar", id));
   }

examples/example-subgraph/src/types/schema.ts

@@ -29,6 +29,10 @@ export class ExampleEntity extends Entity {
     }
   }
 
+  static loadInBlock(id: string): ExampleEntity | null {
+    return changetype<ExampleEntity | null>(store.get_in_block('ExampleEntity', id));
+  }
+
   static load(id: string): ExampleEntity | null {
     return changetype<ExampleEntity | null>(store.get('ExampleEntity', id));
   }
@@ -417,6 +421,10 @@ export class OtherEntity extends Entity {
     }
   }
 
+  static loadInBlock(id: string): OtherEntity | null {
+    return changetype<OtherEntity | null>(store.get_in_block('OtherEntity', id));
+  }
+
   static load(id: string): OtherEntity | null {
     return changetype<OtherEntity | null>(store.get('OtherEntity', id));
   }

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

@@ -127,6 +127,15 @@ describe('Schema code generator', () => {
               }
             `,
           },
+          {
+            name: 'loadInBlock',
+            static: true,
+            params: [new Param('id', new NamedType('string'))],
+            returnType: new NullableType(new NamedType('Account')),
+            body: `
+              return changetype<Account | null>(store.get_in_block('Account', id))
+            `,
+          },
           {
             name: 'load',
             static: true,
@@ -307,6 +316,15 @@ describe('Schema code generator', () => {
               }
             `,
           },
+          {
+            name: 'loadInBlock',
+            static: true,
+            params: [new Param('id', new NamedType('string'))],
+            returnType: new NullableType(new NamedType('Wallet')),
+            body: `
+              return changetype<Wallet | null>(store.get_in_block('Wallet', id))
+            `,
+          },
           {
             name: 'load',
             static: true,
@@ -426,6 +444,16 @@ describe('Schema code generator', () => {
             "          store.set('Task', id.toBytes().toHexString(), this)\n" +
             '        }',
         },
+        {
+          name: 'loadInBlock',
+          static: true,
+          params: [new Param('id', new NamedType('Bytes'))],
+          returnType: new NullableType(new NamedType('Task')),
+          body:
+            '\n' +
+            "        return changetype<Task | null>(store.get_in_block('Task', id.toHexString()))\n" +
+            '        ',
+        },
         {
           name: 'load',
           static: true,

packages/cli/src/codegen/schema.ts

@@ -171,6 +171,15 @@ export default class SchemaCodeGenerator {
         }`,
       ),
 
+      tsCodegen.staticMethod(
+        'loadInBlock',
+        [tsCodegen.param('id', tsCodegen.namedType(idField.typeName()))],
+        tsCodegen.nullableType(tsCodegen.namedType(entityName)),
+        `
+        return changetype<${entityName} | null>(store.get_in_block('${entityName}', ${idField.tsToString()}))
+        `,
+      ),
+
       tsCodegen.staticMethod(
         'load',
         [tsCodegen.param('id', tsCodegen.namedType(idField.typeName()))],

packages/cli/tests/cli/validation/call-handler-with-tuple/generated/schema.ts

@@ -29,13 +29,21 @@ export class MyEntity extends Entity {
     }
   }
 
+  static loadInBlock(id: string): MyEntity | null {
+    return changetype<MyEntity | null>(store.get_in_block("MyEntity", id));
+  }
+
   static load(id: string): MyEntity | null {
     return changetype<MyEntity | null>(store.get("MyEntity", id));
   }
 
   get id(): string {
     let value = this.get("id");
-    return value!.toString();
+    if (!value || value.kind == ValueKind.NULL) {
+      throw new Error("Cannot return null for a required field.");
+    } else {
+      return value.toString();
+    }
   }
 
   set id(value: string) {

packages/cli/tests/cli/validation/example-values-found/generated/schema.ts

@@ -29,13 +29,21 @@ export class MyEntity extends Entity {
     }
   }
 
+  static loadInBlock(id: string): MyEntity | null {
+    return changetype<MyEntity | null>(store.get_in_block("MyEntity", id));
+  }
+
   static load(id: string): MyEntity | null {
     return changetype<MyEntity | null>(store.get("MyEntity", id));
   }
 
   get id(): string {
     let value = this.get("id");
-    return value!.toString();
+    if (!value || value.kind == ValueKind.NULL) {
+      throw new Error("Cannot return null for a required field.");
+    } else {
+      return value.toString();
+    }
   }
 
   set id(value: string) {

packages/ts/index.ts

@@ -23,6 +23,9 @@ export * from './common/value';
  */
 export declare namespace store {
   function get(entity: string, id: string): Entity | null;
+  /** 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 set(entity: string, id: string, data: Entity): void;
   function remove(entity: string, id: string): void;
 }