From 67cfe2b9bfa57fe80e1bc87ad5f1810023dbb612 Mon Sep 17 00:00:00 2001
From: Ben <ben@zfben.com>
Date: Sat, 8 Feb 2025 05:10:38 +0000
Subject: [PATCH] add jsonb selector

---
 src/__tests__/query-builder/query.test.ts | 17 +++++
 src/query-builder.ts                      | 78 ++++++++++++++++++++---
 2 files changed, 87 insertions(+), 8 deletions(-)

diff --git a/src/__tests__/query-builder/query.test.ts b/src/__tests__/query-builder/query.test.ts
index b8f2cd3..ad49258 100644
--- a/src/__tests__/query-builder/query.test.ts
+++ b/src/__tests__/query-builder/query.test.ts
@@ -62,6 +62,23 @@ describe('QueryBuilder/query', () => {
       ])
       expectTypeOf(result).toEqualTypeOf<Pick<User, 'name' | 'metadata'>>()
     })
+
+    it('select jsonb', async () => {
+      const result = await new QueryBuilder(client, 'query')
+        .select({
+          column: 'metadata',
+          fields: ['age'],
+        })
+        .orderBy('id', 'ASC')
+
+      expect(result).toEqual([
+        { metadata: { age: 100 } },
+        { metadata: { age: null } },
+      ])
+      expectTypeOf(result).toEqualTypeOf<{
+        metadata: Pick<User['metadata'], 'age'>
+      }>()
+    })
   })
 
   describe('where', () => {
diff --git a/src/query-builder.ts b/src/query-builder.ts
index f3921a5..6a9e73f 100644
--- a/src/query-builder.ts
+++ b/src/query-builder.ts
@@ -1,5 +1,5 @@
 import type { Client } from './client'
-import type { ColumnName, ColumnValue, TableName, TableType } from './types'
+import type { ColumnName, ColumnValue, TableName, Tables, TableType } from './types'
 import { escapeIdentifier } from './utils'
 
 const NormalOperators = ['=', '!=', '<', '<=', '>', '>='] as const
@@ -20,14 +20,62 @@ const QueryOrderDirections = ['ASC', 'DESC'] as const
 
 type QueryOrderDirection = (typeof QueryOrderDirections)[number]
 
+type IsObject<T> = T extends object ? true : false
+
+type GetTableType<T extends TableName | string> = T extends keyof Tables
+  ? Tables[T]
+  : never
+
+type JsonbColumns<T extends TableName | string, Table = TableType<T>> = {
+  [K in keyof Table]: IsObject<Table[K]> extends true ? K : never
+}[keyof Table]
+
+type JsonbFields<
+  T extends TableName | string,
+  C extends JsonbColumns<T>,
+> = keyof GetTableType<T>[C & keyof GetTableType<T>]
+
+type JsonSelectField<T extends TableName | string> = {
+  column: JsonbColumns<T>
+  fields: JsonbFields<T, JsonbColumns<T>>[]
+  alias?: string
+}
+
+type InferJsonFields<
+  T extends TableName | string,
+  C extends JsonbColumns<T>,
+  Fields extends JsonbFields<T, C>[]
+> = {
+    [K in C & keyof GetTableType<T>]: Pick<GetTableType<T>[K], Fields[number]>
+  }
+
+type InferColumnType<
+  T extends TableName | string,
+  C extends ColumnName<T> | JsonSelectField<T>
+> = C extends JsonSelectField<T>
+  ? InferJsonFields<T, C['column'], C['fields']>
+  : C extends keyof GetTableType<T>
+  ? { [K in C]: GetTableType<T>[K] }
+  : never
+
+type Flatten<T> = { [K in keyof T]: T[K] }
+
+type UnionToIntersection<U> = (
+  U extends any ? (k: U) => void : never
+) extends (k: infer I) => void
+  ? I
+  : never
+
+type MergeTypes<T> = T extends any[] ? Flatten<UnionToIntersection<T[number]>> : T
+
 type InferTResult<
   TName extends TableName | string,
-  ColumnNames extends ColumnName<TName>[] = ColumnName<TName>[],
+  ColumnNames extends (ColumnName<TName> | JsonSelectField<TName>)[] = ColumnName<TName>[],
 > = ColumnNames extends ['*']
   ? TableType<TName>
-  : ColumnNames[number] extends keyof TableType<TName>
-  ? Pick<TableType<TName>, ColumnNames[number]>
-  : Record<string, any>[]
+  : MergeTypes<{
+    [K in keyof ColumnNames]: InferColumnType<TName, ColumnNames[K]>
+  }>
 
 export class QueryBuilder<
   T extends TableName | string = string,
@@ -35,7 +83,7 @@ export class QueryBuilder<
 > {
   private client: Client
   private table: T
-  private selectColumns: ColumnName<T>[] = []
+  private selectColumns: (ColumnName<T> | JsonSelectField<T>)[] = []
   private whereConditions: {
     column: ColumnName<T> | string
     operator: Operator
@@ -53,7 +101,21 @@ export class QueryBuilder<
     this.table = table
   }
 
-  select<ColumnNames extends ColumnName<T>[]>(
+  /**
+   * Selects specific columns for the query.
+   *
+   * @template ColumnNames - An array of column names or JSON select fields.
+   * @param {...ColumnNames} columns - The columns to select.
+   * @returns {QueryBuilder<T, InferTResult<T, ColumnNames>>} The query builder instance with the selected columns.
+   *
+   * @example
+   * ```ts
+   * const users = await db('users').select('id', 'name') // SELECT id, name FROM users
+   *
+   * const users = await db('users').select('id', { column: 'data', fields: ['email'] }) // SELECT id, jsonb_build_object('email', data->'email') AS data FROM users
+   * ```
+   */
+  select<ColumnNames extends (ColumnName<T> | JsonSelectField<T>)[]>(
     ...columns: ColumnNames
   ): QueryBuilder<T, InferTResult<T, ColumnNames>> {
     if (columns?.length > 0) this.selectColumns = columns
@@ -167,7 +229,7 @@ export class QueryBuilder<
     const params: any[] = []
 
     // Add columns
-    sql.push(this.selectColumns.map(escapeIdentifier).join(',') || '*')
+    sql.push(this.selectColumns.map(c => typeof c === 'string' ? escapeIdentifier(c) : `jsonb_build_object(${c.fields.map(f => `'${f as string}', ${escapeIdentifier(c.column as string)}->'${f as string}'`).join(',')}) AS ${escapeIdentifier(c.alias || c.column as string)}`).join(',') || '*')
 
     // Add table
     sql.push('FROM', escapeIdentifier(this.table))