feat(orm): support where enum

Author: DoctorFTB

packages/orm/src/client/executor/name-mapper.ts

@@ -1,6 +1,7 @@
 import { invariant } from '@zenstackhq/common-helpers';
 import {
     AliasNode,
+    BinaryOperationNode,
     CaseWhenBuilder,
     ColumnNode,
     ColumnUpdateNode,
@@ -10,20 +11,21 @@ import {
     FromNode,
     IdentifierNode,
     InsertQueryNode,
+    type OperationNode,
     OperationNodeTransformer,
     PrimitiveValueListNode,
     ReferenceNode,
     ReturningNode,
     SelectAllNode,
     SelectionNode,
     SelectQueryNode,
+    type SimpleReferenceExpressionNode,
     TableNode,
     UpdateQueryNode,
     ValueListNode,
     ValueNode,
     ValuesNode,
-    type OperationNode,
-    type SimpleReferenceExpressionNode,
+    WhereNode,
 } from 'kysely';
 import type { EnumDef, EnumField, FieldDef, ModelDef, SchemaDef } from '../../schema';
 import {
@@ -183,6 +185,38 @@ export class QueryNameMapper extends OperationNodeTransformer {
         return ColumnNode.create(mappedName);
     }
 
+    protected override transformWhere(node: WhereNode): WhereNode {
+        if (
+            BinaryOperationNode.is(node.where) &&
+            ReferenceNode.is(node.where.leftOperand) &&
+            ColumnNode.is(node.where.leftOperand.column) &&
+            node.where.leftOperand.table &&
+            TableNode.is(node.where.leftOperand.table) &&
+            // where: { enumColumn: Enum.VALUE }
+            (ValueNode.is(node.where.rightOperand) ||
+                // where: { enumColumn: { in: [Enum.VALUE] } }
+                PrimitiveValueListNode.is(node.where.rightOperand))
+        ) {
+            const tableName = node.where.leftOperand.table.table.identifier.name;
+            const columnNode = node.where.leftOperand.column;
+            const valueNode = node.where.rightOperand;
+
+            let resultValue: OperationNode = valueNode;
+
+            if (ValueNode.is(valueNode)) {
+                resultValue = this.processEnumMappingForValue(tableName, columnNode, valueNode) as OperationNode;
+            } else if (PrimitiveValueListNode.is(valueNode)) {
+                resultValue = PrimitiveValueListNode.create(
+                    this.processEnumMappingForValues(tableName, [columnNode], valueNode.values),
+                );
+            }
+
+            return WhereNode.create(BinaryOperationNode.create(node.where.leftOperand, node.where.operator, resultValue));
+        }
+
+        return super.transformWhere(node);
+    }
+
     protected override transformUpdateQuery(node: UpdateQueryNode) {
         if (!node.table) {
             return super.transformUpdateQuery(node);

tests/e2e/orm/client-api/name-mapping.test.ts

@@ -46,6 +46,16 @@ describe('Name mapping tests', () => {
             user_role: 'role_user',
         });
 
+        rawRead = await db.$qbRaw
+            .selectFrom('users')
+            .where('user_role', '=', 'role_user')
+            .selectAll()
+            .executeTakeFirst();
+        await expect(rawRead).toMatchObject({
+            user_email: 'u1@test.com',
+            user_role: 'role_user',
+        });
+
         await expect(
             db.user.create({
                 data: {
@@ -66,6 +76,15 @@ describe('Name mapping tests', () => {
             user_role: 'MODERATOR',
         });
 
+        rawRead = await db.$qbRaw
+            .selectFrom('users')
+            .where('user_role', '=', 'MODERATOR')
+            .selectAll()
+            .executeTakeFirst();
+        await expect(rawRead).toMatchObject({
+            user_role: 'MODERATOR',
+        });
+
         await expect(
             db.$qb
                 .insertInto('User')
@@ -146,6 +165,52 @@ describe('Name mapping tests', () => {
             posts: [{ title: 'Post1' }],
         });
 
+        await expect(
+            db.user.findFirst({
+                where: { role: 'USER' },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toMatchObject({
+            email: 'u1@test.com',
+            role: 'USER',
+        });
+
+        await expect(
+            db.user.findMany({
+                where: { role: 'USER' },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toEqual([expect.objectContaining({ email: 'u1@test.com', role: 'USER' })]);
+
+        await expect(
+            db.user.findFirst({
+                where: { role: { in: ['USER'] } },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toMatchObject({
+            email: 'u1@test.com',
+            role: 'USER',
+        });
+
+        await expect(
+            db.user.findMany({
+                where: { role: { in: ['USER'] } },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toEqual([expect.objectContaining({ email: 'u1@test.com', role: 'USER' })]);
+
         // select all
         await expect(
             db.user.findFirst({