Skip to content

Commit

Permalink
add jsonb selector
Browse files Browse the repository at this point in the history
  • Loading branch information
zfben committed Feb 8, 2025
1 parent c8bd82c commit 67cfe2b
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 8 deletions.
17 changes: 17 additions & 0 deletions src/__tests__/query-builder/query.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
78 changes: 70 additions & 8 deletions src/query-builder.ts
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -20,22 +20,70 @@ 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,
TResult = InferTResult<T>[],
> {
private client: Client
private table: T
private selectColumns: ColumnName<T>[] = []
private selectColumns: (ColumnName<T> | JsonSelectField<T>)[] = []
private whereConditions: {
column: ColumnName<T> | string
operator: Operator
Expand All @@ -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
Expand Down Expand Up @@ -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))
Expand Down

0 comments on commit 67cfe2b

Please sign in to comment.