feat(core): New plugin utils withState (#2607)

* feat(core): New plugin utils `withState`

* changeset

* fix changeset

Author: EmrysMyrddin

.changeset/sharp-goats-trade.md

@@ -0,0 +1,63 @@
+---
+'@envelop/core': minor
+---
+
+Added new `withState` plugin utility for easy data sharing between hooks.
+
+## New plugin utility to ease data sharing between hooks
+
+Sometimes, plugins can grow in complexity and need to share data between its hooks.
+
+A way to solve this can be to mutate the graphql context, but this context is not always available
+in all hooks in Yoga or Hive Gateway plugins. Moreover, mutating the context gives access to your
+internal data to all other plugins and graphql resolvers, without mentioning performance impact on
+field access on this object.
+
+The recommended approach to this problem was to use a `WeakMap` with a stable key (often the
+`context` or `request` object). While it works, it's not very convenient for plugin developers, and
+is prone to error with the choice of key.
+
+The new `withState` utility solves this DX issue by providing an easy and straightforward API for
+data sharing between hooks.
+
+```ts
+import { withState } from '@envelop/core'
+
+type State = { foo: string }
+
+const myPlugin = () =>
+  withState<Plugin, State>(() => ({
+    onParse({ state }) {
+      state.forOperation.foo = 'foo'
+    },
+    onValidate({ state }) {
+      const { foo } = state.forOperation
+      console.log('foo', foo)
+    }
+  }))
+```
+
+The `state` payload field will be available in all relevant hooks, making it easy to access shared
+data. It also forces the developer to choose the scope for the data:
+
+- `forOperation` for a data scoped to GraphQL operation (Envelop, Yoga and Hive Gateway)
+- `forRequest` for a data scoped to HTTP request (Yoga and Hive Gateway)
+- `forSubgraphExecution` for a data scoped to the subgraph execution (Hive Gateway)
+
+Not all scopes are available in all hooks, the type reflects which scopes are available
+
+Under the hood, those states are kept in memory using `WeakMap`, which avoid any memory leaks.
+
+It is also possible to manually retrieve the state with the `getState` function:
+
+```ts
+const myPlugin = () =>
+  withState(getState => ({
+    onParse({ context }) {
+      // You can provide a payload, which will dictate which scope you have access to.
+      // The scope can contain `context`, `request` and `executionRequest` fields.
+      const state = getState({ context })
+      // Use the state elsewhere.
+    }
+  }))
+```

packages/core/src/index.ts

@@ -12,4 +12,5 @@ export * from './plugins/use-payload-formatter.js';
 export * from './plugins/use-masked-errors.js';
 export * from './plugins/use-engine.js';
 export * from './plugins/use-validation-rule.js';
+export * from './plugin-with-state.js';
 export { getDocumentString } from './document-string-map.js';

packages/core/src/plugin-with-state.ts

@@ -0,0 +1,164 @@
+import { MaybePromise } from '@whatwg-node/promise-helpers';
+
+export function withState<
+  P extends { instrumentation?: GenericInstrumentation },
+  HttpState = object,
+  GraphqlState = object,
+  SubExecState = object,
+>(
+  pluginFactory: (
+    getState: <SP extends {}>(
+      payload: SP,
+    ) => PayloadWithState<SP, HttpState, GraphqlState, SubExecState>['state'],
+  ) => PluginWithState<P, HttpState, GraphqlState, SubExecState>,
+): P {
+  const states: {
+    forRequest?: WeakMap<Request, Partial<HttpState>>;
+    forOperation?: WeakMap<object, Partial<GraphqlState>>;
+    forSubgraphExecution?: WeakMap<{ context: any }, Partial<SubExecState>>;
+  } = {};
+
+  function getProp(scope: keyof typeof states, key: any): PropertyDescriptor {
+    return {
+      get() {
+        if (!states[scope]) states[scope] = new WeakMap<any, any>();
+        let value = states[scope].get(key as any);
+        if (!value) states[scope].set(key, (value = {}));
+        return value;
+      },
+      enumerable: true,
+    };
+  }
+
+  function getState(payload: Payload) {
+    if (!payload) {
+      return undefined;
+    }
+    let { executionRequest, context, request } = payload;
+    const state = {};
+    const defineState = (scope: keyof typeof states, key: any) =>
+      Object.defineProperty(state, scope, getProp(scope, key));
+
+    if (executionRequest) {
+      defineState('forSubgraphExecution', executionRequest);
+      // ExecutionRequest can happen outside of any Graphql Operation for Gateway internal usage like Introspection queries.
+      // We check for `params` to be present, which means it's actually a GraphQL context.
+      if (executionRequest.context?.params) context = executionRequest.context;
+    }
+    if (context) {
+      defineState('forOperation', context);
+      if (context.request) request = context.request;
+    }
+    if (request) {
+      defineState('forRequest', request);
+    }
+    return state;
+  }
+
+  function addStateGetters(src: any) {
+    const result: any = {};
+    for (const [hookName, hook] of Object.entries(src) as any) {
+      if (typeof hook !== 'function') {
+        result[hookName] = hook;
+      } else {
+        result[hookName] = {
+          [hook.name](payload: any, ...args: any[]) {
+            if (payload && (payload.request || payload.context || payload.executionRequest)) {
+              return hook(
+                {
+                  ...payload,
+                  get state() {
+                    return getState(payload);
+                  },
+                },
+                ...args,
+              );
+            } else {
+              return hook(payload, ...args);
+            }
+          },
+        }[hook.name];
+      }
+    }
+    return result;
+  }
+
+  const { instrumentation, ...hooks } = pluginFactory(getState as any);
+
+  const pluginWithState = addStateGetters(hooks);
+  if (instrumentation) {
+    pluginWithState.instrumentation = addStateGetters(instrumentation);
+  }
+
+  return pluginWithState as P;
+}
+
+export type HttpState<T> = {
+  forRequest: Partial<T>;
+};
+
+export type GraphQLState<T> = {
+  forOperation: Partial<T>;
+};
+
+export type GatewayState<T> = {
+  forSubgraphExecution: Partial<T>;
+};
+
+export function getMostSpecificState<T>(
+  state: Partial<HttpState<T> & GraphQLState<T> & GatewayState<T>> = {},
+): Partial<T> | undefined {
+  const { forOperation, forRequest, forSubgraphExecution } = state;
+  return forSubgraphExecution ?? forOperation ?? forRequest;
+}
+
+type Payload = {
+  request?: Request;
+  context?: any;
+  executionRequest?: { context: any };
+};
+
+type GenericInstrumentation = Record<
+  string,
+  (payload: any, wrapped: () => MaybePromise<void>) => MaybePromise<void>
+>;
+
+// Brace yourself! TS Wizardry is coming!
+
+type PayloadWithState<T, Http, GraphQL, Gateway> = T extends {
+  executionRequest: any;
+}
+  ? T & {
+      state: Partial<HttpState<Http> & GraphQLState<GraphQL>> & GatewayState<Gateway>;
+    }
+  : T extends {
+        executionRequest?: any;
+      }
+    ? T & {
+        state: Partial<HttpState<Http> & GraphQLState<GraphQL> & GatewayState<Gateway>>;
+      }
+    : T extends { context: any }
+      ? T & { state: HttpState<Http> & GraphQLState<GraphQL> }
+      : T extends { request: any }
+        ? T & { state: HttpState<Http> }
+        : T extends { request?: any }
+          ? T & { state: Partial<HttpState<Http>> }
+          : T;
+
+export type PluginWithState<P, Http = object, GraphQL = object, Gateway = object> = {
+  [K in keyof P]: K extends 'instrumentation'
+    ? P[K] extends infer Instrumentation | undefined
+      ? {
+          [I in keyof Instrumentation]: Instrumentation[I] extends
+            | ((payload: infer IP, ...args: infer Args) => infer IR)
+            | undefined
+            ?
+                | ((payload: PayloadWithState<IP, Http, GraphQL, Gateway>, ...args: Args) => IR)
+                | undefined
+            : Instrumentation[I];
+        }
+      : P[K]
+    : P[K] extends ((payload: infer T) => infer R) | undefined
+      ? ((payload: PayloadWithState<T, Http, GraphQL, Gateway>) => R) | undefined
+      : P[K];
+};