feat: select callback (#154)

Author: samwillis

.changeset/full-steaks-wait.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/db": patch
+---
+
+the live query select clause can now be a callback function that receives each row as a context object returning a new object with the selected fields. This also allows the for the callback to make more expressive changes to the returned data.

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

@@ -198,8 +198,9 @@ export class BaseQueryBuilder<TContext extends Context<Schema>> {
   /**
    * Specify what columns to select.
    * Overwrites any previous select clause.
+   * Also supports callback functions that receive the row context and return selected data.
    *
-   * @param selects The columns to select
+   * @param selects The columns to select (can include callbacks)
    * @returns A new QueryBuilder with the select clause set
    */
   select<TSelects extends Array<Select<TContext>>>(

packages/db/src/query/schema.ts

@@ -191,6 +191,11 @@ export type Select<TContext extends Context = Context> =
         | AggregateFunctionCall<TContext>
     }
   | WildcardReferenceString<TContext>
+  | SelectCallback<TContext>
+
+export type SelectCallback<TContext extends Context = Context> = (
+  context: TContext extends { schema: infer S } ? S : any
+) => any
 
 export type As<TContext extends Context = Context> = string
 

packages/db/src/query/select.ts

@@ -3,7 +3,7 @@ import {
   evaluateOperandOnNamespacedRow,
   extractValueFromNamespacedRow,
 } from "./extractors"
-import type { ConditionOperand, Query } from "./schema"
+import type { ConditionOperand, Query, SelectCallback } from "./schema"
 import type { KeyedStream, NamespacedAndKeyedStream } from "../types"
 
 export function processSelect(
@@ -31,6 +31,29 @@ export function processSelect(
       }
 
       for (const item of query.select) {
+        // Handle callback functions
+        if (typeof item === `function`) {
+          const callback = item as SelectCallback
+          const callbackResult = callback(namespacedRow)
+
+          // If the callback returns an object, merge its properties into the result
+          if (
+            callbackResult &&
+            typeof callbackResult === `object` &&
+            !Array.isArray(callbackResult)
+          ) {
+            Object.assign(result, callbackResult)
+          } else {
+            // If the callback returns a primitive value, we can't merge it
+            // This would need a specific key, but since we don't have one, we'll skip it
+            // In practice, select callbacks should return objects with keys
+            console.warn(
+              `SelectCallback returned a non-object value. SelectCallbacks should return objects with key-value pairs.`
+            )
+          }
+          continue
+        }
+
         if (typeof item === `string`) {
           // Handle wildcard select - all columns from all tables
           if ((item as string) === `@*`) {

packages/db/tests/query/conditions.test.ts

@@ -648,5 +648,124 @@ describe(`Query`, () => {
         expect(result.price).toBeLessThan(1000)
       })
     })
+
+    test(`select callback function`, () => {
+      const query: Query<Context> = {
+        select: [
+          ({ products }) => ({
+            displayName: `${products.name} (${products.category})`,
+            priceLevel: products.price > 500 ? `expensive` : `affordable`,
+            availability: products.inStock ? `in-stock` : `out-of-stock`,
+          }),
+        ],
+        from: `products`,
+        where: [[`@id`, `<=`, 3]], // First three products
+      }
+
+      const graph = new D2({ initialFrontier: v([0, 0]) })
+      const input = graph.newInput<[number, Product]>()
+      const pipeline = compileQueryPipeline(query, { [query.from]: input })
+
+      const messages: Array<Message<any>> = []
+      pipeline.pipe(
+        output((message) => {
+          messages.push(message)
+        })
+      )
+
+      graph.finalize()
+
+      input.sendData(
+        v([1, 0]),
+        new MultiSet(
+          sampleProducts.map((product) => [[product.id, product], 1])
+        )
+      )
+      input.sendFrontier(new Antichain([v([1, 0])]))
+
+      graph.run()
+
+      // Check the transformed results
+      const dataMessages = messages.filter((m) => m.type === MessageType.DATA)
+      const results = dataMessages[0]!.data.collection
+        .getInner()
+        .map(([data]) => data)
+
+      expect(results).toHaveLength(3) // First three products
+
+      // Verify the callback transformation
+      results.forEach(([_key, result]) => {
+        expect(result).toHaveProperty(`displayName`)
+        expect(result).toHaveProperty(`priceLevel`)
+        expect(result).toHaveProperty(`availability`)
+        expect(typeof result.displayName).toBe(`string`)
+        expect([`expensive`, `affordable`]).toContain(result.priceLevel)
+        expect([`in-stock`, `out-of-stock`]).toContain(result.availability)
+      })
+
+      // Check specific transformations for known products
+      const laptop = results.find(([_key, r]) =>
+        r.displayName.includes(`Laptop`)
+      )
+      expect(laptop).toBeDefined()
+      expect(laptop![1].priceLevel).toBe(`expensive`)
+      expect(laptop![1].availability).toBe(`in-stock`)
+    })
+
+    test(`mixed select: traditional columns and callback`, () => {
+      const query: Query<Context> = {
+        select: [
+          `@id`,
+          `@name`,
+          ({ products }) => ({
+            computedField: `${products.name}_computed`,
+            doublePrice: products.price * 2,
+          }),
+        ],
+        from: `products`,
+        where: [[`@id`, `=`, 1]], // Just the laptop
+      }
+
+      const graph = new D2({ initialFrontier: v([0, 0]) })
+      const input = graph.newInput<[number, Product]>()
+      const pipeline = compileQueryPipeline(query, { [query.from]: input })
+
+      const messages: Array<Message<any>> = []
+      pipeline.pipe(
+        output((message) => {
+          messages.push(message)
+        })
+      )
+
+      graph.finalize()
+
+      input.sendData(
+        v([1, 0]),
+        new MultiSet(
+          sampleProducts.map((product) => [[product.id, product], 1])
+        )
+      )
+      input.sendFrontier(new Antichain([v([1, 0])]))
+
+      graph.run()
+
+      // Check the mixed results
+      const dataMessages = messages.filter((m) => m.type === MessageType.DATA)
+      const results = dataMessages[0]!.data.collection
+        .getInner()
+        .map(([data]) => data)
+
+      expect(results).toHaveLength(1)
+
+      const [_key, result] = results[0]!
+
+      // Check traditional columns
+      expect(result.id).toBe(1)
+      expect(result.name).toBe(`Laptop`)
+
+      // Check callback-generated fields
+      expect(result.computedField).toBe(`Laptop_computed`)
+      expect(result.doublePrice).toBe(2400) // 1200 * 2
+    })
   })
 })

packages/db/tests/query/query-builder/select.test.ts

@@ -39,8 +39,8 @@ describe(`QueryBuilder.select`, () => {
 
     const builtQuery = query._query
     expect(builtQuery.select).toHaveLength(2)
-    expect(builtQuery.select[0]).toBe(`@id`)
-    expect(builtQuery.select[1]).toHaveProperty(`employee_name`, `@name`)
+    expect(builtQuery.select![0]).toBe(`@id`)
+    expect(builtQuery.select![1]).toHaveProperty(`employee_name`, `@name`)
   })
 
   it(`handles function calls`, () => {
@@ -52,7 +52,7 @@ describe(`QueryBuilder.select`, () => {
 
     const builtQuery = query._query
     expect(builtQuery.select).toHaveLength(2)
-    expect(builtQuery.select[1]).toHaveProperty(`upper_name`)
+    expect(builtQuery.select![1]).toHaveProperty(`upper_name`)
   })
 
   it(`overrides previous select calls`, () => {
@@ -85,4 +85,59 @@ describe(`QueryBuilder.select`, () => {
     const builtQuery = query._query
     expect(builtQuery.select).toEqual([`@id`, `@name`])
   })
+
+  it(`supports callback functions`, () => {
+    const callback = ({ employees }: any) => ({
+      fullInfo: `${employees.name} (ID: ${employees.id})`,
+      salaryLevel: employees.salary > 50000 ? `high` : `low`,
+    })
+
+    const query = queryBuilder<TestSchema>().from(`employees`).select(callback)
+
+    const builtQuery = query._query
+    expect(builtQuery.select).toHaveLength(1)
+    expect(builtQuery.select).toBeDefined()
+    expect(typeof builtQuery.select![0]).toBe(`function`)
+    expect(builtQuery.select![0]).toBe(callback)
+  })
+
+  it(`combines callback with traditional selects`, () => {
+    const callback = ({ employees }: any) => ({
+      computed: employees.salary * 1.1,
+    })
+
+    const query = queryBuilder<TestSchema>()
+      .from(`employees`)
+      .select(`@id`, `@name`, callback, { department_name: `@employees.name` })
+
+    const builtQuery = query._query
+    expect(builtQuery.select).toHaveLength(4)
+    expect(builtQuery.select).toBeDefined()
+    expect(builtQuery.select![0]).toBe(`@id`)
+    expect(builtQuery.select![1]).toBe(`@name`)
+    expect(typeof builtQuery.select![2]).toBe(`function`)
+    expect(builtQuery.select![3]).toHaveProperty(`department_name`)
+  })
+
+  it(`supports multiple callback functions`, () => {
+    const callback1 = ({ employees }: any) => ({
+      displayName: employees.name.toUpperCase(),
+    })
+    const callback2 = ({ employees }: any) => ({
+      isActive: employees.active,
+      experience: new Date().getFullYear() - 2020,
+    })
+
+    const query = queryBuilder<TestSchema>()
+      .from(`employees`)
+      .select(callback1, callback2)
+
+    const builtQuery = query._query
+    expect(builtQuery.select).toHaveLength(2)
+    expect(builtQuery.select).toBeDefined()
+    expect(typeof builtQuery.select![0]).toBe(`function`)
+    expect(typeof builtQuery.select![1]).toBe(`function`)
+    expect(builtQuery.select![0]).toBe(callback1)
+    expect(builtQuery.select![1]).toBe(callback2)
+  })
 })