add type-specific decorators

subsquid · Apr 1, 2024 · a67d902 · a67d902
1 parent 8269dcb
commit a67d902
Show file tree

Hide file tree

Showing 15 changed files with 157 additions and 105 deletions.
diff --git a/test/erc20-transfers/src/model/generated/marshal.ts b/test/erc20-transfers/src/model/generated/marshal.ts
@@ -127,36 +127,6 @@ export function nonNull<T>(val: T | undefined | null): T {
 }
 
 
-export const bigintTransformer = {
-    to(x?: bigint) {
-        return x?.toString()
-    },
-    from(s?: string): bigint | undefined {
-        return s == null ? undefined : BigInt(s)
-    }
-}
-
-
-export const floatTransformer = {
-    to(x?: number) {
-        return x?.toString()
-    },
-    from(s?: string): number | undefined {
-        return s == null ? undefined : Number(s)
-    }
-}
-
-
-export const bigdecimalTransformer = {
-    to(x?: any) {
-        return x?.toString()
-    },
-    from(s?: any): any | undefined {
-        return s == null ? undefined : decimal.BigDecimal(s)
-    }
-}
-
-
 export function enumFromJson<E extends object>(json: unknown, enumObject: E): E[keyof E] {
     assert(typeof json == 'string', 'invalid enum value')
     let val = (enumObject as any)[json]

diff --git a/test/erc20-transfers/src/model/generated/transfer.model.ts b/test/erc20-transfers/src/model/generated/transfer.model.ts
@@ -1,5 +1,4 @@
-import {Entity as Entity_, Column as Column_, PrimaryColumn as PrimaryColumn_} from "@subsquid/typeorm-store"
-import * as marshal from "./marshal"
+import {Entity as Entity_, Column as Column_, PrimaryColumn as PrimaryColumn_, IntColumn as IntColumn_, DateTimeColumn as DateTimeColumn_, StringColumn as StringColumn_, BigIntColumn as BigIntColumn_} from "@subsquid/typeorm-store"
 
 @Entity_()
 export class Transfer {
@@ -10,21 +9,21 @@ export class Transfer {
     @PrimaryColumn_()
     id!: string
 
-    @Column_("int4", {nullable: false})
+    @IntColumn_({nullable: false})
     blockNumber!: number
 
-    @Column_("timestamp with time zone", {nullable: false})
+    @DateTimeColumn_({nullable: false})
     timestamp!: Date
 
-    @Column_("text", {nullable: false})
+    @StringColumn_({nullable: false})
     tx!: string
 
-    @Column_("text", {nullable: false})
+    @StringColumn_({nullable: false})
     from!: string
 
-    @Column_("text", {nullable: false})
+    @StringColumn_({nullable: false})
     to!: string
 
-    @Column_("numeric", {transformer: marshal.bigintTransformer, nullable: false})
+    @BigIntColumn_({nullable: false})
     amount!: bigint
 }
diff --git a/typeorm/typeorm-codegen/src/codegen.ts b/typeorm/typeorm-codegen/src/codegen.ts
@@ -59,34 +59,11 @@ export function generateOrmModels(model: Model, dir: OutDir): void {
                         if (key === 'id') {
                             out.line('@PrimaryColumn_()')
                         } else {
+                            const decorator = getDecorator(prop.type.name)
+                            imports.useTypeormStore(decorator)
+
                             addIndexAnnotation(entity, key, imports, out)
-                            switch (prop.type.name) {
-                                case 'BigInt':
-                                    imports.useMarshal()
-                                    out.line(
-                                        `@Column_("${getDbType(prop.type.name)}", {transformer: marshal.bigintTransformer, nullable: ${prop.nullable}})`
-                                    )
-                                    break
-                                case 'BigDecimal':
-                                    imports.useMarshal()
-                                    out.line(
-                                        `@Column_("${getDbType(prop.type.name)}", {transformer: marshal.bigdecimalTransformer, nullable: ${prop.nullable}})`
-                                    )
-                                    break
-                                case 'Float':
-                                    imports.useMarshal()
-                                    out.line(
-                                        `@Column_("${getDbType(prop.type.name)}", {transformer: marshal.floatTransformer, nullable: ${prop.nullable}})`
-                                    )
-                                    break
-                                default:
-                                    out.line(
-                                        `@Column_("${getDbType(prop.type.name)}", {nullable: ${
-                                            prop.nullable
-                                        }})`
-                                    )
-                                    break
-                            }
+                            out.line(`@${decorator}_({nullable: ${prop.nullable}})`)
                         }
                         break
                     case 'enum':
@@ -144,8 +121,12 @@ export function generateOrmModels(model: Model, dir: OutDir): void {
                                 if (scalar == 'BigInt' || scalar == 'BigDecimal') {
                                     throw new Error(`Property ${name}.${key} has unsupported type: can't generate code for native ${scalar} arrays.`)
                                 }
+
+                                const decorator = getDecorator(scalar)
+                                imports.useTypeormStore(decorator)
+
                                 out.line(
-                                    `@Column_("${getDbType(scalar)}", {array: true, nullable: ${prop.nullable}})`
+                                    `@${decorator}_({array: true, nullable: ${prop.nullable}})`
                                 )
                                 break
                             }
@@ -186,26 +167,27 @@ export function generateOrmModels(model: Model, dir: OutDir): void {
         out.write()
     }
 
-    function getDbType(scalar: string): string {
+    function getDecorator(scalar: string): string {
         switch(scalar) {
             case 'ID':
             case 'String':
-                return 'text'
+                return 'StringColumn'
             case 'Int':
-                return 'int4'
+                return 'IntColumn'
             case 'Float':
-                return 'numeric'
+                return 'FloatColumn'
             case 'Boolean':
-                return 'bool'
+                return 'BooleanColumn'
             case 'DateTime':
-                return 'timestamp with time zone'
+                return 'DateTimeColumn'
             case 'BigInt':
+                return 'BigIntColumn'
             case 'BigDecimal':
-                return 'numeric'
+                return 'BigDecimalColumn'
             case 'Bytes':
-                return 'bytea'
+                return 'BytesColumn'
             case 'JSON':
-                return 'jsonb'
+                return 'JSONColumn'
             default:
                 throw unexpectedCase(scalar)
         }

diff --git a/typeorm/typeorm-codegen/src/marshal.ts b/typeorm/typeorm-codegen/src/marshal.ts
@@ -127,36 +127,6 @@ export function nonNull<T>(val: T | undefined | null): T {
 }
 
 
-export const bigintTransformer = {
-    to(x?: bigint) {
-        return x?.toString()
-    },
-    from(s?: string): bigint | undefined {
-        return s == null ? undefined : BigInt(s)
-    }
-}
-
-
-export const floatTransformer = {
-    to(x?: number) {
-        return x?.toString()
-    },
-    from(s?: string): number | undefined {
-        return s == null ? undefined : Number(s)
-    }
-}
-
-
-export const bigdecimalTransformer = {
-    to(x?: any) {
-        return x?.toString()
-    },
-    from(s?: any): any | undefined {
-        return s == null ? undefined : decimal.BigDecimal(s)
-    }
-}
-
-
 export function enumFromJson<E extends object>(json: unknown, enumObject: E): E[keyof E] {
     assert(typeof json == 'string', 'invalid enum value')
     let val = (enumObject as any)[json]

diff --git a/typeorm/typeorm-store/src/decorators/columns/BigDecimalColumn.ts b/typeorm/typeorm-store/src/decorators/columns/BigDecimalColumn.ts
@@ -0,0 +1,9 @@
+import {bigdecimalTransformer} from '../../transformers'
+import {Column} from './Column'
+import {ColumnCommonOptions} from './common'
+
+export type BigDecimalColumnOptions = Pick<ColumnCommonOptions, 'name' | 'unique' | 'nullable' | 'default' | 'comment'>
+
+export function BigDecimalColumn(options?: BigDecimalColumnOptions): PropertyDecorator {
+    return Column('numeric', {...options, transformer: bigdecimalTransformer})
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/BigIntColumn.ts b/typeorm/typeorm-store/src/decorators/columns/BigIntColumn.ts
@@ -0,0 +1,9 @@
+import {bigintTransformer} from '../../transformers'
+import {Column} from './Column'
+import {ColumnCommonOptions} from './common'
+
+export type BigIntColumnOptions = Pick<ColumnCommonOptions, 'name' | 'unique' | 'nullable' | 'default' | 'comment'>
+
+export function BigIntColumn(options?: BigIntColumnOptions): PropertyDecorator {
+    return Column('numeric', {...options, transformer: bigintTransformer})
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/BooleanColumn.ts b/typeorm/typeorm-store/src/decorators/columns/BooleanColumn.ts
@@ -0,0 +1,11 @@
+import {Column} from './Column'
+import {ColumnCommonOptions} from './common'
+
+export type BooleanColumnOptions = Pick<
+    ColumnCommonOptions,
+    'name' | 'unique' | 'nullable' | 'default' | 'comment' | 'array'
+>
+
+export function BooleanColumn(options?: BooleanColumnOptions): PropertyDecorator {
+    return Column('bool', options)
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/BytesColumn.ts b/typeorm/typeorm-store/src/decorators/columns/BytesColumn.ts
@@ -0,0 +1,11 @@
+import {Column} from './Column'
+import {ColumnCommonOptions} from './common'
+
+export type BytesColumnOptions = Pick<
+    ColumnCommonOptions,
+    'name' | 'unique' | 'nullable' | 'default' | 'comment' | 'array'
+>
+
+export function BytesColumn(options?: BytesColumnOptions): PropertyDecorator {
+    return Column('bytea', options)
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/DateTimeColumn.ts b/typeorm/typeorm-store/src/decorators/columns/DateTimeColumn.ts
@@ -0,0 +1,11 @@
+import {Column} from './Column'
+import {ColumnCommonOptions} from './common'
+
+export type DateTimeColumnOptions = Pick<
+    ColumnCommonOptions,
+    'name' | 'unique' | 'nullable' | 'default' | 'comment' | 'array'
+>
+
+export function DateTimeColumn(options?: DateTimeColumnOptions): PropertyDecorator {
+    return Column('timestamp with time zone', options)
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/FloatColumn.ts b/typeorm/typeorm-store/src/decorators/columns/FloatColumn.ts
@@ -0,0 +1,9 @@
+import {floatTransformer} from '../../transformers'
+import {Column} from './Column'
+import {ColumnCommonOptions} from './common'
+
+export type FloatColumnOptions = Pick<ColumnCommonOptions, 'name' | 'unique' | 'nullable' | 'default' | 'comment'>
+
+export function FloatColumn(options?: FloatColumnOptions): PropertyDecorator {
+    return Column('numeric', {...options, transformer: floatTransformer})
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/IntColumn.ts b/typeorm/typeorm-store/src/decorators/columns/IntColumn.ts
@@ -0,0 +1,11 @@
+import {Column} from './Column'
+import {ColumnCommonOptions, ColumnOptions} from './common'
+
+export type IntColumnOptions = Pick<
+    ColumnCommonOptions,
+    'name' | 'unique' | 'nullable' | 'default' | 'comment' | 'array'
+>
+
+export function IntColumn(options?: IntColumnOptions): PropertyDecorator {
+    return Column('int4', options)
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/StringColumn.ts b/typeorm/typeorm-store/src/decorators/columns/StringColumn.ts
@@ -0,0 +1,11 @@
+import {Column} from './Column'
+import {ColumnCommonOptions} from './common'
+
+export type StringColumnOptions = Pick<
+    ColumnCommonOptions,
+    'name' | 'unique' | 'nullable' | 'default' | 'comment' | 'array'
+>
+
+export function StringColumn(options?: StringColumnOptions): PropertyDecorator {
+    return Column('text', options)
+}
diff --git a/typeorm/typeorm-store/src/decorators/columns/index.ts b/typeorm/typeorm-store/src/decorators/columns/index.ts
@@ -1,3 +1,11 @@
 export * from './common'
 export * from './Column'
 export * from './PrimaryColumn'
+export * from './BigDecimalColumn'
+export * from './BigIntColumn'
+export * from './BooleanColumn'
+export * from './BytesColumn'
+export * from './DateTimeColumn'
+export * from './FloatColumn'
+export * from './IntColumn'
+export * from './StringColumn'
diff --git a/typeorm/typeorm-store/src/index.ts b/typeorm/typeorm-store/src/index.ts
@@ -1,3 +1,4 @@
 export * from './database'
 export {EntityClass, FindManyOptions, FindOneOptions, Store} from './store'
 export * from './decorators'
+export * from './transformers'
diff --git a/typeorm/typeorm-store/src/transformers.ts b/typeorm/typeorm-store/src/transformers.ts
@@ -0,0 +1,40 @@
+import {ValueTransformer} from 'typeorm'
+
+export const bigintTransformer: ValueTransformer = {
+    to(x?: bigint) {
+        return x?.toString()
+    },
+    from(s?: string): bigint | undefined {
+        return s == null ? undefined : BigInt(s)
+    },
+}
+
+export const floatTransformer: ValueTransformer = {
+    to(x?: number) {
+        return x?.toString()
+    },
+    from(s?: string): number | undefined {
+        return s == null ? undefined : Number(s)
+    },
+}
+
+const decimal = {
+    get BigDecimal(): any {
+        throw new Error('Package `@subsquid/big-decimal` is not installed')
+    },
+}
+
+try {
+    Object.defineProperty(decimal, 'BigDecimal', {
+        value: require('@subsquid/big-decimal').BigDecimal,
+    })
+} catch (e) {}
+
+export const bigdecimalTransformer: ValueTransformer = {
+    to(x?: any) {
+        return x?.toString()
+    },
+    from(s?: any): any | undefined {
+        return s == null ? undefined : decimal.BigDecimal(s)
+    },
+}