feat(orm): support where enum

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,20 @@ import {
     FromNode,
     IdentifierNode,
     InsertQueryNode,
+    type OperationNode,
     OperationNodeTransformer,
     PrimitiveValueListNode,
     ReferenceNode,
     ReturningNode,
     SelectAllNode,
     SelectionNode,
     SelectQueryNode,
+    type SimpleReferenceExpressionNode,
     TableNode,
     UpdateQueryNode,
     ValueListNode,
     ValueNode,
     ValuesNode,
-    type OperationNode,
-    type SimpleReferenceExpressionNode,
 } from 'kysely';
 import type { EnumDef, EnumField, FieldDef, ModelDef, SchemaDef } from '../../schema';
 import {
@@ -183,6 +184,56 @@ export class QueryNameMapper extends OperationNodeTransformer {
         return ColumnNode.create(mappedName);
     }
 
+    protected override transformBinaryOperation(node: BinaryOperationNode) {
+        // transform enum name mapping for enum values used inside binary operations
+        //   1. simple value: column = EnumValue
+        //   2. list value: column IN [EnumValue, EnumValue2]
+
+        // note: Kysely only allows column ref on the left side of a binary operation
+
+        if (
+            ReferenceNode.is(node.leftOperand) &&
+            ColumnNode.is(node.leftOperand.column) &&
+            (ValueNode.is(node.rightOperand) || PrimitiveValueListNode.is(node.rightOperand))
+        ) {
+            const columnNode = node.leftOperand.column;
+
+            // resolve field from scope in case it's not directly qualified with a table name
+            const resolvedScope = this.resolveFieldFromScopes(
+                columnNode.column.name,
+                node.leftOperand.table?.table.identifier.name,
+            );
+
+            if (resolvedScope?.model) {
+                const valueNode = node.rightOperand;
+                let resultValue: OperationNode = valueNode;
+
+                if (ValueNode.is(valueNode)) {
+                    resultValue = this.processEnumMappingForValue(
+                        resolvedScope.model,
+                        columnNode,
+                        valueNode,
+                    ) as OperationNode;
+                } else if (PrimitiveValueListNode.is(valueNode)) {
+                    resultValue = PrimitiveValueListNode.create(
+                        this.processEnumMappingForValues(
+                            resolvedScope.model,
+                            valueNode.values.map(() => columnNode),
+                            valueNode.values,
+                        ),
+                    );
+                }
+
+                return super.transformBinaryOperation({
+                    ...node,
+                    rightOperand: resultValue,
+                });
+            }
+        }
+
+        return super.transformBinaryOperation(node);
+    }
+
     protected override transformUpdateQuery(node: UpdateQueryNode) {
         if (!node.table) {
             return super.transformUpdateQuery(node);

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

@@ -8,10 +8,10 @@ describe('Name mapping tests', () => {
     let db: ClientContract<SchemaType>;
 
     beforeEach(async () => {
-        db = (await createTestClient(schema, {
+        db = await createTestClient(schema, {
             usePrismaPush: true,
             schemaFile: path.join(__dirname, '../schemas/name-mapping/schema.zmodel'),
-        })) as any;
+        });
     });
 
     afterEach(async () => {
@@ -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: '[email protected]',
+            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,64 @@ describe('Name mapping tests', () => {
             posts: [{ title: 'Post1' }],
         });
 
+        await expect(
+            db.user.findFirst({
+                where: { role: 'USER' },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toMatchObject({
+            email: '[email protected]',
+            role: 'USER',
+        });
+
+        await expect(
+            db.user.findMany({
+                where: { role: 'USER' },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toEqual([expect.objectContaining({ email: '[email protected]', role: 'USER' })]);
+
+        await expect(
+            db.user.findFirst({
+                where: { role: { in: ['USER'] } },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toMatchObject({
+            email: '[email protected]',
+            role: 'USER',
+        });
+
+        await expect(
+            db.user.findMany({
+                where: { role: { in: ['USER'] } },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toEqual([expect.objectContaining({ email: '[email protected]', role: 'USER' })]);
+
+        await expect(
+            db.user.findMany({
+                where: {
+                    AND: [{ role: { in: ['USER'] } }, { role: { in: ['USER'] } }, { OR: [{ role: { in: ['USER'] } }] }],
+                },
+                select: {
+                    email: true,
+                    role: true,
+                },
+            }),
+        ).resolves.toEqual([expect.objectContaining({ email: '[email protected]', role: 'USER' })]);
+
         // select all
         await expect(
             db.user.findFirst({
@@ -185,6 +262,38 @@ describe('Name mapping tests', () => {
             role: 'USER',
         });
 
+        // name mapping for enum value in where clause, with unqualified column name
+        await expect(
+            db.$qb.selectFrom('User').select(['User.email', 'User.role']).where('role', '=', 'USER').executeTakeFirst(),
+        ).resolves.toMatchObject({
+            email: '[email protected]',
+            role: 'USER',
+        });
+
+        // name mapping for enum value in simple where clause, with qualified column name
+        await expect(
+            db.$qb
+                .selectFrom('User as u')
+                .select(['u.email', 'u.role'])
+                .where('u.role', '=', 'USER')
+                .executeTakeFirst(),
+        ).resolves.toMatchObject({
+            email: '[email protected]',
+            role: 'USER',
+        });
+
+        // enum value in list
+        await expect(
+            db.$qb
+                .selectFrom('User')
+                .select(['User.email', 'User.role'])
+                .where('role', 'in', ['USER', 'ADMIN'])
+                .executeTakeFirst(),
+        ).resolves.toMatchObject({
+            email: '[email protected]',
+            role: 'USER',
+        });
+
         await expect(
             db.$qb
                 .selectFrom('User')