feat(schema): generate enum type helpers (#91)

Author: ymc9

packages/runtime/test/typing/models.ts

@@ -14,3 +14,5 @@ export type Profile = ModelResult<Schema, "Profile">;
 export type Tag = ModelResult<Schema, "Tag">;
 export type Region = ModelResult<Schema, "Region">;
 export type Meta = ModelResult<Schema, "Meta">;
+export const Role = schema.enums.Role;
+export type Role = (typeof Role)[keyof typeof Role];

packages/runtime/test/typing/schema.ts

@@ -37,6 +37,11 @@ export const schema = {
                     unique: true,
                     attributes: [{ name: "@unique" }]
                 },
+                role: {
+                    type: "Role",
+                    attributes: [{ name: "@default", args: [{ name: "value", value: ExpressionUtils.literal("USER") }] }],
+                    default: "USER"
+                },
                 posts: {
                     type: "Post",
                     array: true,
@@ -241,6 +246,12 @@ export const schema = {
             }
         }
     },
+    enums: {
+        Role: {
+            ADMIN: "ADMIN",
+            USER: "USER"
+        }
+    },
     authType: "User",
     plugins: {}
 } as const satisfies SchemaDef;

packages/runtime/test/typing/typing-test.zmodel

@@ -3,12 +3,18 @@ datasource db {
     url = "file:./test.db"
 }
 
+enum Role {
+    ADMIN
+    USER
+}
+
 model User {
     id        Int      @id @default(autoincrement())
     createdAt DateTime @default(now())
     updatedAt DateTime @updatedAt
     name      String
     email     String   @unique
+    role      Role     @default(USER)
     posts     Post[]
     profile   Profile?
     postCount Int      @computed

packages/runtime/test/typing/verify-typing.ts

@@ -1,4 +1,5 @@
 import { ZenStackClient } from '../../dist';
+import { Role } from './models';
 import { schema } from './schema';
 import SQLite from 'better-sqlite3';
 
@@ -26,12 +27,14 @@ async function main() {
     await aggregate();
     await groupBy();
     await queryBuilder();
+    enums();
 }
 
 async function find() {
     const user1 = await client.user.findFirst({
         where: {
             name: 'Alex',
+            role: Role.USER,
         },
     });
     console.log(user1?.name);
@@ -562,4 +565,12 @@ async function queryBuilder() {
     console.log(r.name);
 }
 
+function enums() {
+    const a: Role = 'ADMIN';
+    console.log(a);
+    let b = Role.ADMIN;
+    b = a;
+    console.log(b);
+}
+
 main();

packages/sdk/src/ts-schema-generator.ts

@@ -1021,20 +1021,49 @@ export class TsSchemaGenerator {
         // generate enums
         const enums = model.declarations.filter(isEnum);
         for (const e of enums) {
-            // generate:
-            // export const enum Enum = {
-            //   value1 = 'value1',
-            //   value2 = 'value2',
-            // }
-            let enumDecl = ts.factory.createEnumDeclaration(
+            // generate: export const Enum = schema.enums.Enum;
+            let enumDecl = ts.factory.createVariableStatement(
                 [ts.factory.createModifier(ts.SyntaxKind.ExportKeyword)],
-                e.name,
-                e.fields.map((f) => ts.factory.createEnumMember(f.name, ts.factory.createStringLiteral(f.name))),
+                ts.factory.createVariableDeclarationList(
+                    [
+                        ts.factory.createVariableDeclaration(
+                            e.name,
+                            undefined,
+                            undefined,
+                            ts.factory.createPropertyAccessExpression(
+                                ts.factory.createPropertyAccessExpression(
+                                    ts.factory.createIdentifier('schema'),
+                                    ts.factory.createIdentifier('enums'),
+                                ),
+                                ts.factory.createIdentifier(e.name),
+                            ),
+                        ),
+                    ],
+                    ts.NodeFlags.Const,
+                ),
             );
             if (e.comments.length > 0) {
                 enumDecl = this.generateDocs(enumDecl, e);
             }
             statements.push(enumDecl);
+
+            // generate: export type Enum = (typeof Enum)[keyof typeof Enum];
+            let typeAlias = ts.factory.createTypeAliasDeclaration(
+                [ts.factory.createModifier(ts.SyntaxKind.ExportKeyword)],
+                e.name,
+                undefined,
+                ts.factory.createIndexedAccessTypeNode(
+                    ts.factory.createTypeQueryNode(ts.factory.createIdentifier(e.name)),
+                    ts.factory.createTypeOperatorNode(
+                        ts.SyntaxKind.KeyOfKeyword,
+                        ts.factory.createTypeQueryNode(ts.factory.createIdentifier(e.name)),
+                    ),
+                ),
+            );
+            if (e.comments.length > 0) {
+                typeAlias = this.generateDocs(typeAlias, e);
+            }
+            statements.push(typeAlias);
         }
 
         this.generateBannerComments(statements);
@@ -1047,7 +1076,10 @@ export class TsSchemaGenerator {
         fs.writeFileSync(outputFile, result);
     }
 
-    private generateDocs<T extends ts.TypeAliasDeclaration | ts.EnumDeclaration>(tsDecl: T, decl: DataModel | Enum): T {
+    private generateDocs<T extends ts.TypeAliasDeclaration | ts.VariableStatement>(
+        tsDecl: T,
+        decl: DataModel | Enum,
+    ): T {
         return ts.addSyntheticLeadingComment(
             tsDecl,
             ts.SyntaxKind.MultiLineCommentTrivia,

samples/blog/zenstack/models.ts

@@ -25,7 +25,8 @@ export type Post = ModelResult<Schema, "Post">;
 /**
  * User roles
  */
-export enum Role {
-    ADMIN = "ADMIN",
-    USER = "USER"
-}
+export const Role = schema.enums.Role;
+/**
+ * User roles
+ */
+export type Role = (typeof Role)[keyof typeof Role];