add findOne() to query builder (#440)

* add findOne() to query builder

* move utility type to db/types

* fix utility type

* add single support to aditional useLiveQuery overloads

* Add support for specifying `.findOne()` outside of a query builder

* WIP addressing of the review

* fix for review

* remove unused offset

* restore text comments

* restore text comments

* disable limit=1 for future enforcing

---------

Co-authored-by: Sam Willis <[email protected]>

Author: fcannizzaro

docs/reference/classes/basequerybuilder.md

@@ -200,6 +200,32 @@ query
 
 ***
 
+### findOne()
+
+```ts
+findOne(): QueryBuilder<TContext>
+```
+
+Specify that the query should return a single row as `data` and not an array.
+
+#### Returns
+
+[`QueryBuilder`](../../type-aliases/querybuilder.md)\<`TContext`\>
+
+A QueryBuilder with single return enabled
+
+#### Example
+
+```ts
+// Get an user by ID
+query
+  .from({ users: usersCollection })
+  .where(({users}) => eq(users.id, 1))
+  .findOne()
+```
+
+***
+
 ### from()
 
 ```ts

packages/db/src/collection/index.ts

@@ -24,7 +24,9 @@ import type {
   InferSchemaInput,
   InferSchemaOutput,
   InsertConfig,
+  NonSingleResult,
   OperationConfig,
+  SingleResult,
   SubscribeChangesOptions,
   Transaction as TransactionType,
   UtilsRecord,
@@ -50,6 +52,7 @@ export interface Collection<
   TInsertInput extends object = T,
 > extends CollectionImpl<T, TKey, TUtils, TSchema, TInsertInput> {
   readonly utils: TUtils
+  readonly singleResult?: true
 }
 
 /**
@@ -132,8 +135,22 @@ export function createCollection<
   options: CollectionConfig<InferSchemaOutput<T>, TKey, T> & {
     schema: T
     utils?: TUtils
-  }
-): Collection<InferSchemaOutput<T>, TKey, TUtils, T, InferSchemaInput<T>>
+  } & NonSingleResult
+): Collection<InferSchemaOutput<T>, TKey, TUtils, T, InferSchemaInput<T>> &
+  NonSingleResult
+
+// Overload for when schema is provided and singleResult is true
+export function createCollection<
+  T extends StandardSchemaV1,
+  TKey extends string | number = string | number,
+  TUtils extends UtilsRecord = {},
+>(
+  options: CollectionConfig<InferSchemaOutput<T>, TKey, T> & {
+    schema: T
+    utils?: TUtils
+  } & SingleResult
+): Collection<InferSchemaOutput<T>, TKey, TUtils, T, InferSchemaInput<T>> &
+  SingleResult
 
 // Overload for when no schema is provided
 // the type T needs to be passed explicitly unless it can be inferred from the getKey function in the config
@@ -145,8 +162,21 @@ export function createCollection<
   options: CollectionConfig<T, TKey, never> & {
     schema?: never // prohibit schema if an explicit type is provided
     utils?: TUtils
-  }
-): Collection<T, TKey, TUtils, never, T>
+  } & NonSingleResult
+): Collection<T, TKey, TUtils, never, T> & NonSingleResult
+
+// Overload for when no schema is provided and singleResult is true
+// the type T needs to be passed explicitly unless it can be inferred from the getKey function in the config
+export function createCollection<
+  T extends object,
+  TKey extends string | number = string | number,
+  TUtils extends UtilsRecord = {},
+>(
+  options: CollectionConfig<T, TKey, never> & {
+    schema?: never // prohibit schema if an explicit type is provided
+    utils?: TUtils
+  } & SingleResult
+): Collection<T, TKey, TUtils, never, T> & SingleResult
 
 // Implementation
 export function createCollection(

packages/db/src/query/builder/index.ts

@@ -16,7 +16,7 @@ import {
   SubQueryMustHaveFromClauseError,
 } from "../../errors.js"
 import { createRefProxy, toExpression } from "./ref-proxy.js"
-import type { NamespacedRow } from "../../types.js"
+import type { NamespacedRow, SingleResult } from "../../types.js"
 import type {
   Aggregate,
   BasicExpression,
@@ -615,6 +615,28 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
     }) as any
   }
 
+  /**
+   * Specify that the query should return a single result
+   * @returns A QueryBuilder that returns the first result
+   *
+   * @example
+   * ```ts
+   * // Get the user matching the query
+   * query
+   *   .from({ users: usersCollection })
+   *   .where(({users}) => eq(users.id, 1))
+   *   .findOne()
+   *```
+   */
+  findOne(): QueryBuilder<TContext & SingleResult> {
+    return new BaseQueryBuilder({
+      ...this.query,
+      // TODO: enforcing return only one result with also a default orderBy if none is specified
+      // limit: 1,
+      singleResult: true,
+    })
+  }
+
   // Helper methods
   private _getCurrentAliases(): Array<string> {
     const aliases: Array<string> = []
@@ -817,4 +839,10 @@ export type ExtractContext<T> =
       : never
 
 // Export the types from types.ts for convenience
-export type { Context, Source, GetResult, RefLeaf as Ref } from "./types.js"
+export type {
+  Context,
+  Source,
+  GetResult,
+  RefLeaf as Ref,
+  InferResultType,
+} from "./types.js"

packages/db/src/query/builder/types.ts

@@ -1,4 +1,5 @@
 import type { CollectionImpl } from "../../collection/index.js"
+import type { SingleResult } from "../../types.js"
 import type {
   Aggregate,
   BasicExpression,
@@ -47,6 +48,8 @@ export interface Context {
   >
   // The result type after select (if select has been called)
   result?: any
+  // Single result only (if findOne has been called)
+  singleResult?: boolean
 }
 
 /**
@@ -571,6 +574,7 @@ export type MergeContextWithJoinType<
     [K in keyof TNewSchema & string]: TJoinType
   }
   result: TContext[`result`]
+  singleResult: TContext[`singleResult`] extends true ? true : false
 }
 
 /**
@@ -621,6 +625,14 @@ export type ApplyJoinOptionalityToMergedSchema<
       TNewSchema[K]
 }
 
+/**
+ * Utility type to infer the query result size (single row or an array)
+ */
+export type InferResultType<TContext extends Context> =
+  TContext extends SingleResult
+    ? GetResult<TContext> | undefined
+    : Array<GetResult<TContext>>
+
 /**
  * GetResult - Determines the final result type of a query
  *

packages/db/src/query/index.ts

@@ -9,6 +9,7 @@ export {
   type Context,
   type Source,
   type GetResult,
+  type InferResultType,
 } from "./builder/index.js"
 
 // Expression functions exports

packages/db/src/query/ir.ts

@@ -17,6 +17,7 @@ export interface QueryIR {
   limit?: Limit
   offset?: Offset
   distinct?: true
+  singleResult?: true
 
   // Functional variants
   fnSelect?: (row: NamespacedRow) => any

packages/db/src/query/live-query-collection.ts

@@ -3,9 +3,29 @@ import { CollectionConfigBuilder } from "./live/collection-config-builder.js"
 import type { LiveQueryCollectionConfig } from "./live/types.js"
 import type { InitialQueryBuilder, QueryBuilder } from "./builder/index.js"
 import type { Collection } from "../collection/index.js"
-import type { CollectionConfig, UtilsRecord } from "../types.js"
+import type {
+  CollectionConfig,
+  CollectionConfigSingleRowOption,
+  NonSingleResult,
+  SingleResult,
+  UtilsRecord,
+} from "../types.js"
 import type { Context, GetResult } from "./builder/types.js"
 
+type CollectionConfigForContext<
+  TContext extends Context,
+  TResult extends object,
+> = TContext extends SingleResult
+  ? CollectionConfigSingleRowOption<TResult> & SingleResult
+  : CollectionConfigSingleRowOption<TResult> & NonSingleResult
+
+type CollectionForContext<
+  TContext extends Context,
+  TResult extends object,
+> = TContext extends SingleResult
+  ? Collection<TResult> & SingleResult
+  : Collection<TResult> & NonSingleResult
+
 /**
  * Creates live query collection options for use with createCollection
  *
@@ -35,12 +55,15 @@ export function liveQueryCollectionOptions<
   TResult extends object = GetResult<TContext>,
 >(
   config: LiveQueryCollectionConfig<TContext, TResult>
-): CollectionConfig<TResult> {
+): CollectionConfigForContext<TContext, TResult> {
   const collectionConfigBuilder = new CollectionConfigBuilder<
     TContext,
     TResult
   >(config)
-  return collectionConfigBuilder.getConfig()
+  return collectionConfigBuilder.getConfig() as CollectionConfigForContext<
+    TContext,
+    TResult
+  >
 }
 
 /**
@@ -83,7 +106,7 @@ export function createLiveQueryCollection<
   TResult extends object = GetResult<TContext>,
 >(
   query: (q: InitialQueryBuilder) => QueryBuilder<TContext>
-): Collection<TResult, string | number, {}>
+): CollectionForContext<TContext, TResult>
 
 // Overload 2: Accept full config object with optional utilities
 export function createLiveQueryCollection<
@@ -92,7 +115,7 @@ export function createLiveQueryCollection<
   TUtils extends UtilsRecord = {},
 >(
   config: LiveQueryCollectionConfig<TContext, TResult> & { utils?: TUtils }
-): Collection<TResult, string | number, TUtils>
+): CollectionForContext<TContext, TResult>
 
 // Implementation
 export function createLiveQueryCollection<
@@ -103,7 +126,7 @@ export function createLiveQueryCollection<
   configOrQuery:
     | (LiveQueryCollectionConfig<TContext, TResult> & { utils?: TUtils })
     | ((q: InitialQueryBuilder) => QueryBuilder<TContext>)
-): Collection<TResult, string | number, TUtils> {
+): CollectionForContext<TContext, TResult> {
   // Determine if the argument is a function (query) or a config object
   if (typeof configOrQuery === `function`) {
     // Simple query function case
@@ -113,7 +136,10 @@ export function createLiveQueryCollection<
       ) => QueryBuilder<TContext>,
     }
     const options = liveQueryCollectionOptions<TContext, TResult>(config)
-    return bridgeToCreateCollection(options)
+    return bridgeToCreateCollection(options) as CollectionForContext<
+      TContext,
+      TResult
+    >
   } else {
     // Config object case
     const config = configOrQuery as LiveQueryCollectionConfig<
@@ -124,7 +150,7 @@ export function createLiveQueryCollection<
     return bridgeToCreateCollection({
       ...options,
       utils: config.utils,
-    })
+    }) as CollectionForContext<TContext, TResult>
   }
 }
 

packages/db/src/query/live/collection-config-builder.ts

@@ -7,7 +7,7 @@ import type { RootStreamBuilder } from "@tanstack/db-ivm"
 import type { OrderByOptimizationInfo } from "../compiler/order-by.js"
 import type { Collection } from "../../collection/index.js"
 import type {
-  CollectionConfig,
+  CollectionConfigSingleRowOption,
   KeyedStream,
   ResultStream,
   SyncConfig,
@@ -79,7 +79,7 @@ export class CollectionConfigBuilder<
     this.compileBasePipeline()
   }
 
-  getConfig(): CollectionConfig<TResult> {
+  getConfig(): CollectionConfigSingleRowOption<TResult> {
     return {
       id: this.id,
       getKey:
@@ -93,6 +93,7 @@ export class CollectionConfigBuilder<
       onUpdate: this.config.onUpdate,
       onDelete: this.config.onDelete,
       startSync: this.config.startSync,
+      singleResult: this.query.singleResult,
     }
   }
 

packages/db/src/query/live/types.ts

@@ -90,4 +90,9 @@ export interface LiveQueryCollectionConfig<
    * GC time for the collection
    */
   gcTime?: number
+
+  /**
+   * If enabled the collection will return a single object instead of an array
+   */
+  singleResult?: true
 }

packages/db/src/types.ts

@@ -503,6 +503,28 @@ export interface CollectionConfig<
   sync: SyncConfig<T, TKey>
 }
 
+export type SingleResult = {
+  singleResult: true
+}
+
+export type NonSingleResult = {
+  singleResult?: never
+}
+
+export type MaybeSingleResult = {
+  /**
+   * If enabled the collection will return a single object instead of an array
+   */
+  singleResult?: true
+}
+
+// Only used for live query collections
+export type CollectionConfigSingleRowOption<
+  T extends object = Record<string, unknown>,
+  TKey extends string | number = string | number,
+  TSchema extends StandardSchemaV1 = never,
+> = CollectionConfig<T, TKey, TSchema> & MaybeSingleResult
+
 export type ChangesPayload<T extends object = Record<string, unknown>> = Array<
   ChangeMessage<T>
 >