diff --git a/.changeset/fix-ref-union-collapse.md b/.changeset/fix-ref-union-collapse.md new file mode 100644 index 0000000000..c7c14995d3 --- /dev/null +++ b/.changeset/fix-ref-union-collapse.md @@ -0,0 +1,5 @@ +--- +'@tanstack/db': patch +--- + +fix: preserve discriminated union types when selecting object fields via `.select()` diff --git a/packages/db/src/query/builder/types.ts b/packages/db/src/query/builder/types.ts index 9416d10c99..a1228294c8 100644 --- a/packages/db/src/query/builder/types.ts +++ b/packages/db/src/query/builder/types.ts @@ -322,9 +322,16 @@ export type ResultTypeFromSelect = : // includes subquery (bare QueryBuilder) — produces a child Collection TSelectObject[K] extends QueryBuilder ? Collection> - : // Ref (full object ref or spread with RefBrand) - recursively process properties + : // Ref (full object ref or spread with RefBrand) TSelectObject[K] extends Ref - ? ExtractRef + ? // When the branded type is a union, extract it directly to + // preserve the discriminated union. ExtractRef would collapse + // it via keyof/Prettify which only yield common keys. + true extends IsUnion> + ? IsNullableRef extends true + ? ExtractRefBrand | undefined + : ExtractRefBrand + : ExtractRef : // RefLeaf (simple property ref like user.name) TSelectObject[K] extends RefLeaf ? IsNullableRef extends true @@ -372,6 +379,9 @@ export type SelectResult = ? ResultTypeFromSelect : ResultTypeFromSelectValue +// Extract the raw type stored in a RefLeaf's brand +type ExtractRefBrand = T extends RefLeaf ? U : never + // Extract Ref or subobject with a spread or a Ref type ExtractRef = Prettify>> @@ -1068,6 +1078,14 @@ type IsPlainObject = T extends unknown type IsAny = 0 extends 1 & T ? true : false +// Detects whether T is a union type (e.g., A | B | C returns true, A returns false) +type IsUnion = T extends unknown + ? [U] extends [T] + ? false + : true + : false + + /** * JsBuiltIns - List of JavaScript built-ins */ diff --git a/packages/db/tests/query/select.test-d.ts b/packages/db/tests/query/select.test-d.ts index 225decdfb3..3941dccaf8 100644 --- a/packages/db/tests/query/select.test-d.ts +++ b/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,67 @@ describe(`select types`, () => { expectTypeOf(results).toMatchTypeOf>() }) + 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({ + 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() + }) + + test(`select preserves union when collection type is a union`, () => { + type DocV1 = { version: 1; title: string } + type DocV2 = { version: 2; title: string; subtitle: string } + type Doc = DocV1 | DocV2 + + const docs = createCollection( + mockSyncCollectionOptions({ + id: `union-collection`, + getKey: (d) => d.title, + initialData: [], + }), + ) + + // Without select — union preserved + const col1 = createLiveQueryCollection((q) => q.from({ d: docs })) + const r1 = col1.toArray[0]! + expectTypeOf(r1).toMatchTypeOf() + + // With select on individual fields — per-field unions (not top-level union) + const col2 = createLiveQueryCollection((q) => + q.from({ d: docs }).select(({ d }) => ({ + version: d.version, + title: d.title, + })), + ) + const r2 = col2.toArray[0]! + expectTypeOf(r2).toMatchTypeOf<{ version: 1 | 2; title: string }>() + }) + test(`nested spread preserves object structure types`, () => { const users = createUsers() const col = createLiveQueryCollection((q) => {