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. The real-`Ref` detection uses a strict structural equivalence against the canonical `Ref<U>` shape, so spread-derived objects that keep the same keys but change a field's type (e.g. `{ ...u, code: u.slug }`) or drop an optional key (e.g. `const { nickname, ...rest } = u`) are projected through `ResultTypeFromSelect` instead of being collapsed back to `U`.

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

@@ -478,8 +478,59 @@ 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 that is structurally equivalent to the canonical
+// `Ref<U>` shape the query builder produces for its underlying user type
+// `U` (taking the ref's own nullability into account). When `T` is a true
+// ref, `ExtractRef` can safely return `U` directly; otherwise it must fall
+// through to the recursive projection.
+//
+// Checking only that `T` has "no extra keys" beyond `keyof U` (plus the
+// brand/virtual props) is not sufficient. A spread-derived object can keep
+// exactly the keys of `U` while:
+//   - changing a field's type, e.g. `{ ...u, code: u.slug }`, or
+//   - dropping an optional key, e.g. `const { nickname, ...rest } = u`.
+// Both must be recursively projected, not collapsed back to `U`. We
+// therefore require strict structural equivalence against the canonical ref
+// shape rather than a one-directional key-subset check.
+type IsTrueRef<T> =
+  T extends RefLeaf<infer U>
+    ? RefShapeMatches<T, Ref<U, IsNullableRef<T>>> extends true
+      ? true
+      : false
+    : false
+
+// Strict structural equivalence between two ref shapes. Unlike plain
+// bidirectional assignability, this is sensitive to *key presence* — an
+// object that drops an optional key (e.g. `const { nickname, ...rest } = u`)
+// is not considered equal to one that keeps `nickname?`, even though the two
+// remain mutually assignable. A direct ref (`u.document`, a union member,
+// etc.) is exactly the canonical `Ref` shape and matches here, so it returns
+// `U` via the fast path; any spread-derived object differs (changed field
+// types, dropped keys, or stripped `readonly` modifiers) and instead falls
+// through to the recursive projection, which reconstructs the correct type.
+type RefShapeMatches<A, B> =
+  (<G>() => G extends A ? 1 : 2) extends <G>() => G extends B ? 1 : 2
+    ? true
+    : 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 +821,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,126 @@ 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(`spread with a same-key narrower override projects the override type`, () => {
+    type SpreadUser = {
+      id: number
+      code: string | number
+      slug: string
+      nickname?: string
+    }
+
+    const spreadUsers = createCollection(
+      mockSyncCollectionOptions<SpreadUser>({
+        id: `spread-override-users`,
+        getKey: (u) => u.id,
+        initialData: [],
+      }),
+    )
+
+    const col = createLiveQueryCollection((q) =>
+      q.from({ u: spreadUsers }).select(({ u }) => ({
+        narrowed: { ...u, code: u.slug },
+      })),
+    )
+
+    const result = col.toArray[0]!
+    // `code` was overridden with `u.slug` (string), so the projected
+    // field must be `string`, not the original `string | number`.
+    expectTypeOf(result.narrowed.code).toEqualTypeOf<string>()
+  })
+
+  test(`spread that omits an optional property drops the key`, () => {
+    type SpreadUser = {
+      id: number
+      code: string | number
+      slug: string
+      nickname?: string
+    }
+
+    const spreadUsers = createCollection(
+      mockSyncCollectionOptions<SpreadUser>({
+        id: `spread-omit-users`,
+        getKey: (u) => u.id,
+        initialData: [],
+      }),
+    )
+
+    const col = createLiveQueryCollection((q) =>
+      q.from({ u: spreadUsers }).select(({ u }) => {
+        const { nickname, ...withoutNickname } = u
+        return { trimmed: withoutNickname }
+      }),
+    )
+
+    const result = col.toArray[0]!
+    // `nickname` was destructured out, so the projected object must
+    // not reintroduce the key.
+    type HasNickname = `nickname` extends keyof typeof result.trimmed
+      ? true
+      : false
+    expectTypeOf<HasNickname>().toEqualTypeOf<false>()
+  })
+
   test(`nested spread preserves object structure types`, () => {
     const users = createUsers()
     const col = createLiveQueryCollection((q) => {