Improve type of collection utils for Trailbase, PowerSync, and localOnly collections

.changeset/tough-dragons-fold.md

@@ -0,0 +1,8 @@
+---
+'@tanstack/powersync-db-collection': patch
+'@tanstack/trailbase-db-collection': patch
+'@tanstack/electric-db-collection': patch
+'@tanstack/db': patch
+---
+
+Make type of collection utils more precise for localOnly, PowerSync, Trailbase, and Electric collections

packages/db/src/local-only.ts

@@ -64,7 +64,7 @@ type LocalOnlyCollectionOptionsResult<
   T extends object,
   TKey extends string | number,
   TSchema extends StandardSchemaV1 | never = never,
-> = CollectionConfig<T, TKey, TSchema> & {
+> = CollectionConfig<T, TKey, TSchema, LocalOnlyCollectionUtils> & {
   utils: LocalOnlyCollectionUtils
 }
 

packages/db/tests/local-only.test-d.ts

@@ -2,6 +2,7 @@ import { describe, expectTypeOf, it } from 'vitest'
 import { z } from 'zod'
 import { createCollection } from '../src/index'
 import { localOnlyCollectionOptions } from '../src/local-only'
+import type { LocalOnlyCollectionUtils } from '../src/local-only'
 
 interface TestItem extends Record<string, unknown> {
   id: number
@@ -252,4 +253,18 @@ describe(`LocalOnly Collection Types`, () => {
     // Test that the collection has the correct inferred type from schema
     expectTypeOf(collection.toArray).toEqualTypeOf<Array<ExpectedType>>()
   })
+
+  it(`should type collection.utils as LocalOnlyCollectionUtils`, () => {
+    const collection = createCollection(
+      localOnlyCollectionOptions({
+        id: `test-utils-typing`,
+        getKey: (item: TestItem) => item.id,
+      }),
+    )
+
+    // Verify that collection.utils is typed as LocalOnlyCollectionUtils, not UtilsRecord
+    const utils: LocalOnlyCollectionUtils = collection.utils
+    expectTypeOf(utils.acceptMutations).toBeFunction()
+    expectTypeOf(collection.utils.acceptMutations).toBeFunction()
+  })
 })

packages/db/tests/local-storage.test-d.ts

@@ -4,6 +4,7 @@ import { createCollection } from '../src/index'
 import { localStorageCollectionOptions } from '../src/local-storage'
 import type {
   LocalStorageCollectionConfig,
+  LocalStorageCollectionUtils,
   StorageApi,
   StorageEventApi,
 } from '../src/local-storage'
@@ -394,4 +395,23 @@ describe(`LocalStorage collection type resolution tests`, () => {
       expectTypeOf(draft).toEqualTypeOf<SelectUrlType>()
     })
   })
+
+  it(`should type collection.utils as LocalStorageCollectionUtils after createCollection`, () => {
+    const collection = createCollection(
+      localStorageCollectionOptions<ExplicitType>({
+        storageKey: `test-utils-typing`,
+        storage: mockStorage,
+        storageEventApi: mockStorageEventApi,
+        getKey: (item) => item.id,
+      }),
+    )
+
+    // Verify that collection.utils is typed as LocalStorageCollectionUtils, not UtilsRecord
+    const utils: LocalStorageCollectionUtils = collection.utils
+    expectTypeOf(utils.clearStorage).toBeFunction()
+    expectTypeOf(utils.getStorageSize).toBeFunction()
+    expectTypeOf(utils.acceptMutations).toBeFunction()
+    expectTypeOf(collection.utils.clearStorage).toBeFunction()
+    expectTypeOf(collection.utils.getStorageSize).toBeFunction()
+  })
 })

packages/electric-db-collection/src/electric.ts

@@ -511,22 +511,33 @@ export function electricCollectionOptions<T extends StandardSchemaV1>(
   config: ElectricCollectionConfig<InferSchemaOutput<T>, T> & {
     schema: T
   },
-): Omit<CollectionConfig<InferSchemaOutput<T>, string | number, T>, `utils`> & {
-  id?: string
-  utils: ElectricCollectionUtils<InferSchemaOutput<T>>
-  schema: T
-}
+): Omit<
+  CollectionConfig<InferSchemaOutput<T>, string | number, T>,
+  `utils` | `onInsert` | `onUpdate` | `onDelete`
+> &
+  Pick<
+    ElectricCollectionConfig<InferSchemaOutput<T>, T>,
+    `onInsert` | `onUpdate` | `onDelete`
+  > & {
+    id?: string
+    utils: ElectricCollectionUtils<InferSchemaOutput<T>>
+    schema: T
+  }
 
 // Overload for when no schema is provided
 export function electricCollectionOptions<T extends Row<unknown>>(
   config: ElectricCollectionConfig<T> & {
     schema?: never // prohibit schema
   },
-): Omit<CollectionConfig<T, string | number>, `utils`> & {
-  id?: string
-  utils: ElectricCollectionUtils<T>
-  schema?: never // no schema in the result
-}
+): Omit<
+  CollectionConfig<T, string | number>,
+  `utils` | `onInsert` | `onUpdate` | `onDelete`
+> &
+  Pick<ElectricCollectionConfig<T>, `onInsert` | `onUpdate` | `onDelete`> & {
+    id?: string
+    utils: ElectricCollectionUtils<T>
+    schema?: never // no schema in the result
+  }
 
 export function electricCollectionOptions<T extends Row<unknown>>(
   config: ElectricCollectionConfig<T, any>,

packages/electric-db-collection/tests/electric.test-d.ts

@@ -154,7 +154,6 @@ describe(`Electric collection type resolution tests`, () => {
     // but ElectricCollectionUtils extends UtilsRecord which is Record<string, any> (no number index signature).
     // This causes a constraint error instead of a type mismatch error.
     // Instead, we test via type assignment which will show a proper type error if the types don't match.
-    // Currently this shows that todosCollection.utils is typed as UtilsRecord, not ElectricCollectionUtils<TodoType>
     const testTodosUtils: ElectricCollectionUtils<TodoType> =
       todosCollection.utils
 
@@ -165,6 +164,44 @@ describe(`Electric collection type resolution tests`, () => {
     expectTypeOf(todosCollection.utils.awaitMatch).toBeFunction
   })
 
+  it(`should preserve ElectricCollectionUtils type on collection.utils after createCollection with handlers`, () => {
+    const todoSchema = z.object({
+      id: z.string(),
+      title: z.string(),
+      completed: z.boolean(),
+    })
+
+    type TodoType = z.infer<typeof todoSchema>
+
+    const options = electricCollectionOptions({
+      shapeOptions: {
+        url: `/api/todos`,
+      },
+      schema: todoSchema,
+      getKey: (item) => item.id,
+      onInsert: async () => {
+        return Promise.resolve({ txid: 1 })
+      },
+      onUpdate: async () => {
+        return Promise.resolve({ txid: 1 })
+      },
+      onDelete: async () => {
+        return Promise.resolve({ txid: 1 })
+      },
+    })
+
+    const todosCollection = createCollection(options)
+
+    // After createCollection, utils should be typed as ElectricCollectionUtils<TodoType>
+    // and not widened to UtilsRecord
+    const testUtils: ElectricCollectionUtils<TodoType> = todosCollection.utils
+
+    expectTypeOf(testUtils.awaitTxId).toBeFunction
+    expectTypeOf(testUtils.awaitMatch).toBeFunction
+    expectTypeOf(todosCollection.utils.awaitTxId).toBeFunction
+    expectTypeOf(todosCollection.utils.awaitMatch).toBeFunction
+  })
+
   it(`should properly type the onInsert, onUpdate, and onDelete handlers`, () => {
     const options = electricCollectionOptions<ExplicitType>({
       shapeOptions: {
@@ -195,17 +232,35 @@ describe(`Electric collection type resolution tests`, () => {
       },
     })
 
-    // Verify that the handlers are properly typed
+    // Verify that the handlers are properly typed with ElectricCollectionUtils
     expectTypeOf(options.onInsert).parameters.toEqualTypeOf<
-      [InsertMutationFnParams<ExplicitType>]
+      [
+        InsertMutationFnParams<
+          ExplicitType,
+          string | number,
+          ElectricCollectionUtils<ExplicitType>
+        >,
+      ]
     >()
 
     expectTypeOf(options.onUpdate).parameters.toEqualTypeOf<
-      [UpdateMutationFnParams<ExplicitType>]
+      [
+        UpdateMutationFnParams<
+          ExplicitType,
+          string | number,
+          ElectricCollectionUtils<ExplicitType>
+        >,
+      ]
     >()
 
     expectTypeOf(options.onDelete).parameters.toEqualTypeOf<
-      [DeleteMutationFnParams<ExplicitType>]
+      [
+        DeleteMutationFnParams<
+          ExplicitType,
+          string | number,
+          ElectricCollectionUtils<ExplicitType>
+        >,
+      ]
     >()
   })
 
@@ -279,15 +334,33 @@ describe(`Electric collection type resolution tests`, () => {
       },
     })
 
-    // Verify that the handlers are properly typed
+    // Verify that the handlers are properly typed with ElectricCollectionUtils
     expectTypeOf(options.onDelete).parameters.toEqualTypeOf<
-      [DeleteMutationFnParams<TodoType>]
+      [
+        DeleteMutationFnParams<
+          TodoType,
+          string | number,
+          ElectricCollectionUtils<TodoType>
+        >,
+      ]
     >()
     expectTypeOf(options.onInsert).parameters.toEqualTypeOf<
-      [InsertMutationFnParams<TodoType>]
+      [
+        InsertMutationFnParams<
+          TodoType,
+          string | number,
+          ElectricCollectionUtils<TodoType>
+        >,
+      ]
     >()
     expectTypeOf(options.onUpdate).parameters.toEqualTypeOf<
-      [UpdateMutationFnParams<TodoType>]
+      [
+        UpdateMutationFnParams<
+          TodoType,
+          string | number,
+          ElectricCollectionUtils<TodoType>
+        >,
+      ]
     >()
   })
 

packages/electric-db-collection/tests/electric.test.ts

@@ -528,7 +528,11 @@ describe(`Electric Integration`, () => {
         id: `test-transaction`,
         mutations: [],
       } as unknown as TransactionWithMutations<Row, `insert`>
-      const mockParams: InsertMutationFnParams<Row> = {
+      const mockParams: InsertMutationFnParams<
+        Row,
+        string | number,
+        ElectricCollectionUtils<Row>
+      > = {
         transaction: mockTransaction,
         // @ts-expect-error not relevant to test
         collection: CollectionImpl,

packages/powersync-db-collection/src/definitions.ts

@@ -260,7 +260,12 @@ export type EnhancedPowerSyncCollectionConfig<
   TTable extends Table,
   OutputType extends Record<string, unknown> = Record<string, unknown>,
   TSchema extends StandardSchemaV1 = never,
-> = CollectionConfig<OutputType, string, TSchema> & {
+> = CollectionConfig<
+  OutputType,
+  string,
+  TSchema,
+  PowerSyncCollectionUtils<TTable>
+> & {
   id?: string
   utils: PowerSyncCollectionUtils<TTable>
   schema?: TSchema

packages/powersync-db-collection/tests/powersync.test-d.ts

@@ -0,0 +1,31 @@
+import { describe, expectTypeOf, it } from 'vitest'
+import { Schema, Table, column } from '@powersync/node'
+import { createCollection } from '@tanstack/db'
+import { powerSyncCollectionOptions } from '../src'
+import type { PowerSyncCollectionUtils } from '../src'
+import type { AbstractPowerSyncDatabase } from '@powersync/node'
+
+const APP_SCHEMA = new Schema({
+  documents: new Table({
+    name: column.text,
+    author: column.text,
+  }),
+})
+
+describe(`PowerSync collection type tests`, () => {
+  it(`should type collection.utils as PowerSyncCollectionUtils after createCollection`, () => {
+    const collection = createCollection(
+      powerSyncCollectionOptions({
+        database: {} as AbstractPowerSyncDatabase,
+        table: APP_SCHEMA.props.documents,
+      }),
+    )
+
+    // Verify that collection.utils is typed as PowerSyncCollectionUtils, not UtilsRecord
+    const utils: PowerSyncCollectionUtils<
+      (typeof APP_SCHEMA.props)['documents']
+    > = collection.utils
+    expectTypeOf(utils.getMeta).toBeFunction()
+    expectTypeOf(collection.utils.getMeta).toBeFunction()
+  })
+})

packages/powersync-db-collection/tests/powersync.test.ts

@@ -203,7 +203,9 @@ describe(`PowerSync Integration`, () => {
       const _crudEntries = await db.getAll(`
         SELECT * FROM ps_crud ORDER BY id`)
 
-      const crudEntries = _crudEntries.map((r) => CrudEntry.fromRow(r))
+      const crudEntries = _crudEntries.map((r) =>
+        CrudEntry.fromRow(r as Parameters<typeof CrudEntry.fromRow>[0]),
+      )
 
       expect(crudEntries.length).toBe(6)
       // We can only group transactions for similar operations
@@ -250,7 +252,9 @@ describe(`PowerSync Integration`, () => {
       // There should be a crud entries for this
       const _crudEntries = await db.getAll(`
         SELECT * FROM ps_crud ORDER BY id`)
-      const crudEntries = _crudEntries.map((r) => CrudEntry.fromRow(r))
+      const crudEntries = _crudEntries.map((r) =>
+        CrudEntry.fromRow(r as Parameters<typeof CrudEntry.fromRow>[0]),
+      )
 
       const lastTransactionId =
         crudEntries[crudEntries.length - 1]?.transactionId
@@ -312,7 +316,9 @@ describe(`PowerSync Integration`, () => {
       // There should be a crud entries for this
       const _crudEntries = await db.getAll(`
         SELECT * FROM ps_crud ORDER BY id`)
-      const crudEntries = _crudEntries.map((r) => CrudEntry.fromRow(r))
+      const crudEntries = _crudEntries.map((r) =>
+        CrudEntry.fromRow(r as Parameters<typeof CrudEntry.fromRow>[0]),
+      )
 
       const lastTransactionId =
         crudEntries[crudEntries.length - 1]?.transactionId
@@ -464,7 +470,7 @@ describe(`PowerSync Integration`, () => {
       liveDocuments.subscribeChanges((changes) => {
         changes
           .map((change) => change.value.name)
-          .forEach((change) => bookNames.add(change))
+          .forEach((change) => bookNames.add(change!))
       })
 
       await collection.insert({

packages/query-db-collection/tests/query.test-d.ts

@@ -560,4 +560,26 @@ describe(`Query collection type resolution tests`, () => {
       createCollection(options)
     })
   })
+
+  it(`should type collection.utils as QueryCollectionUtils after createCollection`, () => {
+    const collection = createCollection(
+      queryCollectionOptions<ExplicitType>({
+        id: `test-utils-typing`,
+        queryClient,
+        queryKey: [`test-utils`],
+        queryFn: () => Promise.resolve([]),
+        getKey: (item) => item.id,
+      }),
+    )
+
+    // Verify that collection.utils is typed as QueryCollectionUtils, not UtilsRecord
+    const utils: QueryCollectionUtils<ExplicitType> = collection.utils
+    expectTypeOf(utils.refetch).toBeFunction()
+    expectTypeOf(collection.utils.refetch).toBeFunction()
+    expectTypeOf(collection.utils.writeInsert).toBeFunction()
+    expectTypeOf(collection.utils.writeUpdate).toBeFunction()
+    expectTypeOf(collection.utils.writeDelete).toBeFunction()
+    expectTypeOf(collection.utils.isFetching).toBeBoolean()
+    expectTypeOf(collection.utils.isLoading).toBeBoolean()
+  })
 })

packages/trailbase-db-collection/src/trailbase.ts

@@ -113,7 +113,9 @@ export function trailBaseCollectionOptions<
   TKey extends string | number = string | number,
 >(
   config: TrailBaseCollectionConfig<TItem, TRecord, TKey>,
-): CollectionConfig<TItem, TKey> & { utils: TrailBaseCollectionUtils } {
+): CollectionConfig<TItem, TKey, never, TrailBaseCollectionUtils> & {
+  utils: TrailBaseCollectionUtils
+} {
   const getKey = config.getKey
 
   const parse = (record: TRecord) =>