fix: enhance auth type

packages/orm/src/client/contract.ts

@@ -1,5 +1,13 @@
 import type Decimal from 'decimal.js';
-import { type GetModels, type IsDelegateModel, type ProcedureDef, type SchemaDef } from '../schema';
+import {
+    type FieldIsArray,
+    type GetModels,
+    type IsDelegateModel,
+    type ProcedureDef,
+    type RelationFields,
+    type RelationFieldType,
+    type SchemaDef,
+} from '../schema';
 import type { AnyKysely } from '../utils/kysely-utils';
 import type { OrUndefinedIf, Simplify, UnwrapTuplePromises } from '../utils/type-utils';
 import type { TRANSACTION_UNSUPPORTED_METHODS } from './constants';
@@ -12,6 +20,7 @@ import type {
     CreateArgs,
     CreateManyAndReturnArgs,
     CreateManyArgs,
+    DefaultModelResult,
     DeleteArgs,
     DeleteManyArgs,
     FindFirstArgs,
@@ -810,11 +819,23 @@ export type ModelOperations<Schema extends SchemaDef, Model extends GetModels<Sc
 
 //#region Supporting types
 
+/**
+ * Type for auth context that includes both scalar and relation fields.
+ * Relations are recursively included to allow nested auth data like { user: { profile: { ... } } }
+ */
+type AuthModelType<Schema extends SchemaDef, Model extends GetModels<Schema>> = Partial<
+    DefaultModelResult<Schema, Model>
+> & {
+    [Key in RelationFields<Schema, Model>]?: FieldIsArray<Schema, Model, Key> extends true
+        ? AuthModelType<Schema, RelationFieldType<Schema, Model, Key>>[]
+        : AuthModelType<Schema, RelationFieldType<Schema, Model, Key>>;
+};
+
 export type AuthType<Schema extends SchemaDef> =
     string extends GetModels<Schema>
         ? Record<string, unknown>
         : Schema['authType'] extends GetModels<Schema>
-          ? Partial<ModelResult<Schema, Schema['authType']>>
+          ? AuthModelType<Schema, Schema['authType']>
           : never;
 
 //#endregion

packages/orm/src/client/crud-types.ts

@@ -47,7 +47,7 @@ import type { ToKyselySchema } from './query-builder';
 
 //#region Query results
 
-type DefaultModelResult<
+export type DefaultModelResult<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>,
     Omit = undefined,

tests/regression/test/issue-422/input.ts

@@ -0,0 +1,70 @@
+//////////////////////////////////////////////////////////////////////////////////////////////
+// DO NOT MODIFY THIS FILE                                                                  //
+// This file is automatically generated by ZenStack CLI and should not be manually updated. //
+//////////////////////////////////////////////////////////////////////////////////////////////
+
+/* eslint-disable */
+
+import { type SchemaType as $Schema } from "./schema";
+import type { FindManyArgs as $FindManyArgs, FindUniqueArgs as $FindUniqueArgs, FindFirstArgs as $FindFirstArgs, CreateArgs as $CreateArgs, CreateManyArgs as $CreateManyArgs, CreateManyAndReturnArgs as $CreateManyAndReturnArgs, UpdateArgs as $UpdateArgs, UpdateManyArgs as $UpdateManyArgs, UpdateManyAndReturnArgs as $UpdateManyAndReturnArgs, UpsertArgs as $UpsertArgs, DeleteArgs as $DeleteArgs, DeleteManyArgs as $DeleteManyArgs, CountArgs as $CountArgs, AggregateArgs as $AggregateArgs, GroupByArgs as $GroupByArgs, WhereInput as $WhereInput, SelectInput as $SelectInput, IncludeInput as $IncludeInput, OmitInput as $OmitInput } from "@zenstackhq/orm";
+import type { SimplifiedModelResult as $SimplifiedModelResult, SelectIncludeOmit as $SelectIncludeOmit } from "@zenstackhq/orm";
+export type SessionFindManyArgs = $FindManyArgs<$Schema, "Session">;
+export type SessionFindUniqueArgs = $FindUniqueArgs<$Schema, "Session">;
+export type SessionFindFirstArgs = $FindFirstArgs<$Schema, "Session">;
+export type SessionCreateArgs = $CreateArgs<$Schema, "Session">;
+export type SessionCreateManyArgs = $CreateManyArgs<$Schema, "Session">;
+export type SessionCreateManyAndReturnArgs = $CreateManyAndReturnArgs<$Schema, "Session">;
+export type SessionUpdateArgs = $UpdateArgs<$Schema, "Session">;
+export type SessionUpdateManyArgs = $UpdateManyArgs<$Schema, "Session">;
+export type SessionUpdateManyAndReturnArgs = $UpdateManyAndReturnArgs<$Schema, "Session">;
+export type SessionUpsertArgs = $UpsertArgs<$Schema, "Session">;
+export type SessionDeleteArgs = $DeleteArgs<$Schema, "Session">;
+export type SessionDeleteManyArgs = $DeleteManyArgs<$Schema, "Session">;
+export type SessionCountArgs = $CountArgs<$Schema, "Session">;
+export type SessionAggregateArgs = $AggregateArgs<$Schema, "Session">;
+export type SessionGroupByArgs = $GroupByArgs<$Schema, "Session">;
+export type SessionWhereInput = $WhereInput<$Schema, "Session">;
+export type SessionSelect = $SelectInput<$Schema, "Session">;
+export type SessionInclude = $IncludeInput<$Schema, "Session">;
+export type SessionOmit = $OmitInput<$Schema, "Session">;
+export type SessionGetPayload<Args extends $SelectIncludeOmit<$Schema, "Session", true>> = $SimplifiedModelResult<$Schema, "Session", Args>;
+export type UserFindManyArgs = $FindManyArgs<$Schema, "User">;
+export type UserFindUniqueArgs = $FindUniqueArgs<$Schema, "User">;
+export type UserFindFirstArgs = $FindFirstArgs<$Schema, "User">;
+export type UserCreateArgs = $CreateArgs<$Schema, "User">;
+export type UserCreateManyArgs = $CreateManyArgs<$Schema, "User">;
+export type UserCreateManyAndReturnArgs = $CreateManyAndReturnArgs<$Schema, "User">;
+export type UserUpdateArgs = $UpdateArgs<$Schema, "User">;
+export type UserUpdateManyArgs = $UpdateManyArgs<$Schema, "User">;
+export type UserUpdateManyAndReturnArgs = $UpdateManyAndReturnArgs<$Schema, "User">;
+export type UserUpsertArgs = $UpsertArgs<$Schema, "User">;
+export type UserDeleteArgs = $DeleteArgs<$Schema, "User">;
+export type UserDeleteManyArgs = $DeleteManyArgs<$Schema, "User">;
+export type UserCountArgs = $CountArgs<$Schema, "User">;
+export type UserAggregateArgs = $AggregateArgs<$Schema, "User">;
+export type UserGroupByArgs = $GroupByArgs<$Schema, "User">;
+export type UserWhereInput = $WhereInput<$Schema, "User">;
+export type UserSelect = $SelectInput<$Schema, "User">;
+export type UserInclude = $IncludeInput<$Schema, "User">;
+export type UserOmit = $OmitInput<$Schema, "User">;
+export type UserGetPayload<Args extends $SelectIncludeOmit<$Schema, "User", true>> = $SimplifiedModelResult<$Schema, "User", Args>;
+export type ProfileFindManyArgs = $FindManyArgs<$Schema, "Profile">;
+export type ProfileFindUniqueArgs = $FindUniqueArgs<$Schema, "Profile">;
+export type ProfileFindFirstArgs = $FindFirstArgs<$Schema, "Profile">;
+export type ProfileCreateArgs = $CreateArgs<$Schema, "Profile">;
+export type ProfileCreateManyArgs = $CreateManyArgs<$Schema, "Profile">;
+export type ProfileCreateManyAndReturnArgs = $CreateManyAndReturnArgs<$Schema, "Profile">;
+export type ProfileUpdateArgs = $UpdateArgs<$Schema, "Profile">;
+export type ProfileUpdateManyArgs = $UpdateManyArgs<$Schema, "Profile">;
+export type ProfileUpdateManyAndReturnArgs = $UpdateManyAndReturnArgs<$Schema, "Profile">;
+export type ProfileUpsertArgs = $UpsertArgs<$Schema, "Profile">;
+export type ProfileDeleteArgs = $DeleteArgs<$Schema, "Profile">;
+export type ProfileDeleteManyArgs = $DeleteManyArgs<$Schema, "Profile">;
+export type ProfileCountArgs = $CountArgs<$Schema, "Profile">;
+export type ProfileAggregateArgs = $AggregateArgs<$Schema, "Profile">;
+export type ProfileGroupByArgs = $GroupByArgs<$Schema, "Profile">;
+export type ProfileWhereInput = $WhereInput<$Schema, "Profile">;
+export type ProfileSelect = $SelectInput<$Schema, "Profile">;
+export type ProfileInclude = $IncludeInput<$Schema, "Profile">;
+export type ProfileOmit = $OmitInput<$Schema, "Profile">;
+export type ProfileGetPayload<Args extends $SelectIncludeOmit<$Schema, "Profile", true>> = $SimplifiedModelResult<$Schema, "Profile", Args>;

tests/regression/test/issue-422/models.ts

@@ -0,0 +1,12 @@
+//////////////////////////////////////////////////////////////////////////////////////////////
+// DO NOT MODIFY THIS FILE                                                                  //
+// This file is automatically generated by ZenStack CLI and should not be manually updated. //
+//////////////////////////////////////////////////////////////////////////////////////////////
+
+/* eslint-disable */
+
+import { type SchemaType as $Schema } from "./schema";
+import { type ModelResult as $ModelResult } from "@zenstackhq/orm";
+export type Session = $ModelResult<$Schema, "Session">;
+export type User = $ModelResult<$Schema, "User">;
+export type Profile = $ModelResult<$Schema, "Profile">;

tests/regression/test/issue-422/regression.test.ts

@@ -0,0 +1,18 @@
+import { createTestClient } from '@zenstackhq/testtools';
+import { describe, it } from 'vitest';
+import { schema } from './schema';
+
+describe('Issue 422 regression tests', () => {
+    it('should infer correct auth type', async () => {
+        const db = await createTestClient(schema);
+
+        // all fields optional
+        db.$setAuth({ id: 'session1' });
+
+        // relations are allowed
+        db.$setAuth({ id: 'user1', user: { id: 'user1' } });
+
+        // nested relations are allowed
+        db.$setAuth({ id: 'user1', user: { id: 'user1', profile: { name: 'User1' } } });
+    });
+});

tests/regression/test/issue-422/schema.ts

@@ -0,0 +1,122 @@
+//////////////////////////////////////////////////////////////////////////////////////////////
+// DO NOT MODIFY THIS FILE                                                                  //
+// This file is automatically generated by ZenStack CLI and should not be manually updated. //
+//////////////////////////////////////////////////////////////////////////////////////////////
+
+/* eslint-disable */
+
+import { type SchemaDef, ExpressionUtils } from "@zenstackhq/orm/schema";
+const _schema = {
+    provider: {
+        type: "sqlite"
+    },
+    models: {
+        Session: {
+            name: "Session",
+            fields: {
+                id: {
+                    name: "id",
+                    type: "String",
+                    id: true,
+                    attributes: [{ name: "@id" }]
+                },
+                token: {
+                    name: "token",
+                    type: "String"
+                },
+                userId: {
+                    name: "userId",
+                    type: "String",
+                    foreignKeyFor: [
+                        "user"
+                    ]
+                },
+                user: {
+                    name: "user",
+                    type: "User",
+                    attributes: [{ name: "@relation", args: [{ name: "fields", value: ExpressionUtils.array([ExpressionUtils.field("userId")]) }, { name: "references", value: ExpressionUtils.array([ExpressionUtils.field("id")]) }, { name: "onDelete", value: ExpressionUtils.literal("Cascade") }] }],
+                    relation: { opposite: "sessions", fields: ["userId"], references: ["id"], onDelete: "Cascade" }
+                }
+            },
+            attributes: [
+                { name: "@@auth" }
+            ],
+            idFields: ["id"],
+            uniqueFields: {
+                id: { type: "String" }
+            }
+        },
+        User: {
+            name: "User",
+            fields: {
+                id: {
+                    name: "id",
+                    type: "String",
+                    id: true,
+                    attributes: [{ name: "@id" }, { name: "@default", args: [{ name: "value", value: ExpressionUtils.call("uuid") }] }],
+                    default: ExpressionUtils.call("uuid")
+                },
+                sessions: {
+                    name: "sessions",
+                    type: "Session",
+                    array: true,
+                    relation: { opposite: "user" }
+                },
+                profile: {
+                    name: "profile",
+                    type: "Profile",
+                    optional: true,
+                    relation: { opposite: "user" }
+                }
+            },
+            idFields: ["id"],
+            uniqueFields: {
+                id: { type: "String" }
+            }
+        },
+        Profile: {
+            name: "Profile",
+            fields: {
+                id: {
+                    name: "id",
+                    type: "String",
+                    id: true,
+                    attributes: [{ name: "@id" }, { name: "@default", args: [{ name: "value", value: ExpressionUtils.call("uuid") }] }],
+                    default: ExpressionUtils.call("uuid")
+                },
+                name: {
+                    name: "name",
+                    type: "String",
+                    optional: true
+                },
+                userId: {
+                    name: "userId",
+                    type: "String",
+                    unique: true,
+                    attributes: [{ name: "@unique" }],
+                    foreignKeyFor: [
+                        "user"
+                    ]
+                },
+                user: {
+                    name: "user",
+                    type: "User",
+                    attributes: [{ name: "@relation", args: [{ name: "fields", value: ExpressionUtils.array([ExpressionUtils.field("userId")]) }, { name: "references", value: ExpressionUtils.array([ExpressionUtils.field("id")]) }, { name: "onDelete", value: ExpressionUtils.literal("Cascade") }] }],
+                    relation: { opposite: "profile", fields: ["userId"], references: ["id"], onDelete: "Cascade" }
+                }
+            },
+            idFields: ["id"],
+            uniqueFields: {
+                id: { type: "String" },
+                userId: { type: "String" }
+            }
+        }
+    },
+    authType: "Session",
+    plugins: {}
+} as const satisfies SchemaDef;
+type Schema = typeof _schema & {
+    __brand?: "schema";
+};
+export const schema: Schema = _schema;
+export type SchemaType = Schema;

tests/regression/test/issue-422/schema.zmodel

@@ -0,0 +1,25 @@
+datasource db {
+    provider = "sqlite"
+    url      = "file:./dev.db"
+}
+
+model Session {
+    id     String @id
+    token  String
+    userId String
+    user   User   @relation(fields: [userId], references: [id], onDelete: Cascade)
+    @@auth
+}
+
+model User {
+    id       String    @id @default(uuid())
+    sessions Session[]
+    profile  Profile?
+}
+
+model Profile {
+    id     String  @id @default(uuid())
+    name   String?
+    userId String  @unique
+    user   User    @relation(fields: [userId], references: [id], onDelete: Cascade)
+}