refactor: Improve projection type safety

plexinc · Jan 19, 2024 · 4ef0675 · 4ef0675
1 parent 0b14413
commit 4ef0675
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 16 deletions.
diff --git a/src/__tests__/model.test.ts b/src/__tests__/model.test.ts
@@ -1075,6 +1075,18 @@ describe('model', () => {
       }
     });
 
+    test('with invalid projection', async () => {
+      await simpleModel.find(
+        {},
+        {
+          projection: {
+            // @ts-expect-error `nested.baz` is not present on the schema
+            'nested.baz': 1,
+          },
+        }
+      );
+    });
+
     test('with projection, re-assignable to Pick type', async () => {
       let results = await simpleModel.find({}, { projection });
 

diff --git a/src/model.ts b/src/model.ts
@@ -78,32 +78,32 @@ export interface Model<TSchema extends BaseSchema, TOptions extends SchemaOption
     options?: Omit<FindOptions<TSchema>, 'limit' | 'projection' | 'skip' | 'sort'>
   ) => Promise<boolean>;
 
-  find: <TProjection extends Projection<TSchema> | undefined>(
+  find: <TProjection extends Projection<TProjection, TSchema> | undefined>(
     filter: PaprFilter<TSchema>,
     options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }
   ) => Promise<ProjectionType<TSchema, TProjection>[]>;
 
-  findById: <TProjection extends Projection<TSchema> | undefined>(
+  findById: <TProjection extends Projection<TProjection, TSchema> | undefined>(
     id: TSchema['_id'] | string,
     options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }
   ) => Promise<ProjectionType<TSchema, TProjection> | null>;
 
-  findCursor: <TProjection extends Projection<TSchema> | undefined>(
+  findCursor: <TProjection extends Projection<TProjection, TSchema> | undefined>(
     filter: PaprFilter<TSchema>,
     options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }
   ) => Promise<FindCursor<ProjectionType<TSchema, TProjection>>>;
 
-  findOne: <TProjection extends Projection<TSchema> | undefined>(
+  findOne: <TProjection extends Projection<TProjection, TSchema> | undefined>(
     filter: PaprFilter<TSchema>,
     options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }
   ) => Promise<ProjectionType<TSchema, TProjection> | null>;
 
-  findOneAndDelete: <TProjection extends Projection<TSchema> | undefined>(
+  findOneAndDelete: <TProjection extends Projection<TProjection, TSchema> | undefined>(
     filter: PaprFilter<TSchema>,
     options?: Omit<FindOneAndDeleteOptions, 'projection'> & { projection?: TProjection }
   ) => Promise<ProjectionType<TSchema, TProjection> | null>;
 
-  findOneAndUpdate: <TProjection extends Projection<TSchema> | undefined>(
+  findOneAndUpdate: <TProjection extends Projection<TProjection, TSchema> | undefined>(
     filter: PaprFilter<TSchema>,
     update: PaprUpdateFilter<TSchema>,
     options?: Omit<FindOneAndUpdateOptions, 'projection'> & { projection?: TProjection }
@@ -621,7 +621,7 @@ export function build<TSchema extends BaseSchema, TOptions extends SchemaOptions
   // prettier-ignore
   model.find = wrap(
     model,
-    async function find<TProjection extends Projection<TSchema> | undefined>(
+    async function find<TProjection extends Projection<TProjection, TSchema> | undefined>(
       filter: PaprFilter<TSchema>,
       options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }
     ): Promise<ProjectionType<TSchema, TProjection>[]> {
@@ -663,7 +663,7 @@ export function build<TSchema extends BaseSchema, TOptions extends SchemaOptions
   // prettier-ignore
   model.findById = wrap(
     model,
-    async function findById<TProjection extends Projection<TSchema> | undefined>(
+    async function findById<TProjection extends Projection<TProjection, TSchema> | undefined>(
       id: TSchema['_id'] | string,
       options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }
     ): Promise<ProjectionType<TSchema, TProjection> | null> {
@@ -702,7 +702,7 @@ export function build<TSchema extends BaseSchema, TOptions extends SchemaOptions
    * }
    */
   model.findCursor = wrap(model, async function findCursor<
-    TProjection extends Projection<TSchema> | undefined,
+    TProjection extends Projection<TProjection, TSchema> | undefined,
   >(filter: PaprFilter<TSchema>, options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }): Promise<
     FindCursor<ProjectionType<TSchema, TProjection>>
   > {
@@ -741,7 +741,7 @@ export function build<TSchema extends BaseSchema, TOptions extends SchemaOptions
   // prettier-ignore
   model.findOne = wrap(
     model,
-    async function findOne<TProjection extends Projection<TSchema> | undefined>(
+    async function findOne<TProjection extends Projection<TProjection, TSchema> | undefined>(
       filter: PaprFilter<TSchema>,
       options?: Omit<FindOptions<TSchema>, 'projection'> & { projection?: TProjection }
     ): Promise<ProjectionType<TSchema, TProjection> | null> {
@@ -767,7 +767,7 @@ export function build<TSchema extends BaseSchema, TOptions extends SchemaOptions
    * const user = await User.findOneAndDelete({ firstName: 'John' });
    */
   // prettier-ignore
-  model.findOneAndDelete = wrap(model, async function findOneAndDelete<TProjection extends Projection<TSchema> | undefined>(
+  model.findOneAndDelete = wrap(model, async function findOneAndDelete<TProjection extends Projection<TProjection, TSchema> | undefined>(
     filter: PaprFilter<TSchema>,
     options?: Omit<FindOneAndDeleteOptions, 'projection'> & { projection?: TProjection }
   ): Promise<ProjectionType<TSchema, TProjection> | null> {
@@ -811,7 +811,7 @@ export function build<TSchema extends BaseSchema, TOptions extends SchemaOptions
    * userProjected.lastName; // valid
    */
   // prettier-ignore
-  model.findOneAndUpdate = wrap(model, async function findOneAndUpdate<TProjection extends Projection<TSchema> | undefined>(
+  model.findOneAndUpdate = wrap(model, async function findOneAndUpdate<TProjection extends Projection<TProjection, TSchema> | undefined>(
     filter: PaprFilter<TSchema>,
     update: PaprUpdateFilter<TSchema>,
     options?: Omit<FindOneAndUpdateOptions, 'projection'> & { projection?: TProjection }

diff --git a/src/utils.ts b/src/utils.ts
@@ -134,7 +134,7 @@ type FilterProperties<TObject, TValue> = Pick<TObject, KeysOfAType<TObject, TVal
 export type ProjectionType<
   TSchema extends BaseSchema,
   Projection extends
-    | Partial<Record<Join<NestedPaths<WithId<TSchema>, []>, '.'>, number>>
+    | ExactPartial<Projection, Record<Join<NestedPaths<WithId<TSchema>, []>, '.'>, number>>
     | undefined,
 > = undefined extends Projection
   ? WithId<TSchema>
@@ -147,9 +147,19 @@ export type ProjectionType<
           keyof FilterProperties<Projection, 0>
         >;
 
-export type Projection<TSchema> = Partial<
-  Record<Join<NestedPaths<WithId<TSchema>, []>, '.'>, number>
->;
+type ExactPartial<Subset, BaseObject> = {
+  [K in keyof Subset]: K extends keyof BaseObject
+    ? BaseObject[K] extends Record<number | string | symbol, unknown>
+      ? ExactPartial<Subset[K], BaseObject[K]>
+      : BaseObject[K]
+    : never;
+};
+
+export type Projection<
+  TProjection,
+  TSchema,
+  TPaths = Record<Join<NestedPaths<WithId<TSchema>, []>, '.'>, number>,
+> = ExactPartial<TProjection, TPaths>;
 
 export type PropertyNestedType<
   Type,