fix(builder): fix compiled return type infer

Author: subframe7536

packages/sqlite-builder/src/builder.ts

@@ -1,10 +1,10 @@
 import type { Compilable, CompiledQuery, KyselyPlugin, LogEvent, QueryResult, RawBuilder, Sql, Transaction } from 'kysely'
 import { Kysely, sql } from 'kysely'
 import { SqliteSerializePlugin } from 'kysely-plugin-serialize'
+import type { SimplifySingleResult } from 'kysely/dist/cjs/util/type-utils'
 import { parseTableMap, preCompile, runCreateTable } from './utils'
-import type { AvailableBuilder, Logger, SqliteBuilderOption, Table } from './types'
+import type { AvailableBuilder, BuilderResult, Logger, QueryBuilderOutput, SqliteBuilderOption, Table } from './types'
 import { Stack } from './stack'
-import type { QueryBuilderOutput } from './utils'
 
 const enum DBStatus {
   'needDrop',
@@ -96,15 +96,15 @@ export class SqliteBuilder<DB extends Record<string, any>> {
   public async execOne<O>(
     cb: (db: Kysely<DB> | Transaction<DB>) => AvailableBuilder<DB, O>,
     errorMsg?: string,
-  ): Promise<O | undefined> {
+  ): Promise<SimplifySingleResult<O> | undefined> {
     const resultList = await this.execList(cb, errorMsg)
     return resultList?.length ? resultList[0] : undefined
   }
 
   public async execList<O>(
     cb: (db: Kysely<DB> | Transaction<DB>) => AvailableBuilder<DB, O>,
     errorMsg?: string,
-  ): Promise<O[] | undefined> {
+  ): Promise<Exclude<SimplifySingleResult<O>, undefined>[] | undefined> {
     if (await this.isEmptyTable()) {
       return undefined
     }
@@ -113,36 +113,35 @@ export class SqliteBuilder<DB extends Record<string, any>> {
       .catch((err) => {
         errorMsg && this.logger?.error(errorMsg, err)
         return undefined
-      })
+      }) as any
   }
 
   public preCompile<O>(
     queryBuilder: (db: Kysely<DB> | Transaction<DB>) => QueryBuilderOutput<Compilable<O>>,
-  ): ReturnType<typeof preCompile<DB, O>> {
-    return preCompile(this.kysely, queryBuilder)
+  ): ReturnType<typeof preCompile<O>> {
+    return preCompile(queryBuilder(this.kysely))
   }
 
-  public async execCompiledOne<O>(
+  public async execCompiledRows<O>(
     query: CompiledQuery<O>,
     errorMsg?: string,
-  ): Promise<O | undefined> {
-    const resultList = await this.execCompiledList(query, errorMsg)
-    return resultList?.length ? resultList[0] : undefined
+  ): Promise<BuilderResult<O>['rows'] | undefined> {
+    const result = await this.execCompiled(query, errorMsg)
+    return result?.rows ?? undefined
   }
 
-  public async execCompiledList<O>(
+  public async execCompiled<O>(
     query: CompiledQuery<O>,
     errorMsg?: string,
-  ): Promise<O[] | undefined> {
+  ): Promise<BuilderResult<O> | undefined> {
     if (await this.isEmptyTable()) {
       return undefined
     }
     return this.getDB().executeQuery(query)
-      .then(result => result.rows)
       .catch((err) => {
         errorMsg && this.logger?.error(errorMsg, err)
         return undefined
-      })
+      }) as any
   }
 
   public async toSQL<T extends Compilable>(cb: (db: Kysely<DB>) => T): Promise<CompiledQuery<unknown>> {

packages/sqlite-builder/src/types.ts

@@ -1,4 +1,4 @@
-import type { DeleteQueryBuilder, Dialect, Generated, InsertQueryBuilder, KyselyPlugin, SelectQueryBuilder, Sql, UpdateQueryBuilder } from 'kysely'
+import type { Compilable, DeleteQueryBuilder, Dialect, Generated, InsertQueryBuilder, KyselyPlugin, QueryResult, SelectQueryBuilder, Sql, UpdateQueryBuilder } from 'kysely'
 import type { CompiledQuery } from 'kysely/dist/cjs/query-compiler/compiled-query'
 
 export type TriggerEvent = 'insert' | 'update' | 'delete'
@@ -54,7 +54,11 @@ export type Logger = {
 }
 
 export type AvailableBuilder<DB, O> =
-  | SelectQueryBuilder<DB, keyof DB, O>
-  | UpdateQueryBuilder<DB, keyof DB, keyof DB, O>
-  | InsertQueryBuilder<DB, keyof DB, O>
-  | DeleteQueryBuilder<DB, keyof DB, O>
+  | SelectQueryBuilder<DB, any, O>
+  | UpdateQueryBuilder<DB, any, any, O>
+  | InsertQueryBuilder<DB, any, O>
+  | DeleteQueryBuilder<DB, any, O>
+
+export type QueryBuilderOutput<QB> = QB extends Compilable<infer O> ? O : never
+
+export type BuilderResult<T> = T extends SelectQueryBuilder<any, any, infer P> ? QueryResult<P> : Omit<QueryResult<any>, 'rows'> & { rows: [] }

packages/sqlite-builder/src/utils.ts

@@ -1,8 +1,8 @@
-import type { Compilable, CompiledQuery, Kysely, Transaction } from 'kysely'
+import type { Compilable, CompiledQuery, Kysely, RootOperationNode } from 'kysely'
 import { sql } from 'kysely'
 import type { DataTypeExpression } from 'kysely/dist/cjs/parser/data-type-parser'
 import { defaultSerializer } from 'kysely-plugin-serialize'
-import type { ColumeOption, Table, Tables, TriggerEvent } from './types'
+import type { ColumeOption, QueryBuilderOutput, Table, Tables, TriggerEvent } from './types'
 
 export function isString(value: any): value is string {
   return typeof value === 'string'
@@ -143,11 +143,8 @@ export async function runCreateTable<T>(
   }
 }
 
-export type QueryBuilderOutput<QB> = QB extends Compilable<infer O> ? O : never
-
-export function preCompile<DB, O>(
-  db: Kysely<DB> | Transaction<DB>,
-  queryBuilder: (db: Kysely<DB> | Transaction<DB>) => QueryBuilderOutput<Compilable<O>>,
+export function preCompile<O>(
+  queryBuilder: QueryBuilderOutput<Compilable<O>>,
 ) {
   function getParam<P extends Record<string, any>>(name: keyof P): P[keyof P] {
     return `__precomile_${name as string}` as unknown as P[keyof P]
@@ -160,12 +157,12 @@ export function preCompile<DB, O>(
       ) => Compilable<O>,
     ) {
       let compiled: CompiledQuery<Compilable<O>>
-      return (param: P) => {
+      return (param: P, processRootOperatorNode?: (node: RootOperationNode) => RootOperationNode): CompiledQuery<O> => {
         if (!compiled) {
-          const { parameters, sql, query } = paramBuilder(queryBuilder(db), getParam).compile()
+          const { parameters, sql, query } = paramBuilder(queryBuilder, getParam).compile()
           compiled = {
             sql,
-            query: { kind: query.kind } as any,
+            query: processRootOperatorNode?.(query) || { kind: query.kind } as any,
             parameters,
           }
         }

test/builder.test.ts

@@ -6,6 +6,7 @@ import { SqliteBuilder } from '../packages/sqlite-builder/src'
 
 interface DB {
   test: TestTable
+  pet: PetTable
 }
 interface TestTable {
   id: Generated<number>
@@ -15,6 +16,13 @@ interface TestTable {
   updateAt: Date | null
   buffer: ArrayBuffer | null
 }
+interface PetTable {
+  id: Generated<string>
+  name: string
+  owner_id: string
+  species: 'cat' | 'dog'
+  is_favorite: boolean
+}
 describe('test builder', async () => {
   const db = new SqliteBuilder<DB>({
     dialect: new SqliteDialect({
@@ -36,18 +44,30 @@ describe('test builder', async () => {
           timestamp: true,
         },
       },
+      pet: {
+        columns: {
+          id: { type: 'increments' },
+          is_favorite: { type: 'boolean' },
+          name: { type: 'string' },
+          owner_id: { type: 'string' },
+          species: { type: 'string' },
+        },
+      },
     },
     dropTableBeforeInit: true,
     logger: console,
-    onQuery: (queryInfo, time) => console.log(`${time}ms`, queryInfo.sql, queryInfo.parameters),
+    // onQuery: (queryInfo, time) => console.log(`${time}ms`, queryInfo.sql, queryInfo.parameters),
   })
   beforeAll(async () => {
     // manually generate table
     await db.init(true)
   })
   test('insert', async () => {
     // auto generate table
-    await db.transaction(trx => trx.insertInto('test').values({ gender: false }).execute())
+    console.log(await db.transaction((trx) => {
+      trx.insertInto('test').values([{ gender: false }, { gender: true }]).execute()
+      return trx.updateTable('test').set({ gender: true }).where('id', '=', 2).execute()
+    }))
     const result = await db.execList(d => d.selectFrom('test').selectAll())
     expect(result).toBeInstanceOf(Array)
     expect(result![0].person).toStrictEqual({ name: 'test' })
@@ -62,31 +82,42 @@ describe('test builder', async () => {
     expect(result2!.updateAt).toBeInstanceOf(Date)
   })
   test('raw', async () => {
-    const { sql, parameters } = await db.toSQL(d => d
+    const query = await db.toSQL(d => d
       .selectFrom('test')
       .where('person', '=', { name: '1' })
       .selectAll(),
     )
-    expect(sql).toBe('select * from "test" where "person" = ?')
-    expect(parameters).toStrictEqual(['{"name":"1"}'])
+    expect(query.sql).toBe('select * from "test" where "person" = ?')
+    expect(query.parameters).toStrictEqual(['{"name":"1"}'])
+    const result = await db.execCompiled(query)
+    console.log(result)
   })
   test('preCompile', async () => {
-    const fn = db.preCompile(db => db.selectFrom('test').selectAll())
+    const select = db.preCompile(db => db.selectFrom('test').selectAll())
       .setParam<{ person: { name: string } }>((qb, param) => qb.where('person', '=', param('person')))
+    const insert = db.preCompile(db => db.insertInto('test'))
+      .setParam<{ gender: boolean }>((qb, param) => qb.values({ gender: param('gender') }))
+    const update = db.preCompile(db => db.updateTable('test'))
+      .setParam<{ gender: boolean }>((qb, param) => qb.set({ gender: param('gender') }).where('id', '=', 1))
 
     const start = performance.now()
 
-    const { parameters, sql } = fn({ person: { name: '1' } })
+    const { parameters, sql } = select({ person: { name: '1' } })
     expect(sql).toBe('select * from "test" where "person" = ?')
     expect(parameters).toStrictEqual(['{"name":"1"}'])
 
     const start2 = performance.now()
     console.log('no compiled:', `${(start2 - start).toFixed(2)}ms`)
 
-    const { parameters: p1, sql: s1 } = fn({ person: { name: 'test' } })
+    const { parameters: p1, sql: s1 } = select({ person: { name: 'test' } })
     expect(s1).toBe('select * from "test" where "person" = ?')
     expect(p1).toStrictEqual(['{"name":"test"}'])
 
     console.log('   compiled:', `${(performance.now() - start2).toFixed(2)}ms`)
+
+    const result = await db.execCompiled(insert({ gender: true }))
+    expect(result!.rows).toStrictEqual([])
+    const result2 = await db.execCompiled(update({ gender: false }))
+    expect(result2!.rows).toStrictEqual([])
   })
 })