From 15ec87fdb0593dce1c1c7d356f0d09ae65b521d1 Mon Sep 17 00:00:00 2001
From: Igal Klebanov <igalklebanov@gmail.com>
Date: Fri, 29 Dec 2023 13:00:32 +0200
Subject: [PATCH] fix $if inside with is any. (#793)

---
 src/query-builder/select-query-builder.ts |  8 ++---
 test/typings/test-d/with.test-d.ts        | 41 +++++++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/src/query-builder/select-query-builder.ts b/src/query-builder/select-query-builder.ts
index 5f124be2b..46c357dfe 100644
--- a/src/query-builder/select-query-builder.ts
+++ b/src/query-builder/select-query-builder.ts
@@ -1416,7 +1416,7 @@ export interface SelectQueryBuilder<DB, TB extends keyof DB, O>
   $if<O2>(
     condition: boolean,
     func: (qb: this) => SelectQueryBuilder<any, any, O & O2>
-  ): SelectQueryBuilder<DB, TB, O & Partial<O2>>
+  ): SelectQueryBuilder<DB, TB, O & Partial<Omit<O2, keyof O>>>
 
   /**
    * Change the output type of the query.
@@ -2051,14 +2051,14 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
   $if<O2>(
     condition: boolean,
     func: (qb: this) => SelectQueryBuilder<any, any, O & O2>
-  ): SelectQueryBuilder<DB, TB, O & Partial<O2>> {
+  ): SelectQueryBuilder<DB, TB, O & Partial<Omit<O2, keyof O>>> {
     if (condition) {
       return func(this)
     }
 
-    return new SelectQueryBuilderImpl<DB, TB, O & Partial<O2>>({
+    return new SelectQueryBuilderImpl({
       ...this.#props,
-    })
+    }) as any
   }
 
   $castTo<T>(): SelectQueryBuilder<DB, TB, T> {
diff --git a/test/typings/test-d/with.test-d.ts b/test/typings/test-d/with.test-d.ts
index ad753a401..80da5cdff 100644
--- a/test/typings/test-d/with.test-d.ts
+++ b/test/typings/test-d/with.test-d.ts
@@ -85,6 +85,47 @@ async function testWith(db: Kysely<Database>) {
     }[]
   >(r4)
 
+  // https://github.com/kysely-org/kysely/issues/785
+  const r5 = await db
+    .with('person_projection', (qb) =>
+      qb
+        .selectFrom('person')
+        .select(['first_name', 'last_name'])
+        .$if(true, (qb) => qb.where('first_name', 'is not', null))
+    )
+    .selectFrom('person_projection')
+    .selectAll()
+    .$if(true, (qb) => qb.where('first_name', 'is not', null))
+    .execute()
+
+  expectType<
+    {
+      first_name: string
+      last_name: string | null
+    }[]
+  >(r5)
+
+  // testing fix of https://github.com/kysely-org/kysely/issues/785 didn't break $if that adds columns to the projection.
+  const r6 = await db
+    .with('person_projection', (qb) =>
+      qb
+        .selectFrom('person')
+        .select(['first_name', 'last_name'])
+        .$if(true, (qb) => qb.select('age'))
+    )
+    .selectFrom('person_projection')
+    .selectAll()
+    .$if(true, (qb) => qb.where('first_name', 'is not', null))
+    .execute()
+
+  expectType<
+    {
+      first_name: string
+      last_name: string | null
+      age: number | undefined
+    }[]
+  >(r6)
+
   // Different columns in expression and CTE name.
   expectError(
     db