fix(db): preserve discriminated union types through .select() (#1511)

.changeset/fix-ref-union-collapse-b.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': patch
+---
+
+Fix `.select()` collapsing discriminated-union fields to the intersection of common keys (#1511). `Ref<T>` now distributes over `T` so `keyof (A | B | C)` no longer reduces the union to its common keys, and `ExtractRef<T>` now distinguishes a real branded `Ref` (where the underlying user type `U` can be returned directly) from a spread-produced inline object (which still needs to be projected through `ResultTypeFromSelect`). This preserves discriminated unions both when the field is selected at the top level and when the field is nested inside another selected object.

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

@@ -478,8 +478,44 @@ type ResultTypeFromCaseWhen<T> = T extends unknown
   ? ResultTypeFromSelectValue<T>
   : never
 
-// Extract Ref or subobject with a spread or a Ref
-type ExtractRef<T> = Prettify<ResultTypeFromSelect<WithoutRefBrand<T>>>
+// Extract Ref or subobject with a spread or a Ref.
+type ExtractRef<T> = T extends unknown
+  ? IsTrueRef<T> extends true
+    ? T extends RefLeaf<infer U>
+      ? IsNullableRef<T> extends true
+        ? DeepNullable<U>
+        : U
+      : never
+    : Prettify<ResultTypeFromSelect<WithoutRefBrand<T>>>
+  : never
+
+// A "true" Ref is one whose own keys are exactly the brand+virtual props
+// plus the keys of its underlying user type U (after distribution this means
+// no extra keys from a spread were merged in). When extra keys are present
+// (the case for spread-produced inline objects that pick up the RefBrand
+// symbol from spreading a Ref), we fall through to the recursive projection.
+type IsTrueRef<T> =
+  T extends RefLeaf<infer U>
+    ? [
+        Exclude<
+          keyof T,
+          | typeof RefBrand
+          | typeof NullableBrand
+          | keyof VirtualRowProps
+          | keyof U
+        >,
+      ] extends [never]
+      ? true
+      : false
+    : false
+
+// Propagate nullable-join semantics into the user-data shape.
+type DeepNullable<T> =
+  T extends Record<string, any>
+    ? IsPlainObject<T> extends true
+      ? { [K in keyof T]: DeepNullable<T[K]> }
+      : T | undefined
+    : T | undefined
 
 // Helper type to extract the underlying type from various expression types
 type ExtractExpressionType<T> =
@@ -770,7 +806,11 @@ type VirtualPropsRef<TKey extends string | number = string | number> = {
  * select(({ user }) => ({ ...user })) // Returns User type, not Ref types
  * ```
  */
-export type Ref<T = any, Nullable extends boolean = false> = {
+export type Ref<T = any, Nullable extends boolean = false> = T extends unknown
+  ? RefBranch<T, Nullable>
+  : never
+
+type RefBranch<T, Nullable extends boolean> = {
   [K in keyof T]: IsNonExactOptional<T[K]> extends true
     ? IsNonExactNullable<T[K]> extends true
       ? // Both optional and nullable

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

@@ -1,6 +1,6 @@
 import { describe, expectTypeOf, test } from 'vitest'
 import { createCollection } from '../../src/collection/index.js'
-import { createLiveQueryCollection } from '../../src/query/index.js'
+import { createLiveQueryCollection, eq } from '../../src/query/index.js'
 import { mockSyncCollectionOptions } from '../utils.js'
 import { upper } from '../../src/query/builder/functions.js'
 import type { OutputWithVirtual } from '../utils.js'
@@ -109,6 +109,66 @@ describe(`select types`, () => {
     expectTypeOf(results).toMatchTypeOf<OutputWithVirtualKeyed<Expected>>()
   })
 
+  test(`select preserves union types and where works on common keys`, () => {
+    type ItemDocument =
+      | { type: 'pdf'; url: string; pages: number }
+      | { type: 'image'; url: string; width: number; height: number }
+      | { type: 'legacy'; path: string }
+
+    type Item = { id: number; name: string; document: ItemDocument }
+
+    const items = createCollection(
+      mockSyncCollectionOptions<Item>({
+        id: `union-field-items`,
+        getKey: (i) => i.id,
+        initialData: [],
+      }),
+    )
+
+    // Filtering by a common key of the union should compile,
+    // and the result should preserve the full discriminated union
+    const col = createLiveQueryCollection((q) =>
+      q
+        .from({ i: items })
+        .where(({ i }) => eq(i.document.type, `pdf`))
+        .select(({ i }) => ({
+          id: i.id,
+          document: i.document,
+        })),
+    )
+
+    const result = col.toArray[0]!
+    expectTypeOf(result.document).toEqualTypeOf<ItemDocument>()
+  })
+
+  test(`select preserves union when nested under another field`, () => {
+    type Payload =
+      | { kind: 'text'; body: string }
+      | { kind: 'binary'; bytes: number; mime: string }
+
+    type Envelope = { id: number; payload: { inner: Payload } }
+
+    const envelopes = createCollection(
+      mockSyncCollectionOptions<Envelope>({
+        id: `nested-union-envelopes`,
+        getKey: (e) => e.id,
+        initialData: [],
+      }),
+    )
+
+    // Selecting a nested object whose field is a discriminated union
+    // must preserve the union (not collapse to the intersection of keys).
+    const col = createLiveQueryCollection((q) =>
+      q.from({ e: envelopes }).select(({ e }) => ({
+        id: e.id,
+        payload: e.payload,
+      })),
+    )
+    const r = col.toArray[0]!
+    expectTypeOf(r.payload).toEqualTypeOf<{ inner: Payload }>()
+    expectTypeOf(r.payload.inner).toEqualTypeOf<Payload>()
+  })
+
   test(`nested spread preserves object structure types`, () => {
     const users = createUsers()
     const col = createLiveQueryCollection((q) => {