feat(query): entities api rework

examples/event-board/common/src/stores-di/events.ts

@@ -80,11 +80,11 @@ export function EventsList$() {
       })
     },
     [
-      entities(data => ({
+      entities((capture, data) => ({
         ...data,
         pages: data.pages.map(page => ({
           ...page,
-          events: page.events.map(EventEntity)
+          events: page.events.map(capture(EventEntity))
         }))
       }))
     ]

examples/event-board/preact-nano_kit-intl-ssr/src/stores/events.ts

@@ -80,11 +80,11 @@ export function EventsList$() {
       })
     },
     [
-      entities(data => ({
+      entities((capture, data) => ({
         ...data,
         pages: data.pages.map(page => ({
           ...page,
-          events: page.events.map(EventEntity)
+          events: page.events.map(capture(EventEntity))
         }))
       }))
     ]

examples/event-board/react-nano_kit-intl-ssr/src/stores/events.ts

@@ -80,11 +80,11 @@ export function EventsList$() {
       })
     },
     [
-      entities(data => ({
+      entities((capture, data) => ({
         ...data,
         pages: data.pages.map(page => ({
           ...page,
-          events: page.events.map(EventEntity)
+          events: page.events.map(capture(EventEntity))
         }))
       }))
     ]

examples/event-board/svelte-nano_kit-ssr/src/stores/events.ts

@@ -80,11 +80,11 @@ export function EventsList$() {
       })
     },
     [
-      entities(data => ({
+      entities((capture, data) => ({
         ...data,
         pages: data.pages.map(page => ({
           ...page,
-          events: page.events.map(EventEntity)
+          events: page.events.map(capture(EventEntity))
         }))
       }))
     ]

packages/query/.size-limit.json

@@ -3,7 +3,7 @@
     "name": "All publics",
     "path": "dist/index.js",
     "import": "*",
-    "limit": "3.94 kB"
+    "limit": "3.96 kB"
   },
   {
     "name": "Minimal set",

packages/query/src/settings/entities.spec.ts

@@ -58,9 +58,9 @@ describe('query', () => {
           entities(PostEntity)
         ])
         const [$posts] = query(PostsKey, [], getPosts, [
-          entities(postsPage => ({
+          entities((capture, postsPage) => ({
             ...postsPage,
-            posts: postsPage.posts.map(PostEntity)
+            posts: postsPage.posts.map(capture(PostEntity))
           }))
         ])
         const offPost = effect(() => {
@@ -181,9 +181,9 @@ describe('query', () => {
         )
         const firstPostKey = PostEntity(1)
         const [$posts] = query(PostsKey, [], getPosts, [
-          entities(postsPage => ({
+          entities((capture, postsPage) => ({
             ...postsPage,
-            posts: postsPage.posts.map(PostEntity)
+            posts: postsPage.posts.map(capture(PostEntity))
           }))
         ])
         const [mutate] = mutation<[id: number, title: string, content: string], Post | null>(

packages/query/src/settings/entities.ts

@@ -1,24 +1,37 @@
 import type { PickNonEmptyValue } from '@nano_kit/store'
 import type { ClientSetting } from '../client.types.js'
-import type { CacheKey } from '../CacheStorage.types.js'
+import type {
+  CacheKey,
+  CacheShardKey
+} from '../CacheStorage.types.js'
 import type {
   ClientContext,
   MutationClientContext,
   QueryClientContext
 } from '../ClientContext.js'
 import { queryKey } from '../cache.js'
 
-export interface Entity<T extends {}> {
-  /**
-   * Get the cache key for the entity by its identifier.
-   */
-  (id: number | string): CacheKey<[id: number | string], T | null>
-  /**
-   * Get or upsert the entity in the cache.
-   */
-  (entity: T): T
+export type Entity<T extends {}> = (
+  (id: number | string | T) => CacheKey<[id: number | string], T | null>
+) & CacheShardKey
+
+export interface EntityCapture {
+  <T extends {}>(
+    EntityFn: Entity<T>
+  ): (entity: T) => T
+  <T extends {}>(
+    EntityFn: Entity<T>,
+    entity: T
+  ): T
 }
 
+export type EntityMapper<T> = (
+  (
+    capture: EntityCapture,
+    data: T
+  ) => T
+) | Entity<T extends {} ? T : {}>
+
 const ENTITY_KEY = '#entity'
 const EntityKey = queryKey(ENTITY_KEY)
 
@@ -36,8 +49,6 @@ function isEntityRef<T extends {}>(value: T): value is T & EntityRef {
   return ENTITY_KEY in value
 }
 
-let currentCtx: ClientContext | null = null
-
 /**
  * Create an entity manager for a specific entity type.
  * @param name - The name of the entity type.
@@ -63,31 +74,19 @@ export function entity<T extends { id: number | string }>(
   name: string,
   id = (entity: T) => entity.id
 ) {
-  return (idOrRefOrEntity: number | string | null | undefined | T) => {
+  const entityKey = (idOrRefOrEntity: number | string | null | undefined | T) => {
     if (isIdentifier(idOrRefOrEntity)) {
       return EntityKey(name, idOrRefOrEntity)
     }
 
-    if (!idOrRefOrEntity || !currentCtx) {
+    if (!idOrRefOrEntity) {
       return idOrRefOrEntity
     }
 
-    if (isEntityRef(idOrRefOrEntity)) {
-      return currentCtx.$get(idOrRefOrEntity[ENTITY_KEY]).data
-    }
-
-    const key = EntityKey(name, id(idOrRefOrEntity))
-
-    currentCtx.set(key, {
-      ...currentCtx.initial(),
-      data: idOrRefOrEntity
-    })
-
-    return {
-      ...idOrRefOrEntity,
-      [ENTITY_KEY]: key
-    }
+    return EntityKey(name, id(idOrRefOrEntity))
   }
+
+  return Object.assign(entityKey, EntityKey)
 }
 
 /**
@@ -96,7 +95,7 @@ export function entity<T extends { id: number | string }>(
  * @returns The client setting function.
  */
 export function entities<T>(
-  mapper: (data: NoInfer<PickNonEmptyValue<T>>) => NoInfer<PickNonEmptyValue<T>>
+  mapper: NoInfer<EntityMapper<PickNonEmptyValue<T>>>
 ): ClientSetting<QueryClientContext<T>>
 
 /**
@@ -105,25 +104,44 @@ export function entities<T>(
  * @returns The client setting function.
  */
 export function entities<T>(
-  mapper: (data: NoInfer<PickNonEmptyValue<T>>) => NoInfer<PickNonEmptyValue<T>>
+  mapper: NoInfer<EntityMapper<PickNonEmptyValue<T>>>
 ): ClientSetting<MutationClientContext<T>>
 
 /* @__NO_SIDE_EFFECTS__ */
-export function entities(mapper: (data: unknown) => unknown) {
+export function entities(
+  mapper: EntityMapper<unknown>
+) {
   return (ctx: ClientContext) => {
-    const safeMapper = (data: unknown) => {
-      if (data) {
-        try {
-          currentCtx = ctx
-
-          return mapper(data)
-        } finally {
-          currentCtx = null
-        }
+    const capture = (
+      EntityFn: Entity<{}>,
+      entity: unknown
+    ) => {
+      if (!entity) {
+        return (entity: unknown) => capture(EntityFn, entity)
+      }
+
+      if (isEntityRef(entity)) {
+        return ctx.$get(entity[ENTITY_KEY]).data
       }
 
-      return data
+      const key = EntityFn(entity)
+
+      ctx.set(key, {
+        ...ctx.initial(),
+        // params: key.params,
+        data: entity
+      })
+
+      return {
+        ...entity,
+        [ENTITY_KEY]: key
+      }
     }
+    const safeMapper = (data: {}) => data && (
+      'shard' in mapper
+        ? capture(mapper, data)
+        : mapper(capture as EntityCapture, data)
+    )
 
     ctx.mapComputedData = ctx.mapData = safeMapper
   }

website/src/content/docs/query/advanced.mdx

@@ -525,9 +525,9 @@ const [$posts] = query(PostsKey, [], () => (
 ), [
   /* Map entities in the page to entity references */
   /* Also every refetch will update entities in the cache */
-  entities(page => ({
+  entities((capture, page) => ({
     ...page,
-    posts: page.posts.map(PostEntity)
+    posts: page.posts.map(capture(PostEntity))
   }))
 ])
 
@@ -564,6 +564,8 @@ With this setup:
 2. Fetching specific post via `$post` updates the entity in `$posts` as well.
 3. `updatePost` optimistically updates the entity, instantly reflecting changes in both `$post` and `$posts`.
 
+The `capture` helper stores each entity in the shared entity cache and returns a reference for the query result. Passing `entities(PostEntity)` is the shorthand for a single entity result.
+
 **When to use:** Complex applications where the same data (e.g., a "User" or "Product") appears in multiple places or lists and needs to stay synchronized.
 
 ### `tasks`

website/src/content/docs/tutorial/optimistic-updates.mdx

@@ -45,17 +45,19 @@ infinite(
   lastPage => lastPage.nextCursor,
   (q, category, cursor) => fetchEvents({ q, category, cursor }),
   [
-    entities(data => ({
+    entities((capture, data) => ({
       ...data,
       pages: data.pages.map(page => ({
         ...page,
-        events: page.events.map(EventEntity)
+        events: page.events.map(capture(EventEntity))
       }))
     }))
   ]
 )
 ```
 
+`capture(EventEntity)` stores every event in the shared entity cache and leaves references inside the list pages.
+
 The mutation also needs `$data` from the query client:
 
 ```ts