feat: Invalidate accepts Unions (#3685)

* feat: Invalidate accepts Unions

* fix: Invalidate should denormalize to undefined if entity is not found

Author: ntucker

.changeset/invalidate-union-support.md

@@ -0,0 +1,21 @@
+---
+'@data-client/endpoint': minor
+'@data-client/rest': minor
+'@data-client/graphql': minor
+---
+
+Add [Union](https://dataclient.io/rest/api/Union) support to [schema.Invalidate](https://dataclient.io/rest/api/Invalidate)
+for polymorphic delete operations:
+
+```ts
+new schema.Invalidate(
+  { users: User, groups: Group },
+  'type'
+)
+```
+
+or
+
+```ts
+new schema.Invalidate(MyUnionSchema)
+```

.cursor/commands/changeset.md

@@ -26,6 +26,7 @@ Generate release notes for user-facing changes in published packages.
 - **Breaking**: Prefix with `BREAKING CHANGE:` or `BREAKING:`
 - **Body**: 1–3 lines describing outcome, not implementation
 - **New exports**: Use "New exports:" with bullet list
+- **Documentation links**: Link concepts that have doc pages in @docs (e.g., `[Union](https://dataclient.io/rest/api/Union)`)
 
 ## Code Examples in Changesets
 - Fixes: `// Before: ... ❌` `// After: ... ✓`

docs/rest/api/Invalidate.md

@@ -12,9 +12,22 @@ import EndpointPlayground from '@site/src/components/HTTP/EndpointPlayground';
 Describes entities to be marked as [INVALID](/docs/concepts/expiry-policy#invalid). This removes items from a
 collection, or [forces suspense](/docs/concepts/expiry-policy#invalidate-entity) for endpoints where the entity is required. 
 
-Constructor:
+## Constructor
 
-- `entity` which entity to invalidate. The input is used to compute the pk() for lookup.
+```typescript
+new schema.Invalidate(entity)
+new schema.Invalidate(union)
+new schema.Invalidate(entityMap, schemaAttribute)
+```
+
+- `entity`: A singular [Entity](./Entity.md) to invalidate.
+- `union`: A [Union](./Union.md) schema for polymorphic invalidation.
+- `entityMap`: A mapping of schema keys to [Entities](./Entity.md).
+- `schemaAttribute`: _optional_ (required if `entityMap` is used) The attribute on each entity found that defines what schema, per the entityMap, to use when normalizing.
+  Can be a string or a function. If given a function, accepts the following arguments:
+  - `value`: The input value of the entity.
+  - `parent`: The parent object of the input array.
+  - `key`: The key at which the input array appears on the parent object.
 
 ## Usage
 
@@ -177,6 +190,68 @@ PostResource.deleteMany(['5', '13', '7']);
 
 </EndpointPlayground>
 
+### Polymorphic types
+
+If your endpoint can delete more than one type of entity, you can use polymorphic invalidation.
+
+#### With Union schema
+
+The simplest approach is to pass an existing [Union](./Union.md) schema directly:
+
+```typescript
+class User extends Entity {
+  id = '';
+  name = '';
+  readonly type = 'users';
+}
+class Group extends Entity {
+  id = '';
+  groupname = '';
+  readonly type = 'groups';
+}
+
+const MemberUnion = new schema.Union(
+  { users: User, groups: Group },
+  'type'
+);
+
+const deleteMember = new RestEndpoint({
+  path: '/members/:id',
+  method: 'DELETE',
+  schema: new schema.Invalidate(MemberUnion),
+});
+```
+
+#### string schemaAttribute
+
+Alternatively, define the polymorphic mapping inline with a string attribute:
+
+```typescript
+const deleteMember = new RestEndpoint({
+  path: '/members/:id',
+  method: 'DELETE',
+  schema: new schema.Invalidate(
+    { users: User, groups: Group },
+    'type'
+  ),
+});
+```
+
+#### function schemaAttribute
+
+The return values should match a key in the entity map. This is useful for more complex discrimination logic:
+
+```typescript
+const deleteMember = new RestEndpoint({
+  path: '/members/:id',
+  method: 'DELETE',
+  schema: new schema.Invalidate(
+    { users: User, groups: Group },
+    (input, parent, key) => input.memberType === 'user' ? 'users' : 'groups'
+  ),
+});
+```
+
 ### Impact on useSuspense()
 
 When entities are invalidated in a result currently being presented in React, useSuspense()

packages/endpoint/src/schemas/Invalidate.ts

@@ -1,10 +1,9 @@
-import type {
-  EntityInterface,
-  INormalizeDelegate,
-  SchemaSimple,
-} from '../interface.js';
+import PolymorphicSchema from './Polymorphic.js';
+import type { EntityInterface, INormalizeDelegate } from '../interface.js';
 import type { AbstractInstanceType } from '../normal.js';
 
+type ProcessableEntity = EntityInterface & { process: any };
+
 /**
  * Marks entity as Invalid.
  *
@@ -13,28 +12,29 @@ import type { AbstractInstanceType } from '../normal.js';
  * @see https://dataclient.io/rest/api/Invalidate
  */
 export default class Invalidate<
-  E extends EntityInterface & {
-    process: any;
-  },
-> implements SchemaSimple {
-  declare protected _entity: E;
-
+  E extends ProcessableEntity | Record<string, ProcessableEntity>,
+> extends PolymorphicSchema {
   /**
    * Marks entity as Invalid.
    *
    * This triggers suspense for all endpoints requiring it.
    * Optional (like variable sized Array and Values) will simply remove the item.
    * @see https://dataclient.io/rest/api/Invalidate
    */
-  constructor(entity: E) {
+  constructor(
+    entity: E,
+    schemaAttribute?: E extends Record<string, ProcessableEntity> ?
+      string | ((input: any, parent: any, key: any) => string)
+    : undefined,
+  ) {
     if (process.env.NODE_ENV !== 'production' && !entity) {
       throw new Error('Invalidate schema requires "entity" option.');
     }
-    this._entity = entity;
+    super(entity, schemaAttribute);
   }
 
   get key(): string {
-    return this._entity.key;
+    return this.schemaKey();
   }
 
   normalize(
@@ -44,53 +44,72 @@ export default class Invalidate<
     args: any[],
     visit: (...args: any) => any,
     delegate: INormalizeDelegate,
-  ): string {
-    // TODO: what's store needs to be a differing type from fromJS
-    const processedEntity = this._entity.process(input, parent, key, args);
-    let pk = this._entity.pk(processedEntity, parent, key, args);
+  ): string | { id: string; schema: string } {
+    const entitySchema = this.inferSchema(input, parent, key);
+    if (!entitySchema) return input;
+
+    // Handle string/number input (already processed pk)
+    // Note: This branch is typically not reached through public API as getVisit
+    // handles primitives before calling schema.normalize()
+    let pk: string | number | undefined;
+    /* istanbul ignore if */
+    if (typeof input === 'string' || typeof input === 'number') {
+      pk = input;
+    } else {
+      // Must call process() to get correct pk
+      const processedEntity = entitySchema.process(input, parent, key, args);
+      pk = entitySchema.pk(processedEntity ?? input, parent, key, args);
 
-    if (
-      process.env.NODE_ENV !== 'production' &&
-      (pk === undefined || pk === '' || pk === 'undefined')
-    ) {
-      const error = new Error(
-        `Missing usable primary key when normalizing response.
+      if (
+        process.env.NODE_ENV !== 'production' &&
+        (pk === undefined || pk === '' || pk === 'undefined')
+      ) {
+        const error = new Error(
+          `Missing usable primary key when normalizing response.
 
   This is likely due to a malformed response.
   Try inspecting the network response or fetch() return value.
   Or use debugging tools: https://dataclient.io/docs/getting-started/debugging
   Learn more about schemas: https://dataclient.io/docs/api/schema
 
-  Invalidate(Entity): Invalidate(${this._entity.key})
+  Invalidate(Entity): Invalidate(${entitySchema.key})
   Value (processed): ${input && JSON.stringify(input, null, 2)}
   `,
-      );
-      (error as any).status = 400;
-      throw error;
+        );
+        (error as any).status = 400;
+        throw error;
+      }
     }
     pk = `${pk}`; // ensure pk is a string
 
     // any queued updates are meaningless with delete, so we should just set it
     // and creates will have a different pk
-    delegate.invalidate({ key: this._entity.key }, pk);
-    return pk;
+    delegate.invalidate(entitySchema, pk);
+
+    return this.isSingleSchema ? pk : (
+        { id: pk, schema: this.getSchemaAttribute(input, parent, key) }
+      );
   }
 
-  queryKey(args: any, unvisit: unknown, delegate: unknown): undefined {
+  queryKey(_args: any, _unvisit: unknown, _delegate: unknown): undefined {
     return undefined;
   }
 
   denormalize(
-    id: string,
+    id: string | { id: string; schema: string },
     args: readonly any[],
     unvisit: (schema: any, input: any) => any,
-  ): AbstractInstanceType<E> {
-    // TODO: is this really always going to be the full object - validate that calling fetch will give this even when input is a string
-    return unvisit(this._entity, id) as any;
+  ): E extends ProcessableEntity ? AbstractInstanceType<E>
+  : AbstractInstanceType<E[keyof E]> {
+    // denormalizeValue handles both single entity and polymorphic cases
+    return this.denormalizeValue(id, unvisit) as any;
   }
 
   /* istanbul ignore next */
-  _denormalizeNullable(): AbstractInstanceType<E> | undefined {
+  _denormalizeNullable():
+    | (E extends ProcessableEntity ? AbstractInstanceType<E>
+      : AbstractInstanceType<E[keyof E]>)
+    | undefined {
     return {} as any;
   }
 

packages/endpoint/src/schemas/Polymorphic.ts

@@ -23,8 +23,9 @@ export default class PolymorphicSchema {
   }
 
   define(definition: any) {
-    // sending Union into another Polymorphic gets hoisted
-    if ('_schemaAttribute' in definition && !this._schemaAttribute) {
+    // Only Union opts into hoisting (_hoistable = true)
+    // This prevents Array(Array(...)), Values(Array(...)), Array(Invalidate(...)) issues
+    if (definition._hoistable && !this._schemaAttribute) {
       this.schema = definition.schema;
       this._schemaAttribute = definition._schemaAttribute;
     } else {

packages/endpoint/src/schemas/Union.ts

@@ -6,6 +6,9 @@ import { Visit } from '../interface.js';
  * @see https://dataclient.io/rest/api/Union
  */
 export default class UnionSchema extends PolymorphicSchema {
+  // Union is designed to be transparent; allow hoisting into wrappers (Array, Values)
+  protected readonly _hoistable = true as const;
+
   constructor(definition: any, schemaAttribute: any) {
     if (!schemaAttribute) {
       throw new Error(

packages/endpoint/src/schemas/__tests__/Array.test.js

@@ -452,3 +452,58 @@ describe.each([
     });
   });
 });
+
+describe('nested polymorphic schemas', () => {
+  class User extends IDEntity {
+    type = 'users';
+  }
+  class Group extends IDEntity {
+    type = 'groups';
+  }
+
+  test('Array of Array normalizes without hoisting', () => {
+    const innerArray = new schema.Array(User);
+    const outerArray = new schema.Array(innerArray);
+
+    const input = [[{ id: '1' }, { id: '2' }], [{ id: '3' }]];
+    const output = normalize(outerArray, input);
+
+    expect(output.entities.User).toEqual({
+      1: expect.objectContaining({ id: '1' }),
+      2: expect.objectContaining({ id: '2' }),
+      3: expect.objectContaining({ id: '3' }),
+    });
+    expect(output.result).toEqual([['1', '2'], ['3']]);
+  });
+
+  test('Array of Union normalizes with hoisting', () => {
+    const union = new schema.Union({ users: User, groups: Group }, 'type');
+    const arrayOfUnion = new schema.Array(union);
+
+    const input = [
+      { id: '1', type: 'users' },
+      { id: '2', type: 'groups' },
+    ];
+    const output = normalize(arrayOfUnion, input);
+
+    expect(output.entities.User['1']).toBeDefined();
+    expect(output.entities.Group['2']).toBeDefined();
+    expect(output.result).toEqual([
+      { id: '1', schema: 'users' },
+      { id: '2', schema: 'groups' },
+    ]);
+  });
+
+  test('Array of Invalidate normalizes without hoisting (calls invalidate)', () => {
+    const invalidate = new schema.Invalidate(User);
+    const arrayOfInvalidate = new schema.Array(invalidate);
+
+    const input = [{ id: '1' }, { id: '2' }];
+    const output = normalize(arrayOfInvalidate, input);
+
+    // Invalidate should mark entities as INVALID, not store them as objects
+    expect(output.entities.User['1']).toEqual(expect.any(Symbol));
+    expect(output.entities.User['2']).toEqual(expect.any(Symbol));
+    expect(output.result).toEqual(['1', '2']);
+  });
+});