@tanstack/db 0.1.1 → 0.1.4

package/src/query/builder/types.ts

@@ -1,6 +1,7 @@
 import type { CollectionImpl } from "../../collection.js"
-import type { Aggregate, BasicExpression } from "../ir.js"
+import type { Aggregate, BasicExpression, OrderByDirection } from "../ir.js"
 import type { QueryBuilder } from "./index.js"
+import type { ResolveType } from "../../types.js"
 
 export interface Context {
   // The collections available in the base schema
@@ -27,13 +28,16 @@ export type Source = {
 }
 
 // Helper type to infer collection type from CollectionImpl
+// This uses ResolveType directly to ensure consistency with collection creation logic
 export type InferCollectionType<T> =
-  T extends CollectionImpl<infer U> ? U : never
+  T extends CollectionImpl<infer U, any, any, infer TSchema, any>
+    ? ResolveType<U, TSchema, U>
+    : never
 
 // Helper type to create schema from source
 export type SchemaFromSource<T extends Source> = Prettify<{
-  [K in keyof T]: T[K] extends CollectionImpl<infer U>
-    ? U
+  [K in keyof T]: T[K] extends CollectionImpl<any, any, any, any, any>
+    ? InferCollectionType<T[K]>
     : T[K] extends QueryBuilder<infer TContext>
       ? GetResult<TContext>
       : never
@@ -58,16 +62,18 @@ export type SelectObject<
 // Helper type to get the result type from a select object
 export type ResultTypeFromSelect<TSelectObject> = {
   [K in keyof TSelectObject]: TSelectObject[K] extends RefProxy<infer T>
-    ? // For RefProxy, preserve the type as-is (including optionality from joins)
-      T
+    ? T
     : TSelectObject[K] extends BasicExpression<infer T>
       ? T
       : TSelectObject[K] extends Aggregate<infer T>
         ? T
         : TSelectObject[K] extends RefProxyFor<infer T>
-          ? // For RefProxyFor, preserve the type as-is (including optionality from joins)
-            T
-          : never
+          ? T
+          : TSelectObject[K] extends undefined
+            ? undefined
+            : TSelectObject[K] extends { __type: infer U }
+              ? U
+              : never
 }
 
 // Callback type for orderBy clauses
@@ -75,6 +81,29 @@ export type OrderByCallback<TContext extends Context> = (
   refs: RefProxyForContext<TContext>
 ) => any
 
+export type OrderByOptions = {
+  direction?: OrderByDirection
+  nulls?: `first` | `last`
+} & StringSortOpts
+
+export type StringSortOpts =
+  | {
+      stringSort?: `lexical`
+    }
+  | {
+      stringSort?: `locale`
+      locale?: string
+      localeOptions?: object
+    }
+
+export type CompareOptions = {
+  direction: OrderByDirection
+  nulls: `first` | `last`
+  stringSort: `lexical` | `locale`
+  locale?: string
+  localeOptions?: object
+}
+
 // Callback type for groupBy clauses
 export type GroupByCallback<TContext extends Context> = (
   refs: RefProxyForContext<TContext>
@@ -119,12 +148,11 @@ export type RefProxyFor<T> = OmitRefProxy<
       ? // T is optional (T | undefined) but not exactly undefined
         NonUndefined<T> extends Record<string, any>
         ? {
-            // Properties are accessible and their types become optional
-            [K in keyof NonUndefined<T>]: NonUndefined<T>[K] extends Record<
+            [K in keyof NonUndefined<T>]-?: NonUndefined<T>[K] extends Record<
               string,
               any
             >
-              ? RefProxyFor<NonUndefined<T>[K] | undefined> &
+              ? RefProxyFor<NonUndefined<T>[K]> &
                   RefProxy<NonUndefined<T>[K] | undefined>
               : RefProxy<NonUndefined<T>[K] | undefined>
           } & RefProxy<T>
@@ -132,9 +160,14 @@ export type RefProxyFor<T> = OmitRefProxy<
       : // T is not optional
         T extends Record<string, any>
         ? {
-            [K in keyof T]: T[K] extends Record<string, any>
-              ? RefProxyFor<T[K]> & RefProxy<T[K]>
-              : RefProxy<T[K]>
+            // Make all properties required, but for optional ones, include undefined in the RefProxy type
+            [K in keyof T]-?: undefined extends T[K]
+              ? T[K] extends Record<string, any>
+                ? RefProxyFor<T[K]> & RefProxy<T[K]>
+                : RefProxy<T[K]>
+              : T[K] extends Record<string, any>
+                ? RefProxyFor<T[K]> & RefProxy<T[K]>
+                : RefProxy<T[K]>
           } & RefProxy<T>
         : RefProxy<T>
 >

package/src/query/compiler/evaluators.ts

@@ -20,9 +20,12 @@ export type CompiledSingleRowExpression = (item: Record<string, unknown>) => any
  * Compiles an expression into an optimized evaluator function.
  * This eliminates branching during evaluation by pre-compiling the expression structure.
  */
-export function compileExpression(expr: BasicExpression): CompiledExpression {
-  const compiledFn = compileExpressionInternal(expr, false)
-  return compiledFn as CompiledExpression
+export function compileExpression(
+  expr: BasicExpression,
+  isSingleRow: boolean = false
+): CompiledExpression | CompiledSingleRowExpression {
+  const compiledFn = compileExpressionInternal(expr, isSingleRow)
+  return compiledFn
 }
 
 /**

package/src/query/compiler/group-by.ts

@@ -166,7 +166,9 @@ export function processGroupBy(
   const mapping = validateAndCreateMapping(groupByClause, selectClause)
 
   // Pre-compile groupBy expressions
-  const compiledGroupByExpressions = groupByClause.map(compileExpression)
+  const compiledGroupByExpressions = groupByClause.map((e) =>
+    compileExpression(e)
+  )
 
   // Create a key extractor function using simple __key_X format
   const keyExtractor = ([, row]: [

package/src/query/compiler/index.ts

@@ -7,17 +7,21 @@ import {
   LimitOffsetRequireOrderByError,
   UnsupportedFromTypeError,
 } from "../../errors.js"
+import { PropRef } from "../ir.js"
 import { compileExpression } from "./evaluators.js"
 import { processJoins } from "./joins.js"
 import { processGroupBy } from "./group-by.js"
 import { processOrderBy } from "./order-by.js"
 import { processSelectToResults } from "./select.js"
+import type { OrderByOptimizationInfo } from "./order-by.js"
 import type {
   BasicExpression,
   CollectionRef,
   QueryIR,
   QueryRef,
 } from "../ir.js"
+import type { LazyCollectionCallbacks } from "./joins.js"
+import type { Collection } from "../../collection.js"
 import type {
   KeyedStream,
   NamespacedAndKeyedStream,
@@ -29,6 +33,8 @@ import type { QueryCache, QueryMapping } from "./types.js"
  * Result of query compilation including both the pipeline and collection-specific WHERE clauses
  */
 export interface CompilationResult {
+  /** The ID of the main collection */
+  collectionId: string
   /** The compiled query pipeline */
   pipeline: ResultStream
   /** Map of collection aliases to their WHERE clauses for index optimization */
@@ -46,6 +52,10 @@ export interface CompilationResult {
 export function compileQuery(
   rawQuery: QueryIR,
   inputs: Record<string, KeyedStream>,
+  collections: Record<string, Collection<any, any, any, any, any>>,
+  callbacks: Record<string, LazyCollectionCallbacks>,
+  lazyCollections: Set<string>,
+  optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   cache: QueryCache = new WeakMap(),
   queryMapping: QueryMapping = new WeakMap()
 ): CompilationResult {
@@ -70,9 +80,17 @@ export function compileQuery(
   const tables: Record<string, KeyedStream> = {}
 
   // Process the FROM clause to get the main table
-  const { alias: mainTableAlias, input: mainInput } = processFrom(
+  const {
+    alias: mainTableAlias,
+    input: mainInput,
+    collectionId: mainCollectionId,
+  } = processFrom(
     query.from,
     allInputs,
+    collections,
+    callbacks,
+    lazyCollections,
+    optimizableOrderByCollections,
     cache,
     queryMapping
   )
@@ -96,10 +114,16 @@ export function compileQuery(
       pipeline,
       query.join,
       tables,
+      mainCollectionId,
       mainTableAlias,
       allInputs,
       cache,
-      queryMapping
+      queryMapping,
+      collections,
+      callbacks,
+      lazyCollections,
+      optimizableOrderByCollections,
+      rawQuery
     )
   }
 
@@ -231,8 +255,11 @@ export function compileQuery(
   // Process orderBy parameter if it exists
   if (query.orderBy && query.orderBy.length > 0) {
     const orderedPipeline = processOrderBy(
+      rawQuery,
       pipeline,
       query.orderBy,
+      collections[mainCollectionId]!,
+      optimizableOrderByCollections,
       query.limit,
       query.offset
     )
@@ -249,6 +276,7 @@ export function compileQuery(
     const result = resultPipeline
     // Cache the result before returning (use original query as key)
     const compilationResult = {
+      collectionId: mainCollectionId,
       pipeline: result,
       collectionWhereClauses,
     }
@@ -275,6 +303,7 @@ export function compileQuery(
   const result = resultPipeline
   // Cache the result before returning (use original query as key)
   const compilationResult = {
+    collectionId: mainCollectionId,
     pipeline: result,
     collectionWhereClauses,
   }
@@ -289,16 +318,20 @@ export function compileQuery(
 function processFrom(
   from: CollectionRef | QueryRef,
   allInputs: Record<string, KeyedStream>,
+  collections: Record<string, Collection>,
+  callbacks: Record<string, LazyCollectionCallbacks>,
+  lazyCollections: Set<string>,
+  optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   cache: QueryCache,
   queryMapping: QueryMapping
-): { alias: string; input: KeyedStream } {
+): { alias: string; input: KeyedStream; collectionId: string } {
   switch (from.type) {
     case `collectionRef`: {
       const input = allInputs[from.collection.id]
       if (!input) {
         throw new CollectionInputNotFoundError(from.collection.id)
       }
-      return { alias: from.alias, input }
+      return { alias: from.alias, input, collectionId: from.collection.id }
     }
     case `queryRef`: {
       // Find the original query for caching purposes
@@ -308,6 +341,10 @@ function processFrom(
       const subQueryResult = compileQuery(
         originalQuery,
         allInputs,
+        collections,
+        callbacks,
+        lazyCollections,
+        optimizableOrderByCollections,
         cache,
         queryMapping
       )
@@ -324,7 +361,11 @@ function processFrom(
         })
       )
 
-      return { alias: from.alias, input: extractedInput }
+      return {
+        alias: from.alias,
+        input: extractedInput,
+        collectionId: subQueryResult.collectionId,
+      }
     }
     default:
       throw new UnsupportedFromTypeError((from as any).type)
@@ -380,3 +421,69 @@ function mapNestedQueries(
     }
   }
 }
+
+function getRefFromAlias(
+  query: QueryIR,
+  alias: string
+): CollectionRef | QueryRef | void {
+  if (query.from.alias === alias) {
+    return query.from
+  }
+
+  for (const join of query.join || []) {
+    if (join.from.alias === alias) {
+      return join.from
+    }
+  }
+}
+
+/**
+ * Follows the given reference in a query
+ * until its finds the root field the reference points to.
+ * @returns The collection, its alias, and the path to the root field in this collection
+ */
+export function followRef(
+  query: QueryIR,
+  ref: PropRef<any>,
+  collection: Collection
+): { collection: Collection; path: Array<string> } | void {
+  if (ref.path.length === 0) {
+    return
+  }
+
+  if (ref.path.length === 1) {
+    // This field should be part of this collection
+    const field = ref.path[0]!
+    // is it part of the select clause?
+    if (query.select) {
+      const selectedField = query.select[field]
+      if (selectedField && selectedField.type === `ref`) {
+        return followRef(query, selectedField, collection)
+      }
+    }
+
+    // Either this field is not part of the select clause
+    // and thus it must be part of the collection itself
+    // or it is part of the select but is not a reference
+    // so we can stop here and don't have to follow it
+    return { collection, path: [field] }
+  }
+
+  if (ref.path.length > 1) {
+    // This is a nested field
+    const [alias, ...rest] = ref.path
+    const aliasRef = getRefFromAlias(query, alias!)
+    if (!aliasRef) {
+      return
+    }
+
+    if (aliasRef.type === `queryRef`) {
+      return followRef(aliasRef.query, new PropRef(rest), collection)
+    } else {
+      // This is a reference to a collection
+      // we can't follow it further
+      // so the field must be on the collection itself
+      return { collection: aliasRef.collection, path: rest }
+    }
+  }
+}