feat: make selects optionally lazy (#101)

In a effort to make `workers-qb` more efficient for databases that can support cursors/iterables it now supports lazy selects so that, a query like this doesn't potentially OoM the worker: ```sql SELECT * FROM table ``` The API for this is backwards-compatible and you can enable lazy selects by specifying the lazy parameter: ```ts // this will now return a iterable instead of a list this.db .select('table') .execute({lazy: true}) ``` It also works for `.fetchAll` too ```ts // it will also return a iterable this.db .fetchAll({tableName: "table", lazy: true}) .execute() ``` Because of TS type-system, lazy must be statically known at compile otherwise this won't work properly and must always be true or undefined.
G4brym · Jan 14, 2025 · 99181b4 · 99181b4
1 parent d021135
commit 99181b4
Show file tree

Hide file tree

Showing 8 changed files with 157 additions and 16 deletions.
diff --git a/src/builder.ts b/src/builder.ts
@@ -61,11 +61,15 @@ export class QueryBuilder<GenericResultWrapper, IsAsync extends boolean = true>
     throw new Error('Batch execute method not implemented')
   }
 
+  lazyExecute(query: Query<any, IsAsync>): IsAsync extends true ? Promise<AsyncIterable<any>> : Iterable<any> {
+    throw new Error('Execute lazyExecute not implemented')
+  }
+
   createTable<GenericResult = undefined>(params: {
     tableName: string
     schema: string
     ifNotExists?: boolean
-  }): Query<ArrayResult<GenericResultWrapper, GenericResult>, IsAsync> {
+  }): Query<ArrayResult<GenericResultWrapper, GenericResult, IsAsync>, IsAsync> {
     return new Query(
       (q) => {
         return this.execute(q)
@@ -78,7 +82,7 @@ export class QueryBuilder<GenericResultWrapper, IsAsync extends boolean = true>
   dropTable<GenericResult = undefined>(params: {
     tableName: string
     ifExists?: boolean
-  }): Query<ArrayResult<GenericResultWrapper, GenericResult>, IsAsync> {
+  }): Query<ArrayResult<GenericResultWrapper, GenericResult, IsAsync>, IsAsync> {
     return new Query(
       (q) => {
         return this.execute(q)
@@ -127,20 +131,26 @@ export class QueryBuilder<GenericResultWrapper, IsAsync extends boolean = true>
     )
   }
 
-  fetchAll<GenericResult = DefaultReturnObject>(
-    params: SelectAll
-  ): QueryWithExtra<GenericResultWrapper, ArrayResult<GenericResultWrapper, GenericResult>, IsAsync> {
+  fetchAll<GenericResult = DefaultReturnObject, IsLazy extends true | undefined = undefined>(
+    params: SelectAll<IsLazy>
+  ): QueryWithExtra<GenericResultWrapper, ArrayResult<GenericResultWrapper, GenericResult, IsAsync, IsLazy>, IsAsync> {
     return new QueryWithExtra(
       (q) => {
-        return this.execute(q)
+        return params.lazy
+          ? (this.lazyExecute(q) as unknown as MaybeAsync<
+              IsAsync,
+              ArrayResult<GenericResultWrapper, GenericResult, IsAsync, IsLazy>
+            >)
+          : this.execute(q)
       },
-      this._select(params),
+      this._select({ ...params, lazy: undefined }),
       this._select({
         ...params,
         fields: 'count(*) as total',
         offset: undefined,
         groupBy: undefined,
         limit: 1,
+        lazy: undefined,
       }),
       typeof params.where === 'object' && !Array.isArray(params.where) && params.where?.params
         ? Array.isArray(params.where?.params)
@@ -156,7 +166,7 @@ export class QueryBuilder<GenericResultWrapper, IsAsync extends boolean = true>
   ): Query<OneResult<GenericResultWrapper, GenericResult>, IsAsync>
   raw<GenericResult = DefaultReturnObject>(
     params: RawQueryFetchAll
-  ): Query<ArrayResult<GenericResultWrapper, GenericResult>, IsAsync>
+  ): Query<ArrayResult<GenericResultWrapper, GenericResult, IsAsync>, IsAsync>
   raw<GenericResult = DefaultReturnObject>(params: RawQueryWithoutFetching): Query<GenericResultWrapper, IsAsync>
   raw<GenericResult = DefaultReturnObject>(params: RawQuery): unknown {
     return new Query<any, IsAsync>(
@@ -174,7 +184,7 @@ export class QueryBuilder<GenericResultWrapper, IsAsync extends boolean = true>
   ): Query<OneResult<GenericResultWrapper, GenericResult>, IsAsync>
   insert<GenericResult = DefaultReturnObject>(
     params: InsertMultiple
-  ): Query<ArrayResult<GenericResultWrapper, GenericResult>, IsAsync>
+  ): Query<ArrayResult<GenericResultWrapper, GenericResult, IsAsync>, IsAsync>
   insert<GenericResult = DefaultReturnObject>(params: InsertWithoutReturning): Query<GenericResultWrapper, IsAsync>
   insert<GenericResult = DefaultReturnObject>(params: Insert): unknown {
     let args: any[] = []
@@ -218,7 +228,7 @@ export class QueryBuilder<GenericResultWrapper, IsAsync extends boolean = true>
 
   update<GenericResult = DefaultReturnObject>(
     params: UpdateReturning
-  ): Query<ArrayResult<GenericResultWrapper, GenericResult>, IsAsync>
+  ): Query<ArrayResult<GenericResultWrapper, GenericResult, IsAsync>, IsAsync>
   update<GenericResult = DefaultReturnObject>(params: UpdateWithoutReturning): Query<GenericResultWrapper, IsAsync>
   update<GenericResult = DefaultReturnObject>(params: Update): unknown {
     let args = this._parse_arguments(params.data)
@@ -243,7 +253,7 @@ export class QueryBuilder<GenericResultWrapper, IsAsync extends boolean = true>
 
   delete<GenericResult = DefaultReturnObject>(
     params: DeleteReturning
-  ): Query<ArrayResult<GenericResultWrapper, GenericResult>, IsAsync>
+  ): Query<ArrayResult<GenericResultWrapper, GenericResult, IsAsync>, IsAsync>
   delete<GenericResult = DefaultReturnObject>(params: DeleteWithoutReturning): Query<GenericResultWrapper, IsAsync>
   delete<GenericResult = DefaultReturnObject>(params: Delete): unknown {
     return new Query<any, IsAsync>(

diff --git a/src/databases/do.ts b/src/databases/do.ts
@@ -40,4 +40,16 @@ export class DOQB extends QueryBuilder<{}, false> {
       }
     })
   }
+
+  lazyExecute(query: Query<any, false>): Iterable<any> {
+    return this.loggerWrapper(query, this.options.logger, () => {
+      let cursor
+      if (query.arguments) {
+        cursor = this.db.exec(query.query, ...query.arguments)
+      } else {
+        cursor = this.db.exec(query.query)
+      }
+      return cursor
+    })
+  }
 }
diff --git a/src/interfaces.ts b/src/interfaces.ts
@@ -58,8 +58,9 @@ export type RawQueryFetchAll = Omit<RawQuery, 'fetchType'> & {
 
 export type RawQueryWithoutFetching = Omit<RawQuery, 'fetchType'>
 
-export type SelectAll = SelectOne & {
+export type SelectAll<IsLazy extends true | undefined = undefined> = SelectOne & {
   limit?: number
+  lazy?: IsLazy
 }
 
 export type ConflictUpsert = {
@@ -144,7 +145,32 @@ export type PGResult = {
   rowCount: number
 }
 
-export type ArrayResult<ResultWrapper, Result> = Merge<ResultWrapper, { results?: Array<Result> }>
+export type IterableResult<
+  ResultWrapper,
+  Result,
+  IsAsync extends boolean,
+  IsLazy extends true | undefined = undefined,
+> = IsLazy extends true
+  ? IsAsync extends true
+    ? Promise<Merge<ResultWrapper, { results?: AsyncIterable<Result> }>>
+    : Merge<ResultWrapper, { results?: Iterable<Result> }>
+  : never
+
+export type FullArrayResult<
+  ResultWrapper,
+  Result,
+  IsAsync extends boolean,
+  IsLazy extends true | undefined = undefined,
+> = IsLazy extends undefined
+  ? IsAsync extends true
+    ? Promise<Merge<ResultWrapper, { results?: Array<Result> }>>
+    : Merge<ResultWrapper, { results?: Array<Result> }>
+  : never
+
+export type ArrayResult<ResultWrapper, Result, IsAsync extends boolean, IsLazy extends true | undefined = undefined> =
+  | IterableResult<ResultWrapper, Result, IsAsync, IsLazy>
+  | FullArrayResult<ResultWrapper, Result, IsAsync, IsLazy>
+
 export type OneResult<ResultWrapper, Result> = Merge<ResultWrapper, { results?: Result }>
 
 export type CountResult<GenericResultWrapper> = OneResult<GenericResultWrapper, { total: number }>

diff --git a/src/modularBuilder.ts b/src/modularBuilder.ts
@@ -10,6 +10,10 @@ import {
 } from './interfaces'
 import { Query, QueryWithExtra } from './tools'
 
+export interface SelectExecuteOptions<IsLazy extends true | undefined> {
+  lazy?: IsLazy
+}
+
 export class SelectBuilder<GenericResultWrapper, GenericResult = DefaultReturnObject, IsAsync extends boolean = true> {
   _debugger = false
   _options: Partial<SelectAll> = {}
@@ -203,19 +207,25 @@ export class SelectBuilder<GenericResultWrapper, GenericResult = DefaultReturnOb
     )
   }
 
-  getQueryAll(): Query<ArrayResult<GenericResultWrapper, GenericResult>, IsAsync> {
+  getQueryAll<IsLazy extends true | undefined = undefined>(
+    options?: SelectExecuteOptions<IsLazy>
+  ): Query<ArrayResult<GenericResultWrapper, GenericResult, IsAsync, IsLazy>, IsAsync> {
     return this._fetchAll(this._options as SelectAll)
   }
 
   getQueryOne(): Query<OneResult<GenericResultWrapper, GenericResult>, IsAsync> {
     return this._fetchOne(this._options as SelectAll)
   }
 
-  execute(): MaybeAsync<IsAsync, ArrayResult<GenericResultWrapper, GenericResult>> {
+  execute<IsLazy extends true | undefined = undefined>(
+    options?: SelectExecuteOptions<IsLazy>
+  ): ArrayResult<GenericResultWrapper, GenericResult, IsAsync, IsLazy> {
     return this._fetchAll(this._options as SelectAll).execute()
   }
 
-  all(): MaybeAsync<IsAsync, ArrayResult<GenericResultWrapper, GenericResult>> {
+  all<IsLazy extends true | undefined = undefined>(
+    options?: SelectExecuteOptions<IsLazy>
+  ): ArrayResult<GenericResultWrapper, GenericResult, IsAsync, IsLazy> {
     return this._fetchAll(this._options as SelectAll).execute()
   }
 

diff --git a/tests/tsconfig.json b/tests/tsconfig.json
@@ -1,6 +1,7 @@
 {
   "extends": "../tsconfig.json",
   "compilerOptions": {
+    "module": "ESNext",
     "moduleResolution": "bundler",
     "types": ["@cloudflare/workers-types/experimental", "@cloudflare/vitest-pool-workers"]
   },

diff --git a/tests/unit/select.test-d.ts b/tests/unit/select.test-d.ts
@@ -0,0 +1,41 @@
+import { describe, expectTypeOf, it } from 'vitest'
+import { DefaultReturnObject } from '../../src'
+import { QuerybuilderTest, QuerybuilderTestSync } from '../utils'
+
+describe('Select builder', () => {
+  it('should return a iterable if using async and lazy', async () => {
+    const results = [
+      (await new QuerybuilderTest().select('table').execute({ lazy: true })).results!,
+      (await new QuerybuilderTest().fetchAll({ tableName: 'table', lazy: true }).execute()).results!,
+    ]
+
+    results.forEach((result) => expectTypeOf(result).toEqualTypeOf<AsyncIterable<DefaultReturnObject>>())
+  })
+
+  it('should return a list if using async and not lazy', async () => {
+    const results = [
+      (await new QuerybuilderTest().select('table').execute()).results!,
+      (await new QuerybuilderTest().fetchAll({ tableName: 'table' }).execute()).results!,
+    ]
+
+    results.forEach((result) => expectTypeOf(result).toEqualTypeOf<DefaultReturnObject[]>())
+  })
+
+  it('should return a iterable if using sync and lazy', async () => {
+    const results = [
+      new QuerybuilderTestSync().select('table').execute({ lazy: true }).results!,
+      new QuerybuilderTestSync().fetchAll({ tableName: 'table', lazy: true }).execute().results!,
+    ]
+
+    results.forEach((result) => expectTypeOf(result).toEqualTypeOf<Iterable<DefaultReturnObject>>())
+  })
+
+  it('should return a list if using sync and not lazy', async () => {
+    const results = [
+      new QuerybuilderTestSync().select('table').execute().results!,
+      new QuerybuilderTestSync().fetchAll({ tableName: 'table' }).execute().results!,
+    ]
+
+    results.forEach((result) => expectTypeOf(result).toEqualTypeOf<DefaultReturnObject[]>())
+  })
+})
diff --git a/tests/utils.ts b/tests/utils.ts
@@ -1,3 +1,4 @@
+import { syncLoggerWrapper } from '../src'
 import { QueryBuilder } from '../src/builder'
 import { D1Result } from '../src/interfaces'
 import { Query } from '../src/tools'
@@ -15,6 +16,43 @@ export class QuerybuilderTest extends QueryBuilder<{}> {
       }
     })
   }
+
+  async lazyExecute(query: Query<any, true>): Promise<AsyncIterable<any>> {
+    return this.loggerWrapper(query, this.options.logger, async function* () {
+      yield {
+        query: query.query,
+        arguments: query.arguments,
+        fetchType: query.fetchType,
+      }
+    })
+  }
+}
+
+export class QuerybuilderTestSync extends QueryBuilder<{}, false> {
+  loggerWrapper = syncLoggerWrapper
+
+  execute(query: Query<any, false>) {
+    return this.loggerWrapper(query, this.options.logger, () => {
+      return {
+        // @ts-ignore
+        results: {
+          query: query.query,
+          arguments: query.arguments,
+          fetchType: query.fetchType,
+        },
+      }
+    })
+  }
+
+  lazyExecute(query: Query<any, false>) {
+    return this.loggerWrapper(query, this.options.logger, function* () {
+      yield {
+        query: query.query,
+        arguments: query.arguments,
+        fetchType: query.fetchType,
+      }
+    })
+  }
 }
 
 export class QuerybuilderExceptionTest extends QueryBuilder<{}> {

diff --git a/tests/vitest.config.ts b/tests/vitest.config.ts
@@ -2,6 +2,9 @@ import { defineWorkersConfig } from '@cloudflare/vitest-pool-workers/config'
 
 export default defineWorkersConfig({
   test: {
+    typecheck: {
+      enabled: true,
+    },
     poolOptions: {
       workers: {
         wrangler: {