feat: ensure data utils work well with legacy relationship proxies

Author: runspired

packages/ember/src/-private/promise-state.ts

@@ -44,12 +44,24 @@ export class PromiseState<T = unknown, E = unknown> {
   }
 }
 
+const LegacyPromiseProxy = Symbol.for('LegacyPromiseProxy');
+type LegacyAwaitable<T, E> = { promise: Promise<T> | Awaitable<T, E>; [LegacyPromiseProxy]: true };
+
+function isLegacyAwaitable<T, E>(promise: object): promise is LegacyAwaitable<T, E> {
+  return LegacyPromiseProxy in promise && 'promise' in promise && promise[LegacyPromiseProxy] === true;
+}
+
+function getPromise<T, E>(promise: Promise<T> | Awaitable<T, E> | LegacyAwaitable<T, E>): Promise<T> | Awaitable<T, E> {
+  return isLegacyAwaitable(promise) ? promise.promise : promise;
+}
+
 export function getPromiseState<T = unknown, E = unknown>(promise: Promise<T> | Awaitable<T, E>): PromiseState<T, E> {
-  let state = PromiseCache.get(promise) as PromiseState<T, E> | undefined;
+  const _promise = getPromise(promise);
+  let state = PromiseCache.get(_promise) as PromiseState<T, E> | undefined;
 
   if (!state) {
-    state = new PromiseState(promise);
-    PromiseCache.set(promise, state);
+    state = new PromiseState(_promise);
+    PromiseCache.set(_promise, state);
   }
 
   return state;

packages/model/src/-private/promise-belongs-to.ts

@@ -23,6 +23,8 @@ export interface BelongsToProxyCreateArgs<T = unknown> {
   _belongsToState: BelongsToProxyMeta<T>;
 }
 
+export const LegacyPromiseProxy = Symbol.for('LegacyPromiseProxy');
+
 interface PromiseObjectType<T> extends PromiseProxyMixin<T | null>, ObjectProxy<T> {
   // eslint-disable-next-line @typescript-eslint/no-misused-new
   new <PT>(...args: unknown[]): PromiseObjectType<PT>;
@@ -80,6 +82,8 @@ class PromiseBelongsTo<T = unknown> extends Extended<T> {
     await legacySupport.reloadBelongsTo(key, options);
     return this;
   }
+
+  [LegacyPromiseProxy] = true as const;
 }
 
 export { PromiseBelongsTo };

packages/model/src/-private/promise-many-array.ts

@@ -6,6 +6,7 @@ import { compat } from '@ember-data/tracking';
 import { defineSignal } from '@ember-data/tracking/-private';
 
 import type ManyArray from './many-array';
+import { LegacyPromiseProxy } from './promise-belongs-to';
 
 export interface HasManyProxyCreateArgs<T = unknown> {
   promise: Promise<ManyArray<T>>;
@@ -195,6 +196,8 @@ export default class PromiseManyArray<T = unknown> {
   static create<T>({ promise, content }: HasManyProxyCreateArgs<T>): PromiseManyArray<T> {
     return new this(promise, content);
   }
+
+  [LegacyPromiseProxy] = true as const;
 }
 defineSignal(PromiseManyArray.prototype, 'content', null);
 defineSignal(PromiseManyArray.prototype, 'isPending', false);

tests/warp-drive__ember/tests/integration/get-promise-state-test.gts

@@ -1,11 +1,34 @@
 import { rerender, settled } from '@ember/test-helpers';
 
-import { createDeferred, setPromiseResult } from '@ember-data/request';
+import { createDeferred, setPromiseResult, type Awaitable } from '@ember-data/request';
 import type { RenderingTestContext } from '@warp-drive/diagnostic/ember';
 import { module, setupRenderingTest, test } from '@warp-drive/diagnostic/ember';
 import { getPromiseState } from '@warp-drive/ember';
 
 type PromiseState<T, E> = ReturnType<typeof getPromiseState<T, E>>;
+const SecretSymbol = Symbol.for('LegacyPromiseProxy');
+
+class PromiseProxy<T, E> {
+  [SecretSymbol]: true;
+  promise: Awaitable<T, E>;
+
+  constructor(promise: Awaitable<T, E>) {
+    this[SecretSymbol] = true;
+    this.promise = promise;
+  }
+
+  then(onFulfilled: (value: T) => T, onRejected: (error: E) => void): Promise<T> {
+    return this.promise.then(onFulfilled, onRejected) as Promise<T>;
+  }
+
+  catch(onRejected: (error: E) => void): Promise<unknown> {
+    return this.promise.catch(onRejected) as Promise<unknown>;
+  }
+
+  finally(onFinally: () => void): Promise<unknown> {
+    return this.promise.finally(onFinally) as Promise<unknown>;
+  }
+}
 
 module('Integration | get-promise-state', function (hooks) {
   setupRenderingTest(hooks);
@@ -196,4 +219,119 @@ module('Integration | get-promise-state', function (hooks) {
     assert.equal(counter, 1);
     assert.equal(this.element.textContent?.trim(), 'Our Error\n          Count:\n          1');
   });
+
+  test('it renders only once when the promise already has a result cached', async function (this: RenderingTestContext, assert) {
+    const promise = Promise.resolve().then(() => 'Our Data');
+
+    const result = await promise;
+    setPromiseResult(promise, { result, isError: false });
+
+    let state: PromiseState<string, Error>;
+    function _getPromiseState<T>(p: Promise<T>): PromiseState<T, Error> {
+      state = getPromiseState(p) as PromiseState<string, Error>;
+      return state as PromiseState<T, Error>;
+    }
+    let counter = 0;
+    function countFor(_result: unknown) {
+      return ++counter;
+    }
+
+    await this.render(
+      <template>
+        {{#let (_getPromiseState promise) as |state|}}
+          {{state.result}}<br />Count:
+          {{countFor state.result}}
+        {{/let}}
+      </template>
+    );
+
+    assert.equal(this.element.textContent?.trim(), 'Our DataCount:\n          1');
+    await settled();
+
+    assert.equal(this.element.textContent?.trim(), 'Our DataCount:\n          1');
+  });
+
+  test('it unwraps promise-proxies that utilize the secret symbol for error states', async function (this: RenderingTestContext, assert) {
+    const _promise = Promise.resolve().then(() => {
+      throw new Error('Our Error');
+    });
+    const promise = new PromiseProxy(_promise);
+
+    try {
+      getPromiseState(promise);
+      await promise;
+    } catch {
+      // do nothing
+    }
+
+    let state: PromiseState<string, Error>;
+    function _getPromiseState<T>(p: Promise<T>): PromiseState<T, Error> {
+      state = getPromiseState(p) as PromiseState<string, Error>;
+      return state as PromiseState<T, Error>;
+    }
+    let counter = 0;
+    function countFor(_result: unknown, _error: unknown) {
+      return ++counter;
+    }
+
+    await this.render(
+      <template>
+        {{#let (_getPromiseState promise) as |state|}}
+          {{#if state.isPending}}
+            Pending
+          {{else if state.isError}}
+            {{state.error.message}}
+          {{else if state.isSuccess}}
+            Invalid Success Reached
+          {{/if}}
+          <br />Count:
+          {{countFor state.result state.error}}{{/let}}
+      </template>
+    );
+
+    assert.equal(state!.result, null);
+    assert.true(state!.error instanceof Error);
+    assert.equal((state!.error as Error | undefined)?.message, 'Our Error');
+    assert.equal(counter, 1);
+    assert.equal(this.element.textContent?.trim(), 'Our Error\n          Count:\n          1');
+    await rerender();
+    assert.equal(state!.result, null);
+    assert.true(state!.error instanceof Error);
+    assert.equal((state!.error as Error | undefined)?.message, 'Our Error');
+    assert.equal(counter, 1);
+    assert.equal(this.element.textContent?.trim(), 'Our Error\n          Count:\n          1');
+    assert.strictEqual(state, getPromiseState(_promise));
+  });
+
+  test('it unwraps promise-proxies that utilize the secret symbol for success states', async function (this: RenderingTestContext, assert) {
+    const _promise = Promise.resolve().then(() => 'Our Data');
+    const promise = new PromiseProxy(_promise);
+    getPromiseState(promise);
+    await promise;
+
+    let state: PromiseState<string, Error>;
+    function _getPromiseState<T>(p: Promise<T>): PromiseState<T, Error> {
+      state = getPromiseState(p) as PromiseState<string, Error>;
+      return state as PromiseState<T, Error>;
+    }
+    let counter = 0;
+    function countFor(_result: unknown) {
+      return ++counter;
+    }
+
+    await this.render(
+      <template>
+        {{#let (_getPromiseState promise) as |state|}}
+          {{state.result}}<br />Count:
+          {{countFor state.result}}
+        {{/let}}
+      </template>
+    );
+
+    assert.equal(this.element.textContent?.trim(), 'Our DataCount:\n          1');
+    await settled();
+
+    assert.equal(this.element.textContent?.trim(), 'Our DataCount:\n          1');
+    assert.strictEqual(state, getPromiseState(_promise));
+  });
 });