upsert support and fixes about "omit"

TODO.md

@@ -43,7 +43,7 @@
         -   [x] Nested to-one
         -   [ ] Delta update for numeric fields
         -   [ ] Array update
-    -   [ ] Upsert
+    -   [x] Upsert
     -   [x] Delete
     -   [ ] Aggregation
         -   [x] Count

packages/runtime/src/client/client-impl.ts

@@ -266,7 +266,7 @@ function createModelCrudHandler<
         args: unknown,
         handler: BaseOperationHandler<Schema>,
         postProcess = false,
-        throwIfNotFound = false
+        throwIfNoResult = false
     ) => {
         return createDeferredPromise(async () => {
             let proceed = async (
@@ -275,7 +275,7 @@ function createModelCrudHandler<
             ) => {
                 const _handler = tx ? handler.withClient(tx) : handler;
                 const r = await _handler.handle(operation, _args ?? args);
-                if (!r && throwIfNotFound) {
+                if (!r && throwIfNoResult) {
                     throw new NotFoundError(model);
                 }
                 let result: unknown;
@@ -400,6 +400,15 @@ function createModelCrudHandler<
             );
         },
 
+        upsert: (args: unknown) => {
+            return createPromise(
+                'upsert',
+                args,
+                new UpdateOperationHandler(client, model, inputValidator),
+                true
+            );
+        },
+
         delete: (args: unknown) => {
             return createPromise(
                 'delete',

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

@@ -39,29 +39,35 @@ import type { ToKyselySchema } from './query-builder';
 type DefaultModelResult<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>,
+    Omit = undefined,
     Optional = false,
     Array = false
 > = WrapType<
     {
-        [Key in NonRelationFields<Schema, Model>]: MapFieldType<
-            Schema,
-            Model,
-            Key
-        >;
+        [Key in NonRelationFields<Schema, Model> as Key extends keyof Omit
+            ? Omit[Key] extends true
+                ? never
+                : Key
+            : Key]: MapFieldType<Schema, Model, Key>;
     },
     Optional,
     Array
 >;
 
 type ModelSelectResult<
-    Select,
     Schema extends SchemaDef,
-    Model extends GetModels<Schema>
+    Model extends GetModels<Schema>,
+    Select,
+    Omit
 > = {
     [Key in keyof Select & GetFields<Schema, Model> as Select[Key] extends
         | false
         | undefined
         ? never
+        : Key extends keyof Omit
+        ? Omit[Key] extends true
+            ? never
+            : Key
         : Key]: Key extends ScalarFields<Schema, Model>
         ? MapFieldType<Schema, Model, Key>
         : Key extends RelationFields<Schema, Model>
@@ -80,6 +86,7 @@ type ModelSelectResult<
             : DefaultModelResult<
                   Schema,
                   RelationFieldType<Schema, Model, Key>,
+                  Omit,
                   FieldIsOptional<Schema, Model, Key>,
                   FieldIsArray<Schema, Model, Key>
               >
@@ -89,18 +96,20 @@ type ModelSelectResult<
 export type ModelResult<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>,
-    Args extends SelectInclude<Schema, Model, boolean> = {},
+    Args extends SelectIncludeOmit<Schema, Model, boolean> = {},
     Optional = false,
     Array = false
 > = WrapType<
     Args extends {
         select: infer S;
+        omit?: infer O;
     }
-        ? ModelSelectResult<S, Schema, Model>
+        ? ModelSelectResult<Schema, Model, S, O>
         : Args extends {
               include: infer I;
+              omit?: infer O;
           }
-        ? DefaultModelResult<Schema, Model> & {
+        ? DefaultModelResult<Schema, Model, O> & {
               [Key in keyof I & RelationFields<Schema, Model> as I[Key] extends
                   | false
                   | undefined
@@ -124,6 +133,8 @@ export type ModelResult<
                         FieldIsArray<Schema, Model, Key>
                     >;
           }
+        : Args extends { omit: infer O }
+        ? DefaultModelResult<Schema, Model, O>
         : DefaultModelResult<Schema, Model>,
     Optional,
     Array
@@ -311,7 +322,7 @@ type OmitFields<Schema extends SchemaDef, Model extends GetModels<Schema>> = {
     [Key in ScalarFields<Schema, Model>]?: true;
 };
 
-export type SelectInclude<
+export type SelectIncludeOmit<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>,
     AllowCount extends boolean
@@ -528,14 +539,14 @@ export type FindArgs<
               where?: Where<Schema, Model>;
           }
         : {}) &
-    SelectInclude<Schema, Model, Collection>;
+    SelectIncludeOmit<Schema, Model, Collection>;
 
 export type FindUniqueArgs<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>
 > = {
     where?: WhereUnique<Schema, Model>;
-} & SelectInclude<Schema, Model, true>;
+} & SelectIncludeOmit<Schema, Model, true>;
 
 //#endregion
 
@@ -548,6 +559,7 @@ export type CreateArgs<
     data: CreateInput<Schema, Model>;
     select?: Select<Schema, Model, true>;
     include?: Include<Schema, Model>;
+    omit?: OmitFields<Schema, Model>;
 };
 
 export type CreateManyArgs<
@@ -559,7 +571,7 @@ export type CreateManyAndReturnArgs<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>
 > = CreateManyPayload<Schema, Model> &
-    Omit<SelectInclude<Schema, Model, false>, 'include'>;
+    Omit<SelectIncludeOmit<Schema, Model, false>, 'include'>;
 
 type OptionalWrap<
     Schema extends SchemaDef,
@@ -678,6 +690,7 @@ export type UpdateArgs<
     where: WhereUnique<Schema, Model>;
     select?: Select<Schema, Model, true>;
     include?: Include<Schema, Model>;
+    omit?: OmitFields<Schema, Model>;
 };
 
 export type UpdateManyArgs<
@@ -689,6 +702,18 @@ export type UpdateManyArgs<
     limit?: number;
 };
 
+export type UpsertArgs<
+    Schema extends SchemaDef,
+    Model extends GetModels<Schema>
+> = {
+    create: CreateInput<Schema, Model>;
+    update: UpdateInput<Schema, Model>;
+    where: WhereUnique<Schema, Model>;
+    select?: Select<Schema, Model, true>;
+    include?: Include<Schema, Model>;
+    omit?: OmitFields<Schema, Model>;
+};
+
 export type UpdateScalarInput<
     Schema extends SchemaDef,
     Model extends GetModels<Schema>,
@@ -761,6 +786,7 @@ export type DeleteArgs<
     where: WhereUnique<Schema, Model>;
     select?: Select<Schema, Model, true>;
     include?: Include<Schema, Model>;
+    omit?: OmitFields<Schema, Model>;
 };
 
 export type DeleteManyArgs<
@@ -1104,6 +1130,10 @@ export type ModelOperations<
         args: Subset<T, UpdateManyArgs<Schema, Model>>
     ): Promise<BatchResult>;
 
+    upsert<T extends UpsertArgs<Schema, Model>>(
+        args: SelectSubset<T, UpsertArgs<Schema, Model>>
+    ): Promise<ModelResult<Schema, Model, T>>;
+
     delete<T extends DeleteArgs<Schema, Model>>(
         args: SelectSubset<T, DeleteArgs<Schema, Model>>
     ): Promise<ModelResult<Schema, Model>>;

packages/runtime/src/client/crud/operations/base.ts

@@ -29,7 +29,7 @@ import {
 } from '../../../utils/object-utils';
 import { CONTEXT_COMMENT_PREFIX } from '../../constants';
 import type { CRUD } from '../../contract';
-import type { FindArgs, SelectInclude, Where } from '../../crud-types';
+import type { FindArgs, SelectIncludeOmit, Where } from '../../crud-types';
 import { InternalError, NotFoundError, QueryError } from '../../errors';
 import type { ToKysely } from '../../query-builder';
 import {
@@ -59,6 +59,7 @@ export type CrudOperation =
     | 'createManyAndReturn'
     | 'update'
     | 'updateMany'
+    | 'upsert'
     | 'delete'
     | 'deleteMany'
     | 'count'
@@ -165,8 +166,10 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
 
         // select
         if (args?.select) {
+            // select is mutually exclusive with omit
             query = this.buildFieldSelection(model, query, args?.select, model);
         } else {
+            // include all scalar fields except those in omit
             query = this.buildSelectAllScalarFields(model, query, args?.omit);
         }
 
@@ -1783,7 +1786,7 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
 
     protected trimResult(
         data: any,
-        args: SelectInclude<Schema, GetModels<Schema>, boolean>
+        args: SelectIncludeOmit<Schema, GetModels<Schema>, boolean>
     ) {
         if (!args.select) {
             return data;
@@ -1796,7 +1799,7 @@ export abstract class BaseOperationHandler<Schema extends SchemaDef> {
 
     protected needReturnRelations(
         model: string,
-        args: SelectInclude<Schema, GetModels<Schema>, boolean>
+        args: SelectIncludeOmit<Schema, GetModels<Schema>, boolean>
     ) {
         let returnRelation = false;
 

packages/runtime/src/client/crud/operations/create.ts

@@ -45,6 +45,7 @@ export class CreateOperationHandler<
             return this.readUnique(tx, this.model, {
                 select: args.select,
                 include: args.include,
+                omit: args.omit,
                 where: getIdValues(this.schema, this.model, createResult),
             });
         });

packages/runtime/src/client/crud/operations/delete.ts

@@ -31,6 +31,7 @@ export class DeleteOperationHandler<
         const existing = await this.readUnique(this.kysely, this.model, {
             select: args.select,
             include: args.include,
+            omit: args.omit,
             where: args.where,
         });
         if (!existing) {

packages/runtime/src/client/crud/operations/update.ts

@@ -1,14 +1,14 @@
 import { match } from 'ts-pattern';
 import { RejectedByPolicyError } from '../../../plugins/policy/errors';
 import type { GetModels, SchemaDef } from '../../../schema';
-import type { UpdateArgs, UpdateManyArgs } from '../../crud-types';
-import { BaseOperationHandler } from './base';
+import type { UpdateArgs, UpdateManyArgs, UpsertArgs } from '../../crud-types';
 import { getIdValues } from '../../query-utils';
+import { BaseOperationHandler } from './base';
 
 export class UpdateOperationHandler<
     Schema extends SchemaDef
 > extends BaseOperationHandler<Schema> {
-    async handle(operation: 'update' | 'updateMany', args: unknown) {
+    async handle(operation: 'update' | 'updateMany' | 'upsert', args: unknown) {
         return match(operation)
             .with('update', () =>
                 this.runUpdate(
@@ -20,6 +20,11 @@ export class UpdateOperationHandler<
                     this.inputValidator.validateUpdateManyArgs(this.model, args)
                 )
             )
+            .with('upsert', () =>
+                this.runUpsert(
+                    this.inputValidator.validateUpsertArgs(this.model, args)
+                )
+            )
             .exhaustive();
     }
 
@@ -34,6 +39,7 @@ export class UpdateOperationHandler<
             return this.readUnique(tx, this.model, {
                 select: args.select,
                 include: args.include,
+                omit: args.omit,
                 where: getIdValues(this.schema, this.model, updated),
             });
         });
@@ -59,4 +65,39 @@ export class UpdateOperationHandler<
             args.limit
         );
     }
+
+    private async runUpsert(args: UpsertArgs<Schema, GetModels<Schema>>) {
+        const result = await this.safeTransaction(async (tx) => {
+            let mutationResult = await this.update(
+                tx,
+                this.model,
+                args.where,
+                args.update,
+                undefined,
+                true,
+                false
+            );
+
+            if (!mutationResult) {
+                // non-existing, create
+                mutationResult = await this.create(tx, this.model, args.create);
+            }
+
+            return this.readUnique(tx, this.model, {
+                select: args.select,
+                include: args.include,
+                omit: args.omit,
+                where: getIdValues(this.schema, this.model, mutationResult),
+            });
+        });
+
+        if (!result) {
+            throw new RejectedByPolicyError(
+                this.model,
+                'result is not allowed to be read back'
+            );
+        }
+
+        return result;
+    }
 }