@tanstack/optimistic/src/query/query-builder.ts

Core optimistic updates library

import type { Collection } from "../collection"
import type {
  Comparator,
  Condition,
  From,
  JoinClause,
  Limit,
  LiteralValue,
  Offset,
  OrderBy,
  Query,
  Select,
  WithQuery,
} from "./schema.js"
import type {
  Context,
  Flatten,
  InferResultTypeFromSelectTuple,
  Input,
  InputReference,
  PropertyReference,
  PropertyReferenceString,
  RemoveIndexSignature,
  Schema,
} from "./types.js"

type CollectionRef = { [K: string]: Collection<any> }

export class BaseQueryBuilder<TContext extends Context<Schema>> {
  private readonly query: Partial<Query<TContext>> = {}

  /**
   * Create a new QueryBuilder instance.
   */
  constructor(query: Partial<Query<TContext>> = {}) {
    this.query = query
  }

  from<TCollectionRef extends CollectionRef>(
    collectionRef: TCollectionRef
  ): QueryBuilder<{
    baseSchema: Flatten<
      TContext[`baseSchema`] & {
        [K in keyof TCollectionRef & string]: RemoveIndexSignature<
          (TCollectionRef[keyof TCollectionRef] extends Collection<infer T>
            ? T
            : never) &
            Input
        >
      }
    >
    schema: Flatten<{
      [K in keyof TCollectionRef & string]: RemoveIndexSignature<
        (TCollectionRef[keyof TCollectionRef] extends Collection<infer T>
          ? T
          : never) &
          Input
      >
    }>
    default: keyof TCollectionRef & string
  }>

  from<
    T extends InputReference<{
      baseSchema: TContext[`baseSchema`]
      schema: TContext[`baseSchema`]
    }>,
  >(
    collection: T
  ): QueryBuilder<{
    baseSchema: TContext[`baseSchema`]
    schema: {
      [K in T]: RemoveIndexSignature<TContext[`baseSchema`][T]>
    }
    default: T
  }>

  from<
    T extends InputReference<{
      baseSchema: TContext[`baseSchema`]
      schema: TContext[`baseSchema`]
    }>,
    TAs extends string,
  >(
    collection: T,
    as: TAs
  ): QueryBuilder<{
    baseSchema: TContext[`baseSchema`]
    schema: {
      [K in TAs]: RemoveIndexSignature<TContext[`baseSchema`][T]>
    }
    default: TAs
  }>

  /**
   * Specify the collection to query from.
   * This is the first method that must be called in the chain.
   *
   * @param collection The collection name to query from
   * @param as Optional alias for the collection
   * @returns A new QueryBuilder with the from clause set
   */
  from<
    T extends
      | InputReference<{
          baseSchema: TContext[`baseSchema`]
          schema: TContext[`baseSchema`]
        }>
      | CollectionRef,
    TAs extends string | undefined,
  >(collection: T, as?: TAs) {
    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
    if (typeof collection === `object` && collection !== null) {
      return this.fromCollectionRef(collection)
    } else if (typeof collection === `string`) {
      return this.fromInputReference(
        collection as InputReference<{
          baseSchema: TContext[`baseSchema`]
          schema: TContext[`baseSchema`]
        }>,
        as
      )
    } else {
      throw new Error(`Invalid collection type`)
    }
  }

  private fromCollectionRef<TCollectionRef extends CollectionRef>(
    collectionRef: TCollectionRef
  ) {
    const keys = Object.keys(collectionRef)
    if (keys.length !== 1) {
      throw new Error(`Expected exactly one key`)
    }

    const key = keys[0]!
    const collection = collectionRef[key]!

    const newBuilder = new BaseQueryBuilder()
    Object.assign(newBuilder.query, this.query)
    newBuilder.query.from = key as From<TContext>
    newBuilder.query.collections ??= {}
    newBuilder.query.collections[key] = collection

    return newBuilder as unknown as QueryBuilder<{
      baseSchema: TContext[`baseSchema`] & {
        [K in keyof TCollectionRef &
          string]: (TCollectionRef[keyof TCollectionRef] extends Collection<
          infer T
        >
          ? T
          : never) &
          Input
      }
      schema: {
        [K in keyof TCollectionRef &
          string]: (TCollectionRef[keyof TCollectionRef] extends Collection<
          infer T
        >
          ? T
          : never) &
          Input
      }
      default: keyof TCollectionRef & string
    }>
  }

  private fromInputReference<
    T extends InputReference<{
      baseSchema: TContext[`baseSchema`]
      schema: TContext[`baseSchema`]
    }>,
    TAs extends string | undefined,
  >(collection: T, as?: TAs) {
    const newBuilder = new BaseQueryBuilder()
    Object.assign(newBuilder.query, this.query)
    newBuilder.query.from = collection as From<TContext>
    if (as) {
      newBuilder.query.as = as
    }

    // Calculate the result type without deep nesting
    type ResultSchema = TAs extends undefined
      ? { [K in T]: TContext[`baseSchema`][T] }
      : { [K in string & TAs]: TContext[`baseSchema`][T] }

    type ResultDefault = TAs extends undefined ? T : string & TAs

    // Use simpler type assertion to avoid excessive depth
    return newBuilder as unknown as QueryBuilder<{
      baseSchema: TContext[`baseSchema`]
      schema: ResultSchema
      default: ResultDefault
    }>
  }

  /**
   * Specify what columns to select.
   * Overwrites any previous select clause.
   *
   * @param selects The columns to select
   * @returns A new QueryBuilder with the select clause set
   */
  select<TSelects extends Array<Select<TContext>>>(
    this: QueryBuilder<TContext>,
    ...selects: TSelects
  ) {
    // Validate function calls in the selects
    // Need to use a type assertion to bypass deep recursive type checking
    const validatedSelects = selects.map((select) => {
      // If the select is an object with aliases, validate each value
      if (
        typeof select === `object` &&
        // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
        select !== null &&
        !Array.isArray(select)
      ) {
        const result: Record<string, any> = {}

        for (const [key, value] of Object.entries(select)) {
          // If it's a function call (object with a single key that is an allowed function name)
          if (
            typeof value === `object` &&
            value !== null &&
            !Array.isArray(value)
          ) {
            const keys = Object.keys(value)
            if (keys.length === 1) {
              const funcName = keys[0]!
              // List of allowed function names from AllowedFunctionName
              const allowedFunctions = [
                `SUM`,
                `COUNT`,
                `AVG`,
                `MIN`,
                `MAX`,
                `DATE`,
                `JSON_EXTRACT`,
                `JSON_EXTRACT_PATH`,
                `UPPER`,
                `LOWER`,
                `COALESCE`,
                `CONCAT`,
                `LENGTH`,
                `ORDER_INDEX`,
              ]

              if (!allowedFunctions.includes(funcName)) {
                console.warn(
                  `Unsupported function: ${funcName}. Expected one of: ${allowedFunctions.join(`, `)}`
                )
              }
            }
          }

          result[key] = value
        }

        return result
      }

      return select
    })

    const newBuilder = new BaseQueryBuilder<TContext>(
      (this as BaseQueryBuilder<TContext>).query
    )
    newBuilder.query.select = validatedSelects as Array<Select<TContext>>

    return newBuilder as QueryBuilder<
      Flatten<
        Omit<TContext, `result`> & {
          result: InferResultTypeFromSelectTuple<TContext, TSelects>
        }
      >
    >
  }

  /**
   * Add a where clause comparing two values.
   */
  where(
    left: PropertyReferenceString<TContext> | LiteralValue,
    operator: Comparator,
    right: PropertyReferenceString<TContext> | LiteralValue
  ): QueryBuilder<TContext>

  /**
   * Add a where clause with a complete condition object.
   */
  where(condition: Condition<TContext>): QueryBuilder<TContext>

  /**
   * Add a where clause to filter the results.
   * Can be called multiple times to add AND conditions.
   *
   * @param leftOrCondition The left operand or complete condition
   * @param operator Optional comparison operator
   * @param right Optional right operand
   * @returns A new QueryBuilder with the where clause added
   */
  where(
    leftOrCondition: any,
    operator?: any,
    right?: any
  ): QueryBuilder<TContext> {
    // Create a new builder with a copy of the current query
    // Use simplistic approach to avoid deep type errors
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    let condition: any

    // Determine if this is a complete condition or individual parts
    if (operator !== undefined && right !== undefined) {
      // Create a condition from parts
      condition = [leftOrCondition, operator, right]
    } else {
      // Use the provided condition directly
      condition = leftOrCondition
    }

    if (!newBuilder.query.where) {
      newBuilder.query.where = condition
    } else {
      // Create a composite condition with AND
      // Use any to bypass type checking issues
      const andArray: any = [newBuilder.query.where, `and`, condition]
      newBuilder.query.where = andArray
    }

    return newBuilder as unknown as QueryBuilder<TContext>
  }

  /**
   * Add a having clause comparing two values.
   * For filtering results after they have been grouped.
   */
  having(
    left: PropertyReferenceString<TContext> | LiteralValue,
    operator: Comparator,
    right: PropertyReferenceString<TContext> | LiteralValue
  ): QueryBuilder<TContext>

  /**
   * Add a having clause with a complete condition object.
   * For filtering results after they have been grouped.
   */
  having(condition: Condition<TContext>): QueryBuilder<TContext>

  /**
   * Add a having clause to filter the grouped results.
   * Can be called multiple times to add AND conditions.
   *
   * @param leftOrCondition The left operand or complete condition
   * @param operator Optional comparison operator
   * @param right Optional right operand
   * @returns A new QueryBuilder with the having clause added
   */
  having(
    leftOrCondition: any,
    operator?: any,
    right?: any
  ): QueryBuilder<TContext> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    let condition: any

    // Determine if this is a complete condition or individual parts
    if (operator !== undefined && right !== undefined) {
      // Create a condition from parts
      condition = [leftOrCondition, operator, right]
    } else {
      // Use the provided condition directly
      condition = leftOrCondition
    }

    if (!newBuilder.query.having) {
      newBuilder.query.having = condition
    } else {
      // Create a composite condition with AND
      // Use any to bypass type checking issues
      const andArray: any = [newBuilder.query.having, `and`, condition]
      newBuilder.query.having = andArray
    }

    return newBuilder as QueryBuilder<TContext>
  }

  /**
   * Add a join clause to the query using a CollectionRef.
   */
  join<TCollectionRef extends CollectionRef>(joinClause: {
    type: `inner` | `left` | `right` | `full` | `cross`
    from: TCollectionRef
    on: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: TContext[`schema`] & {
          [K in keyof TCollectionRef & string]: RemoveIndexSignature<
            (TCollectionRef[keyof TCollectionRef] extends Collection<infer T>
              ? T
              : never) &
              Input
          >
        }
      }>
    >
    where?: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: {
          [K in keyof TCollectionRef & string]: RemoveIndexSignature<
            (TCollectionRef[keyof TCollectionRef] extends Collection<infer T>
              ? T
              : never) &
              Input
          >
        }
      }>
    >
  }): QueryBuilder<
    Flatten<
      Omit<TContext, `schema`> & {
        schema: TContext[`schema`] & {
          [K in keyof TCollectionRef & string]: RemoveIndexSignature<
            (TCollectionRef[keyof TCollectionRef] extends Collection<infer T>
              ? T
              : never) &
              Input
          >
        }
      }
    >
  >

  /**
   * Add a join clause to the query without specifying an alias.
   * The collection name will be used as the default alias.
   */
  join<
    T extends InputReference<{
      baseSchema: TContext[`baseSchema`]
      schema: TContext[`baseSchema`]
    }>,
  >(joinClause: {
    type: `inner` | `left` | `right` | `full` | `cross`
    from: T
    on: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: TContext[`schema`] & {
          [K in T]: RemoveIndexSignature<TContext[`baseSchema`][T]>
        }
      }>
    >
    where?: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: { [K in T]: RemoveIndexSignature<TContext[`baseSchema`][T]> }
      }>
    >
  }): QueryBuilder<
    Flatten<
      Omit<TContext, `schema`> & {
        schema: TContext[`schema`] & {
          [K in T]: RemoveIndexSignature<TContext[`baseSchema`][T]>
        }
      }
    >
  >

  /**
   * Add a join clause to the query with a specified alias.
   */
  join<
    TFrom extends InputReference<{
      baseSchema: TContext[`baseSchema`]
      schema: TContext[`baseSchema`]
    }>,
    TAs extends string,
  >(joinClause: {
    type: `inner` | `left` | `right` | `full` | `cross`
    from: TFrom
    as: TAs
    on: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: TContext[`schema`] & {
          [K in TAs]: RemoveIndexSignature<TContext[`baseSchema`][TFrom]>
        }
      }>
    >
    where?: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: {
          [K in TAs]: RemoveIndexSignature<TContext[`baseSchema`][TFrom]>
        }
      }>
    >
  }): QueryBuilder<
    Flatten<
      Omit<TContext, `schema`> & {
        schema: TContext[`schema`] & {
          [K in TAs]: RemoveIndexSignature<TContext[`baseSchema`][TFrom]>
        }
      }
    >
  >

  join<
    TFrom extends
      | InputReference<{
          baseSchema: TContext[`baseSchema`]
          schema: TContext[`baseSchema`]
        }>
      | CollectionRef,
    TAs extends string | undefined = undefined,
  >(joinClause: {
    type: `inner` | `left` | `right` | `full` | `cross`
    from: TFrom
    as?: TAs
    on: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: TContext[`schema`] &
          (TFrom extends CollectionRef
            ? {
                [K in keyof TFrom & string]: RemoveIndexSignature<
                  (TFrom[keyof TFrom] extends Collection<infer T> ? T : never) &
                    Input
                >
              }
            : TFrom extends InputReference<infer TRef>
              ? {
                  [K in keyof TRef & string]: RemoveIndexSignature<
                    TRef[keyof TRef]
                  >
                }
              : never)
      }>
    >
    where?: Condition<
      Flatten<{
        baseSchema: TContext[`baseSchema`]
        schema: TContext[`schema`] &
          (TFrom extends CollectionRef
            ? {
                [K in keyof TFrom & string]: RemoveIndexSignature<
                  (TFrom[keyof TFrom] extends Collection<infer T> ? T : never) &
                    Input
                >
              }
            : TFrom extends InputReference<infer TRef>
              ? {
                  [K in keyof TRef & string]: RemoveIndexSignature<
                    TRef[keyof TRef]
                  >
                }
              : never)
      }>
    >
  }): QueryBuilder<any> {
    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
    if (typeof joinClause.from === `object` && joinClause.from !== null) {
      return this.joinCollectionRef(
        joinClause as {
          type: `inner` | `left` | `right` | `full` | `cross`
          from: CollectionRef
          on: Condition<any>
          where?: Condition<any>
        }
      )
    } else {
      return this.joinInputReference(
        joinClause as {
          type: `inner` | `left` | `right` | `full` | `cross`
          from: InputReference<{
            baseSchema: TContext[`baseSchema`]
            schema: TContext[`baseSchema`]
          }>
          as?: TAs
          on: Condition<any>
          where?: Condition<any>
        }
      )
    }
  }

  private joinCollectionRef<TCollectionRef extends CollectionRef>(joinClause: {
    type: `inner` | `left` | `right` | `full` | `cross`
    from: TCollectionRef
    on: Condition<any>
    where?: Condition<any>
  }): QueryBuilder<any> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Get the collection key
    const keys = Object.keys(joinClause.from)
    if (keys.length !== 1) {
      throw new Error(`Expected exactly one key in CollectionRef`)
    }
    const key = keys[0]!
    const collection = joinClause.from[key]
    if (!collection) {
      throw new Error(`Collection not found for key: ${key}`)
    }

    // Create a copy of the join clause for the query
    const joinClauseCopy = {
      type: joinClause.type,
      from: key,
      on: joinClause.on,
      where: joinClause.where,
    } as JoinClause<TContext>

    // Add the join clause to the query
    if (!newBuilder.query.join) {
      newBuilder.query.join = [joinClauseCopy]
    } else {
      newBuilder.query.join = [...newBuilder.query.join, joinClauseCopy]
    }

    // Add the collection to the collections map
    newBuilder.query.collections ??= {}
    newBuilder.query.collections[key] = collection

    // Return the new builder with updated schema type
    return newBuilder as QueryBuilder<
      Flatten<
        Omit<TContext, `schema`> & {
          schema: TContext[`schema`] & {
            [K in keyof TCollectionRef & string]: RemoveIndexSignature<
              (TCollectionRef[keyof TCollectionRef] extends Collection<infer T>
                ? T
                : never) &
                Input
            >
          }
        }
      >
    >
  }

  private joinInputReference<
    TFrom extends InputReference<{
      baseSchema: TContext[`baseSchema`]
      schema: TContext[`baseSchema`]
    }>,
    TAs extends string | undefined = undefined,
  >(joinClause: {
    type: `inner` | `left` | `right` | `full` | `cross`
    from: TFrom
    as?: TAs
    on: Condition<any>
    where?: Condition<any>
  }): QueryBuilder<any> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Create a copy of the join clause for the query
    const joinClauseCopy = { ...joinClause } as JoinClause<TContext>

    // Add the join clause to the query
    if (!newBuilder.query.join) {
      newBuilder.query.join = [joinClauseCopy]
    } else {
      newBuilder.query.join = [...newBuilder.query.join, joinClauseCopy]
    }

    // Determine the alias or use the collection name as default
    const _effectiveAlias = joinClause.as ?? joinClause.from

    // Return the new builder with updated schema type
    return newBuilder as QueryBuilder<
      Flatten<
        Omit<TContext, `schema`> & {
          schema: TContext[`schema`] & {
            [K in typeof _effectiveAlias]: TContext[`baseSchema`][TFrom]
          }
        }
      >
    >
  }

  /**
   * Add an orderBy clause to sort the results.
   * Overwrites any previous orderBy clause.
   *
   * @param orderBy The order specification
   * @returns A new QueryBuilder with the orderBy clause set
   */
  orderBy(orderBy: OrderBy<TContext>): QueryBuilder<TContext> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Set the orderBy clause
    newBuilder.query.orderBy = orderBy

    return newBuilder as QueryBuilder<TContext>
  }

  /**
   * Set a limit on the number of results returned.
   *
   * @param limit Maximum number of results to return
   * @returns A new QueryBuilder with the limit set
   */
  limit(limit: Limit<TContext>): QueryBuilder<TContext> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Set the limit
    newBuilder.query.limit = limit

    return newBuilder as QueryBuilder<TContext>
  }

  /**
   * Set an offset to skip a number of results.
   *
   * @param offset Number of results to skip
   * @returns A new QueryBuilder with the offset set
   */
  offset(offset: Offset<TContext>): QueryBuilder<TContext> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Set the offset
    newBuilder.query.offset = offset

    return newBuilder as QueryBuilder<TContext>
  }

  /**
   * Specify which column(s) to use as keys in the output keyed stream.
   *
   * @param keyBy The column(s) to use as keys
   * @returns A new QueryBuilder with the keyBy clause set
   */
  keyBy(
    keyBy: PropertyReference<TContext> | Array<PropertyReference<TContext>>
  ): QueryBuilder<TContext> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Set the keyBy clause
    newBuilder.query.keyBy = keyBy

    return newBuilder as QueryBuilder<TContext>
  }

  /**
   * Add a groupBy clause to group the results by one or more columns.
   *
   * @param groupBy The column(s) to group by
   * @returns A new QueryBuilder with the groupBy clause set
   */
  groupBy(
    groupBy: PropertyReference<TContext> | Array<PropertyReference<TContext>>
  ): QueryBuilder<TContext> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Set the groupBy clause
    newBuilder.query.groupBy = groupBy

    return newBuilder as QueryBuilder<TContext>
  }

  /**
   * Define a Common Table Expression (CTE) that can be referenced in the main query.
   * This allows referencing the CTE by name in subsequent from/join clauses.
   *
   * @param name The name of the CTE
   * @param queryBuilderCallback A function that builds the CTE query
   * @returns A new QueryBuilder with the CTE added
   */
  with<TName extends string, TResult = Record<string, unknown>>(
    name: TName,
    queryBuilderCallback: (
      builder: InitialQueryBuilder<{
        baseSchema: TContext[`baseSchema`]
        schema: {}
      }>
    ) => QueryBuilder<any>
  ): InitialQueryBuilder<{
    baseSchema: TContext[`baseSchema`] & { [K in TName]: TResult }
    schema: TContext[`schema`]
  }> {
    // Create a new builder with a copy of the current query
    const newBuilder = new BaseQueryBuilder<TContext>()
    Object.assign(newBuilder.query, this.query)

    // Create a new builder for the CTE
    const cteBuilder = new BaseQueryBuilder<{
      baseSchema: TContext[`baseSchema`]
      schema: {}
    }>()

    // Get the CTE query from the callback
    const cteQueryBuilder = queryBuilderCallback(
      cteBuilder as InitialQueryBuilder<{
        baseSchema: TContext[`baseSchema`]
        schema: {}
      }>
    )

    // Get the query from the builder
    const cteQuery = cteQueryBuilder._query

    // Add an 'as' property to the CTE
    const withQuery: WithQuery<any> = {
      ...cteQuery,
      as: name,
    }

    // Add the CTE to the with array
    if (!newBuilder.query.with) {
      newBuilder.query.with = [withQuery]
    } else {
      newBuilder.query.with = [...newBuilder.query.with, withQuery]
    }

    // Use a type cast that simplifies the type structure to avoid recursion
    return newBuilder as unknown as InitialQueryBuilder<{
      baseSchema: TContext[`baseSchema`] & { [K in TName]: TResult }
      schema: TContext[`schema`]
    }>
  }

  get _query(): Query<TContext> {
    return this.query as Query<TContext>
  }
}

export type InitialQueryBuilder<TContext extends Context<Schema>> = Pick<
  BaseQueryBuilder<TContext>,
  `from` | `with`
>

export type QueryBuilder<TContext extends Context<Schema>> = Omit<
  BaseQueryBuilder<TContext>,
  `from`
>

/**
 * Create a new query builder with the given schema
 */
export function queryBuilder<TBaseSchema extends Schema = {}>() {
  return new BaseQueryBuilder<{
    baseSchema: TBaseSchema
    schema: {}
  }>() as InitialQueryBuilder<{
    baseSchema: TBaseSchema
    schema: {}
  }>
}

export type ResultsFromContext<TContext extends Context<Schema>> = Flatten<
  TContext[`result`] extends object
    ? TContext[`result`]
    : TContext[`result`] extends undefined
      ? TContext[`schema`]
      : object
>

export type ResultFromQueryBuilder<TQueryBuilder> = Flatten<
  TQueryBuilder extends QueryBuilder<infer C>
    ? C extends { result: infer R }
      ? R
      : never
    : never
>