From 7c41775c7f8d7e3bd0bb210657e784435b28e61f Mon Sep 17 00:00:00 2001
From: CatsJuice <wp1298554944@hotmail.com>
Date: Tue, 31 Dec 2024 03:44:58 +0000
Subject: [PATCH] feat(infra): orm f.enum() support (#9323)

---
 .../infra/src/orm/core/__tests__/schema.spec.ts | 12 ++++++++++++
 packages/common/infra/src/orm/core/schema.ts    | 11 ++++++++---
 .../infra/src/orm/core/validators/data.ts       | 17 +++++++++++++++--
 .../core/src/modules/db/schema/schema.ts        |  2 +-
 4 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/packages/common/infra/src/orm/core/__tests__/schema.spec.ts b/packages/common/infra/src/orm/core/__tests__/schema.spec.ts
index d31e5580dc995..ad05d8a1492e5 100644
--- a/packages/common/infra/src/orm/core/__tests__/schema.spec.ts
+++ b/packages/common/infra/src/orm/core/__tests__/schema.spec.ts
@@ -61,6 +61,7 @@ describe('Entity validations', () => {
       id: f.string().primaryKey().default(nanoid),
       name: f.string(),
       color: f.string(),
+      status: f.enum('active', 'inactive').optional(),
     },
   });
 
@@ -130,4 +131,15 @@ describe('Entity validations', () => {
       expect(tag.info).toBe(null);
     });
   });
+
+  test('should throw when trying to create entity with invalid enum value', () => {
+    const client = createTagsClient();
+
+    expect(() =>
+      // @ts-expect-error test
+      client.tags.create({ name: 'test', status: 'not-active' })
+    ).toThrow(
+      "[Table(tags)]: Field 'status' value 'not-active' is not valid. Expected one of [active, inactive]."
+    );
+  });
 });
diff --git a/packages/common/infra/src/orm/core/schema.ts b/packages/common/infra/src/orm/core/schema.ts
index af28ef6ecad7f..e49f7ca6212a2 100644
--- a/packages/common/infra/src/orm/core/schema.ts
+++ b/packages/common/infra/src/orm/core/schema.ts
@@ -1,10 +1,11 @@
-export type FieldType = 'string' | 'number' | 'boolean' | 'json';
+export type FieldType = 'string' | 'number' | 'boolean' | 'json' | 'enum';
 
 export interface FieldSchema<Type = unknown> {
   type: FieldType;
   optional: boolean;
   isPrimaryKey: boolean;
   default?: () => Type;
+  values?: Type[];
 }
 
 export type TableSchema = Record<string, FieldSchema>;
@@ -28,10 +29,12 @@ export class FieldSchemaBuilder<
     optional: false,
     isPrimaryKey: false,
     default: undefined,
+    values: undefined,
   };
 
-  constructor(type: FieldType) {
+  constructor(type: FieldType, values?: string[]) {
     this.schema.type = type;
+    this.schema.values = values;
   }
 
   optional() {
@@ -56,7 +59,9 @@ export const f = {
   number: () => new FieldSchemaBuilder<number>('number'),
   boolean: () => new FieldSchemaBuilder<boolean>('boolean'),
   json: <T = any>() => new FieldSchemaBuilder<T>('json'),
-} satisfies Record<FieldType, () => FieldSchemaBuilder<any>>;
+  enum: <T extends string>(...values: T[]) =>
+    new FieldSchemaBuilder<T>('enum', values),
+} as const;
 
 export const t = {
   document: <T extends TableSchemaBuilder>(schema: T) => {
diff --git a/packages/common/infra/src/orm/core/validators/data.ts b/packages/common/infra/src/orm/core/validators/data.ts
index 1944ae7116567..5e549a2539bba 100644
--- a/packages/common/infra/src/orm/core/validators/data.ts
+++ b/packages/common/infra/src/orm/core/validators/data.ts
@@ -70,7 +70,14 @@ export const dataValidators = {
           }
 
           const typeGet = inputType(val);
-          if (!typeMatches(field.type, typeGet)) {
+
+          if (field.type === 'enum') {
+            if (!field.values?.includes(val)) {
+              throw new Error(
+                `[Table(${table.name})]: Field '${key}' value '${val}' is not valid. Expected one of [${field.values?.join(', ')}].`
+              );
+            }
+          } else if (!typeMatches(field.type, typeGet)) {
             throw new Error(
               `[Table(${table.name})]: Field '${key}' type mismatch. Expected ${field.type} got ${typeGet}.`
             );
@@ -103,7 +110,13 @@ export const dataValidators = {
           }
 
           const typeGet = inputType(val);
-          if (!typeMatches(field.type, typeGet)) {
+          if (field.type === 'enum') {
+            if (!field.values?.includes(val)) {
+              throw new Error(
+                `[Table(${table.name})]: Field '${key}' value '${val}' is not valid. Expected one of [${field.values?.join(', ')}].`
+              );
+            }
+          } else if (!typeMatches(field.type, typeGet)) {
             throw new Error(
               `[Table(${table.name})]: Field '${key}' type mismatch. Expected type '${field.type}' but got '${typeGet}'.`
             );
diff --git a/packages/frontend/core/src/modules/db/schema/schema.ts b/packages/frontend/core/src/modules/db/schema/schema.ts
index 857cf20172bd6..c2ad8ccb23c3b 100644
--- a/packages/frontend/core/src/modules/db/schema/schema.ts
+++ b/packages/frontend/core/src/modules/db/schema/schema.ts
@@ -26,7 +26,7 @@ export const AFFiNE_WORKSPACE_DB_SCHEMA = {
     id: f.string().primaryKey().optional().default(nanoid),
     name: f.string().optional(),
     type: f.string(),
-    show: f.string().optional(),
+    show: f.enum('always-show', 'always-hide', 'hide-when-empty').optional(),
     index: f.string().optional(),
     icon: f.string().optional(),
     additionalData: f.json().optional(),