From ad194e1aa15e014d0aea3009039b8caa6ef92767 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sami=20Koskim=C3=A4ki?= <koskomi@gmail.com>
Date: Sat, 30 Dec 2023 23:47:39 +0200
Subject: [PATCH] Add

---
 src/expression/expression-wrapper.ts          | 19 +++++++--
 .../aggregate-function-builder.ts             | 23 +++++++++++
 src/query-builder/delete-query-builder.ts     |  6 +--
 src/query-builder/insert-query-builder.ts     |  6 +--
 src/query-builder/json-path-builder.ts        |  6 ++-
 src/query-builder/select-query-builder.ts     |  8 ++--
 src/query-builder/update-query-builder.ts     |  6 +--
 src/raw-builder/raw-builder.ts                | 19 ++++++++-
 test/typings/test-d/expression.test-d.ts      | 12 ++++++
 test/typings/test-d/postgres-json.test-d.ts   | 41 +++++++++++++++++++
 10 files changed, 127 insertions(+), 19 deletions(-)

diff --git a/src/expression/expression-wrapper.ts b/src/expression/expression-wrapper.ts
index 16ff40dd4..a3b239cd0 100644
--- a/src/expression/expression-wrapper.ts
+++ b/src/expression/expression-wrapper.ts
@@ -239,7 +239,20 @@ export class ExpressionWrapper<DB, TB extends keyof DB, T>
    * This method call doesn't change the SQL in any way. This methods simply
    * returns a copy of this `ExpressionWrapper` with a new output type.
    */
-  $castTo<T>(): ExpressionWrapper<DB, TB, T> {
+  $castTo<C>(): ExpressionWrapper<DB, TB, C> {
+    return new ExpressionWrapper(this.#node)
+  }
+
+  /**
+   * Omit null from the expression's type.
+   *
+   * This function can be useful in cases where you know an expression can't be
+   * null, but Kysely is unable to infer it.
+   *
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of `this` with a new output type.
+   */
+  $notNull(): ExpressionWrapper<DB, TB, Exclude<T, null>> {
     return new ExpressionWrapper(this.#node)
   }
 
@@ -352,7 +365,7 @@ export class OrWrapper<DB, TB extends keyof DB, T extends SqlBool>
    * This method call doesn't change the SQL in any way. This methods simply
    * returns a copy of this `OrWrapper` with a new output type.
    */
-  $castTo<T extends SqlBool>(): OrWrapper<DB, TB, T> {
+  $castTo<C extends SqlBool>(): OrWrapper<DB, TB, C> {
     return new OrWrapper(this.#node)
   }
 
@@ -436,7 +449,7 @@ export class AndWrapper<DB, TB extends keyof DB, T extends SqlBool>
    * This method call doesn't change the SQL in any way. This methods simply
    * returns a copy of this `AndWrapper` with a new output type.
    */
-  $castTo<T extends SqlBool>(): AndWrapper<DB, TB, T> {
+  $castTo<C extends SqlBool>(): AndWrapper<DB, TB, C> {
     return new AndWrapper(this.#node)
   }
 
diff --git a/src/query-builder/aggregate-function-builder.ts b/src/query-builder/aggregate-function-builder.ts
index 9d7bfb074..e1d36bb2d 100644
--- a/src/query-builder/aggregate-function-builder.ts
+++ b/src/query-builder/aggregate-function-builder.ts
@@ -272,6 +272,29 @@ export class AggregateFunctionBuilder<DB, TB extends keyof DB, O = unknown>
     return func(this)
   }
 
+  /**
+   * Casts the expression to the given type.
+   *
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of this `AggregateFunctionBuilder` with a new output type.
+   */
+  $castTo<C>(): AggregateFunctionBuilder<DB, TB, C> {
+    return new AggregateFunctionBuilder(this.#props)
+  }
+
+  /**
+   * Omit null from the expression's type.
+   *
+   * This function can be useful in cases where you know an expression can't be
+   * null, but Kysely is unable to infer it.
+   *
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of `this` with a new output type.
+   */
+  $notNull(): AggregateFunctionBuilder<DB, TB, Exclude<O, null>> {
+    return new AggregateFunctionBuilder(this.#props)
+  }
+
   toOperationNode(): AggregateFunctionNode {
     return this.#props.aggregateFunctionNode
   }
diff --git a/src/query-builder/delete-query-builder.ts b/src/query-builder/delete-query-builder.ts
index 8197851ee..5f476317f 100644
--- a/src/query-builder/delete-query-builder.ts
+++ b/src/query-builder/delete-query-builder.ts
@@ -716,10 +716,10 @@ export class DeleteQueryBuilder<DB, TB extends keyof DB, O>
   /**
    * Change the output type of the query.
    *
-   * You should only use this method as the last resort if the types
-   * don't support your use case.
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of this `DeleteQueryBuilder` with a new output type.
    */
-  $castTo<T>(): DeleteQueryBuilder<DB, TB, T> {
+  $castTo<C>(): DeleteQueryBuilder<DB, TB, C> {
     return new DeleteQueryBuilder(this.#props)
   }
 
diff --git a/src/query-builder/insert-query-builder.ts b/src/query-builder/insert-query-builder.ts
index 46ec4c0ce..6c75ea670 100644
--- a/src/query-builder/insert-query-builder.ts
+++ b/src/query-builder/insert-query-builder.ts
@@ -677,10 +677,10 @@ export class InsertQueryBuilder<DB, TB extends keyof DB, O>
   /**
    * Change the output type of the query.
    *
-   * You should only use this method as the last resort if the types
-   * don't support your use case.
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of this `InsertQueryBuilder` with a new output type.
    */
-  $castTo<T>(): InsertQueryBuilder<DB, TB, T> {
+  $castTo<C>(): InsertQueryBuilder<DB, TB, C> {
     return new InsertQueryBuilder(this.#props)
   }
 
diff --git a/src/query-builder/json-path-builder.ts b/src/query-builder/json-path-builder.ts
index 4d4abe873..c729bb139 100644
--- a/src/query-builder/json-path-builder.ts
+++ b/src/query-builder/json-path-builder.ts
@@ -246,7 +246,11 @@ export class TraversedJSONPathBuilder<S, O>
    * This method call doesn't change the SQL in any way. This methods simply
    * returns a copy of this `JSONPathBuilder` with a new output type.
    */
-  $castTo<T>(): JSONPathBuilder<T> {
+  $castTo<C>(): JSONPathBuilder<C> {
+    return new JSONPathBuilder(this.#node)
+  }
+
+  $notNull(): JSONPathBuilder<Exclude<O, null>> {
     return new JSONPathBuilder(this.#node)
   }
 
diff --git a/src/query-builder/select-query-builder.ts b/src/query-builder/select-query-builder.ts
index 541cd2ff2..21005ab4e 100644
--- a/src/query-builder/select-query-builder.ts
+++ b/src/query-builder/select-query-builder.ts
@@ -1439,10 +1439,10 @@ export interface SelectQueryBuilder<DB, TB extends keyof DB, O>
   /**
    * Change the output type of the query.
    *
-   * You should only use this method as the last resort if the types
-   * don't support your use case.
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of this `SelectQueryBuilder` with a new output type.
    */
-  $castTo<T>(): SelectQueryBuilder<DB, TB, T>
+  $castTo<C>(): SelectQueryBuilder<DB, TB, C>
 
   /**
    * Changes the output type from an object to a tuple.
@@ -2095,7 +2095,7 @@ class SelectQueryBuilderImpl<DB, TB extends keyof DB, O>
     }) as any
   }
 
-  $castTo<T>(): SelectQueryBuilder<DB, TB, T> {
+  $castTo<C>(): SelectQueryBuilder<DB, TB, C> {
     return new SelectQueryBuilderImpl(this.#props)
   }
 
diff --git a/src/query-builder/update-query-builder.ts b/src/query-builder/update-query-builder.ts
index 2f51eea46..c43fe483e 100644
--- a/src/query-builder/update-query-builder.ts
+++ b/src/query-builder/update-query-builder.ts
@@ -681,10 +681,10 @@ export class UpdateQueryBuilder<DB, UT extends keyof DB, TB extends keyof DB, O>
   /**
    * Change the output type of the query.
    *
-   * You should only use this method as the last resort if the types
-   * don't support your use case.
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of this `UpdateQueryBuilder` with a new output type.
    */
-  $castTo<T>(): UpdateQueryBuilder<DB, UT, TB, T> {
+  $castTo<C>(): UpdateQueryBuilder<DB, UT, TB, C> {
     return new UpdateQueryBuilder(this.#props)
   }
 
diff --git a/src/raw-builder/raw-builder.ts b/src/raw-builder/raw-builder.ts
index eacecf3d5..bce69ab82 100644
--- a/src/raw-builder/raw-builder.ts
+++ b/src/raw-builder/raw-builder.ts
@@ -88,7 +88,18 @@ export interface RawBuilder<O> extends AliasableExpression<O> {
    * This method call doesn't change the SQL in any way. This methods simply
    * returns a copy of this `RawBuilder` with a new output type.
    */
-  $castTo<T>(): RawBuilder<T>
+  $castTo<C>(): RawBuilder<C>
+
+  /**
+   * Omit null from the expression's type.
+   *
+   * This function can be useful in cases where you know an expression can't be
+   * null, but Kysely is unable to infer it.
+   *
+   * This method call doesn't change the SQL in any way. This methods simply
+   * returns a copy of `this` with a new output type.
+   */
+  $notNull(): RawBuilder<Exclude<O, null>>
 
   /**
    * Adds a plugin for this SQL snippet.
@@ -140,10 +151,14 @@ class RawBuilderImpl<O> implements RawBuilder<O> {
     return new AliasedRawBuilderImpl(this, alias)
   }
 
-  $castTo<T>(): RawBuilder<T> {
+  $castTo<C>(): RawBuilder<C> {
     return new RawBuilderImpl({ ...this.#props })
   }
 
+  $notNull(): RawBuilder<Exclude<O, null>> {
+    return new RawBuilderImpl(this.#props)
+  }
+
   withPlugin(plugin: KyselyPlugin): RawBuilder<O> {
     return new RawBuilderImpl({
       ...this.#props,
diff --git a/test/typings/test-d/expression.test-d.ts b/test/typings/test-d/expression.test-d.ts
index 68a8a3cc3..6abd4944a 100644
--- a/test/typings/test-d/expression.test-d.ts
+++ b/test/typings/test-d/expression.test-d.ts
@@ -102,6 +102,18 @@ async function testExpressionBuilder(
     })
   )
 
+  expectAssignable<Expression<number | null>>(
+    eb.case().when('age', '=', 10).then(1).else(null).end()
+  )
+
+  expectNotAssignable<Expression<number>>(
+    eb.case().when('age', '=', 10).then(1).else(null).end()
+  )
+
+  expectAssignable<Expression<number>>(
+    eb.case().when('age', '=', 10).then(1).else(null).end().$notNull()
+  )
+
   expectType<
     KyselyTypeError<'or() method can only be called on boolean expressions'>
   >(eb('age', '+', 1).or('age', '=', 1))
diff --git a/test/typings/test-d/postgres-json.test-d.ts b/test/typings/test-d/postgres-json.test-d.ts
index 36f226bb4..add33942e 100644
--- a/test/typings/test-d/postgres-json.test-d.ts
+++ b/test/typings/test-d/postgres-json.test-d.ts
@@ -128,6 +128,47 @@ async function testPostgresJsonAgg(db: Kysely<Database>) {
       pets: Selectable<Pet>[] | null
     }[]
   >(r3)
+
+  const db2 = db.withTables<{
+    acquisition: {
+      id: number
+    }
+    transaction: {
+      id: number
+      acquisitionId: number
+      status: string
+    }
+  }>()
+
+  const r4 = await db2
+    .selectFrom('acquisition')
+    .leftJoin('transaction', 'transaction.acquisitionId', 'acquisition.id')
+    .select(({ ref, fn }) => [
+      'acquisition.id',
+      fn
+        .coalesce(
+          fn
+            .jsonAgg(
+              jsonBuildObject({
+                id: ref('transaction.id').$notNull(),
+                status: ref('transaction.status'),
+              })
+            )
+            .filterWhere('transaction.id', 'is not', null),
+          sql`'[]'`
+        )
+        .as('transactions'),
+    ])
+    .groupBy('acquisition.id')
+    .executeTakeFirstOrThrow()
+
+  expectType<{
+    id: number
+    transactions: {
+      id: number
+      status: string | null
+    }[]
+  }>(r4)
 }
 
 async function testPostgresToJson(db: Kysely<Database>) {