Working with discriminated unions inside of collection causes TS troubles

[pawelblaszczyk5]

Hello 👋

I'm building some small app with tanstack db and one things that bugs me is working with discriminated unions inside of collections. The app is AI chat, so for example I have MessagePart shape, which can be TextMessagePart, ReasoningMessagePart etc. Unfortunately, it's not properly narrowed in many places when using it.

I made a small reproduction, where you can see these issues, that's much easier then to describe it:

https://github.com/pawelblaszczyk5/tanstack-db-react-compiler-reproduction/blob/c2470dd6a18e869625c531cc9ccbc0bb997c84cb/src/lib/collections.ts

As an example I chosen a schema where discriminator is also TKey to showcase also the 2. issue, but it also doesn't work in cases where it's just one of the fields. In my app I have shape with 3 fields basically id, type, data and data is different based on type, I can't properly query this without typecasting/asserting atm

But for completeness:

1. Can't discriminate union in either .where or .fn.where. It always result in the full model
2. Collection TKey is only inferred properly for few collection types, doesn't work e.g. with local storage or electric ones. They don't pass necessary generics
3. get and update in cases where TKey is discriminator doesn't properly discriminate