Handle enum type fields in subgraph entities

cerc-io · Dec 22, 2021 · 4e86485 · 4e86485
1 parent f30704f
commit 4e86485
Show file tree

Hide file tree

Showing 5 changed files with 129 additions and 19 deletions.
diff --git a/packages/codegen/src/entity.ts b/packages/codegen/src/entity.ts
@@ -241,7 +241,7 @@ export class Entity {
       });
 
       // Add subgraph entity specific columns.
-      entityObject = this._addSubgraphColumns(entityObject, def);
+      entityObject = this._addSubgraphColumns(subgraphTypeDefs, entityObject, def);
 
       // Add bigintTransformer column option if required.
       this._addBigIntTransformerOption(entityObject);
@@ -308,7 +308,7 @@ export class Entity {
     });
   }
 
-  _addSubgraphColumns (entityObject: any, def: any): any {
+  _addSubgraphColumns (subgraphTypeDefs: any, entityObject: any, def: any): any {
     def.fields.forEach((field: any) => {
       let name = field.name.value;
 
@@ -324,7 +324,8 @@ export class Entity {
 
       const columnObject: any = {
         name,
-        columnOptions: []
+        columnOptions: [],
+        columnType: 'Column'
       };
 
       const { typeName, array, nullable } = this._getFieldType(field.type);
@@ -350,29 +351,58 @@ export class Entity {
 
       const pgType = getPgForTs(tsType);
 
-      // If basic type: create a column. If unknown: create a relation.
+      // If basic type: create a column.
       if (pgType) {
-        columnObject.columnType = 'Column';
         columnObject.pgType = pgType;
       } else {
-        columnObject.columnType = 'ManyToOne';
-        columnObject.lhs = '()';
-        columnObject.rhs = tsType;
-
-        entityObject.imports[0].toImport.add('ManyToOne');
-
-        // Check if type import already added.
-        const importObject = entityObject.imports.find((element: any) => {
-          return element.from === `./${tsType}`;
-        });
+        if (subgraphTypeDefs.some((typeDef: any) => typeDef.kind === 'EnumTypeDefinition' && typeDef.name.value === typeName)) {
+          // Create enum type column.
+
+          const entityImport = entityObject.imports.find(({ from }: any) => from === '../types');
+
+          if (!entityImport) {
+            entityObject.imports.push(
+              {
+                toImport: new Set([typeName]),
+                from: '../types'
+              }
+            );
+          } else {
+            entityImport.toImport.add(typeName);
+          }
 
-        if (!importObject) {
-          entityObject.imports.push(
+          columnObject.columnOptions.push(
             {
-              toImport: new Set([tsType]),
-              from: `./${tsType}`
+              option: 'type',
+              value: "'enum'"
+            },
+            {
+              option: 'enum',
+              value: typeName
             }
           );
+        } else {
+          // Create a relation.
+
+          columnObject.columnType = 'ManyToOne';
+          columnObject.lhs = '()';
+          columnObject.rhs = tsType;
+
+          entityObject.imports[0].toImport.add('ManyToOne');
+
+          // Check if type import already added.
+          const importObject = entityObject.imports.find((element: any) => {
+            return element.from === `./${tsType}`;
+          });
+
+          if (!importObject) {
+            entityObject.imports.push(
+              {
+                toImport: new Set([tsType]),
+                from: `./${tsType}`
+              }
+            );
+          }
         }
       }
 

diff --git a/packages/codegen/src/generate-code.ts b/packages/codegen/src/generate-code.ts
@@ -263,6 +263,11 @@ function generateWatcher (contractStrings: string[], visitor: Visitor, argv: any
     : process.stdout;
   exportFill(outStream);
 
+  outStream = outputDir
+    ? fs.createWriteStream(path.join(outputDir, 'src/types.ts'))
+    : process.stdout;
+  visitor.exportTypes(outStream);
+
   let rcOutStream, ignoreOutStream;
 
   if (outputDir) {

diff --git a/packages/codegen/src/templates/types-template.handlebars b/packages/codegen/src/templates/types-template.handlebars
@@ -0,0 +1,11 @@
+//
+// Copyright 2021 Vulcanize, Inc.
+//
+
+{{#each types as | type |}}
+export enum {{type.name}} {
+  {{#each type.values as | value |}}
+  {{value}} = '{{value}}',
+  {{/each}}
+}
+{{/each}}
diff --git a/packages/codegen/src/types.ts b/packages/codegen/src/types.ts
@@ -0,0 +1,52 @@
+//
+// Copyright 2021 Vulcanize, Inc.
+//
+
+import fs from 'fs';
+import path from 'path';
+import Handlebars from 'handlebars';
+import { Writable } from 'stream';
+
+const TEMPLATE_FILE = './templates/types-template.handlebars';
+
+export class Types {
+  _types: Array<any>;
+  _templateString: string;
+
+  constructor () {
+    this._types = [];
+    this._templateString = fs.readFileSync(path.resolve(__dirname, TEMPLATE_FILE)).toString();
+  }
+
+  /**
+   * Writes the generated types files from a template to a stream.
+   * @param outStream A writable output stream to write the types file to.
+   */
+  exportTypes (outStream: Writable): void {
+    const template = Handlebars.compile(this._templateString);
+    const obj = {
+      types: this._types
+    };
+    const database = template(obj);
+    outStream.write(database);
+  }
+
+  addSubgraphTypes (subgraphSchemaDocument: any): void {
+    const subgraphTypeDefs = subgraphSchemaDocument.definitions;
+
+    subgraphTypeDefs.forEach((def: any) => {
+      if (def.kind !== 'EnumTypeDefinition') {
+        return;
+      }
+
+      console.log('def.values', def.values);
+
+      const typeObject: any = {
+        name: def.name.value,
+        values: def.values.map((value: any) => value.name.value)
+      };
+
+      this._types.push(typeObject);
+    });
+  }
+}
diff --git a/packages/codegen/src/visitor.ts b/packages/codegen/src/visitor.ts
@@ -14,6 +14,7 @@ import { Reset } from './reset';
 import { Param } from './utils/types';
 import { MODE_ETH_CALL, MODE_STORAGE } from './utils/constants';
 import { parseSubgraphSchema } from './utils/subgraph';
+import { Types } from './types';
 
 export class Visitor {
   _schema: Schema;
@@ -23,6 +24,7 @@ export class Visitor {
   _database: Database;
   _client: Client;
   _reset: Reset;
+  _types: Types;
 
   constructor () {
     this._schema = new Schema();
@@ -32,6 +34,7 @@ export class Visitor {
     this._database = new Database();
     this._client = new Client();
     this._reset = new Reset();
+    this._types = new Types();
   }
 
   /**
@@ -131,6 +134,7 @@ export class Visitor {
     const subgraphSchemaDocument = parseSubgraphSchema(subgraphPath);
 
     this._schema.addSubgraphSchema(subgraphSchemaDocument);
+    this._types.addSubgraphTypes(subgraphSchemaDocument);
     this._entity.addSubgraphEntities(subgraphSchemaDocument);
     this._resolvers.addSubgraphResolvers(subgraphSchemaDocument);
     this._reset.addSubgraphEntities(subgraphSchemaDocument);
@@ -197,4 +201,12 @@ export class Visitor {
   exportReset (resetOutStream: Writable, resetJQOutStream: Writable, resetStateOutStream: Writable): void {
     this._reset.exportReset(resetOutStream, resetJQOutStream, resetStateOutStream);
   }
+
+  /**
+   * Writes the types file generated from a template to a stream.
+   * @param outStream A writable output stream to write the database file to.
+   */
+  exportTypes (outStream: Writable): void {
+    this._types.exportTypes(outStream);
+  }
 }