Fix type of utils in collection options

Author: kevin-dp

packages/db/src/collection/index.ts

@@ -127,32 +127,85 @@ export interface Collection<
  *
  */
 
-// Overload for when schema is provided
+// Overload for when schema is provided and utils is required (not optional)
+// We can't infer the Utils type from the CollectionConfig because it will always be optional
+// So we omit it from that type and instead infer it from the extension `& { utils: TUtils }`
+// such that we have the real, non-optional Utils type
 export function createCollection<
   T extends StandardSchemaV1,
-  TKey extends string | number = string | number,
-  TUtils extends UtilsRecord = UtilsRecord,
+  TKey extends string | number,
+  TUtils extends UtilsRecord,
 >(
-  options: CollectionConfig<InferSchemaOutput<T>, TKey, T, TUtils> & {
+  options: Omit<
+    CollectionConfig<InferSchemaOutput<T>, TKey, T, TUtils>,
+    `utils`
+  > & {
     schema: T
-    utils?: TUtils
+    utils: TUtils // Required utils
   } & NonSingleResult
 ): Collection<InferSchemaOutput<T>, TKey, TUtils, T, InferSchemaInput<T>> &
   NonSingleResult
 
+// Overload for when schema is provided and utils is optional
+// In this case we can simply infer the Utils type from the CollectionConfig type
+export function createCollection<
+  T extends StandardSchemaV1,
+  TKey extends string | number,
+  TUtils extends UtilsRecord,
+>(
+  options: CollectionConfig<InferSchemaOutput<T>, TKey, T, TUtils> & {
+    schema: T
+  } & NonSingleResult
+): Collection<
+  InferSchemaOutput<T>,
+  TKey,
+  Exclude<TUtils, undefined>,
+  T,
+  InferSchemaInput<T>
+> &
+  NonSingleResult
+
+// Overload for when schema is provided, singleResult is true, and utils is required
+export function createCollection<
+  T extends StandardSchemaV1,
+  TKey extends string | number,
+  TUtils extends UtilsRecord,
+>(
+  options: Omit<
+    CollectionConfig<InferSchemaOutput<T>, TKey, T, TUtils>,
+    `utils`
+  > & {
+    schema: T
+    utils: TUtils // Required utils
+  } & SingleResult
+): Collection<InferSchemaOutput<T>, TKey, TUtils, T, InferSchemaInput<T>> &
+  SingleResult
+
 // 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 = UtilsRecord,
+  TKey extends string | number,
+  TUtils extends UtilsRecord,
 >(
   options: CollectionConfig<InferSchemaOutput<T>, TKey, T, TUtils> & {
     schema: T
-    utils?: TUtils
   } & SingleResult
 ): Collection<InferSchemaOutput<T>, TKey, TUtils, T, InferSchemaInput<T>> &
   SingleResult
 
+// Overload for when no schema is provided and utils is required
+// 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,
+  TUtils extends UtilsRecord,
+>(
+  options: Omit<CollectionConfig<T, TKey, never, TUtils>, `utils`> & {
+    schema?: never // prohibit schema if an explicit type is provided
+    utils: TUtils // Required utils
+  } & NonSingleResult
+): Collection<T, TKey, TUtils, never, T> & NonSingleResult
+
 // 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
 export function createCollection<
@@ -162,10 +215,22 @@ export function createCollection<
 >(
   options: CollectionConfig<T, TKey, never, TUtils> & {
     schema?: never // prohibit schema if an explicit type is provided
-    utils?: TUtils
   } & NonSingleResult
 ): Collection<T, TKey, TUtils, never, T> & NonSingleResult
 
+// Overload for when no schema is provided, singleResult is true, and utils is required
+// 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 = UtilsRecord,
+>(
+  options: Omit<CollectionConfig<T, TKey, never, TUtils>, `utils`> & {
+    schema?: never // prohibit schema if an explicit type is provided
+    utils: TUtils // Required utils
+  } & SingleResult
+): Collection<T, TKey, TUtils, never, T> & SingleResult
+
 // 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<
@@ -175,15 +240,13 @@ export function createCollection<
 >(
   options: CollectionConfig<T, TKey, never, TUtils> & {
     schema?: never // prohibit schema if an explicit type is provided
-    utils?: TUtils
   } & SingleResult
 ): Collection<T, TKey, TUtils, never, T> & SingleResult
 
 // Implementation
 export function createCollection(
-  options: CollectionConfig<any, string | number, any> & {
+  options: CollectionConfig<any, string | number, any, UtilsRecord> & {
     schema?: StandardSchemaV1
-    utils?: UtilsRecord
   }
 ): Collection<any, string | number, UtilsRecord, any, any> {
   const collection = new CollectionImpl<any, string | number, any, any, any>(

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

@@ -294,9 +294,9 @@ export function electricCollectionOptions<T extends StandardSchemaV1>(
   config: ElectricCollectionConfig<InferSchemaOutput<T>, T> & {
     schema: T
   }
-): CollectionConfig<InferSchemaOutput<T>, string | number, T> & {
+): Omit<CollectionConfig<InferSchemaOutput<T>, string | number, T>, `utils`> & {
   id?: string
-  utils: ElectricCollectionUtils
+  utils: ElectricCollectionUtils<InferSchemaOutput<T>>
   schema: T
 }
 
@@ -305,15 +305,15 @@ export function electricCollectionOptions<T extends Row<unknown>>(
   config: ElectricCollectionConfig<T> & {
     schema?: never // prohibit schema
   }
-): CollectionConfig<T, string | number> & {
+): Omit<CollectionConfig<T, string | number>, `utils`> & {
   id?: string
-  utils: ElectricCollectionUtils
+  utils: ElectricCollectionUtils<T>
   schema?: never // no schema in the result
 }
 
 export function electricCollectionOptions(
   config: ElectricCollectionConfig<any, any>
-): CollectionConfig<any, string | number, any> & {
+): Omit<CollectionConfig<any, string | number>, `utils`> & {
   id?: string
   utils: ElectricCollectionUtils
   schema?: any

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

@@ -8,7 +8,10 @@ import {
   gt,
 } from "@tanstack/db"
 import { electricCollectionOptions } from "../src/electric"
-import type { ElectricCollectionConfig } from "../src/electric"
+import type {
+  ElectricCollectionConfig,
+  ElectricCollectionUtils,
+} from "../src/electric"
 import type {
   DeleteMutationFnParams,
   InsertMutationFnParams,
@@ -97,6 +100,59 @@ describe(`Electric collection type resolution tests`, () => {
     expectTypeOf(options.getKey).parameters.toEqualTypeOf<[FallbackType]>()
   })
 
+  it(`should type collection.utils as ElectricCollectionUtils<T>`, () => {
+    const todoSchema = z.object({
+      id: z.string(),
+      title: z.string(),
+      completed: z.boolean(),
+    })
+
+    type TodoType = z.infer<typeof todoSchema>
+
+    const options = electricCollectionOptions({
+      id: `todos`,
+      getKey: (item) => item.id,
+      shapeOptions: {
+        url: `/api/todos`,
+        params: { table: `todos` },
+      },
+      schema: todoSchema,
+      /*
+      onInsert: async ({ collection }) => {
+        const testCollectionUtils: ElectricCollectionUtils<TodoType> =
+          collection.utils
+        expectTypeOf(testCollectionUtils.awaitTxId).toBeFunction
+        expectTypeOf(collection.utils.awaitTxId).toBeFunction
+        return Promise.resolve({ txid: 1 })
+      },
+      */
+    })
+
+    // ✅ Test that options.utils is typed as ElectricCollectionUtils<TodoType>
+    // The options object should have the correct type from electricCollectionOptions
+    const testOptionsUtils: ElectricCollectionUtils<TodoType> = options.utils
+
+    expectTypeOf(testOptionsUtils.awaitTxId).toBeFunction
+
+    const todosCollection = createCollection(options)
+
+    // Test that todosCollection.utils is ElectricCollectionUtils<TodoType>
+    // Note: We can't use expectTypeOf(...).toEqualTypeOf<ElectricCollectionUtils<T>> because
+    // expectTypeOf's toEqualTypeOf has a constraint that requires { [x: string]: any; [x: number]: never; },
+    // 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
+
+    expectTypeOf(testTodosUtils.awaitTxId).toBeFunction
+
+    // Verify the specific properties that define ElectricCollectionUtils exist and are functions
+    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: {