opa-react: introduce batching (#123)

This change introduces **batching** as an option for `@styra/opa-react`.

It is enabled by passing `batch={true}` to `<AuthzProvider>` and will cause multiple calls
to `useAuthz` or usages of `<Authz>` to be combined into groups of batch requests.

It is only applicable when using Enterprise OPA (which has a [Batch API](https://docs.styra.com/enterprise-opa/reference/api-reference/batch-api)),
or when implementing a compatible API yourself.

Notes:
1. `fromResult` is exempt from caching (same as with the non-batching cache feature):
   `{path: "foo", input: {user: "alice"}, fromResult: (x) => x.foo}` and
   `{path: "foo", input: {user: "alice"}, fromResult: (x) => x.bar}` **yield the same result**
   (which one is undefined), so this should **be avoided**.
2. The batch feature doesn't support query cancellation via `signal` (`AbortController`),
   the non-batching querier does.
3. Some props fields have been removed from `AuthzContext`, since only a subset of
   the `AuthzProviderProps` are useful from the hooks etc. It's a general cleanup: you
   can **set** `retry` and `batch` on `AuthzProvider`, but those will be used to create
   the `queryClient` that's employed in `useAuthz`, so only `queryClient` (and not
   `retry` and `batch`) is retrievable from `AuthzContext`.

.changeset/cold-bikes-sit.md

@@ -0,0 +1,22 @@
+---
+"@styra/opa-react": minor
+---
+
+Support batching requests sent to the backend (optional)
+
+When used with [Enterprise OPA's Batch API](https://docs.styra.com/enterprise-opa/reference/api-reference/batch-api), this mode allows for sending much
+fewer requests to the backend. It's enabled by setting `batch={true}` on `<AuthzProvider>`.
+
+Note that the Batch API has no notion of "default query", so it's not possible
+to use batching without having either `defaultPath` (`<AuthzProvider>`) or
+`path` (`useAuthz()`, `<Authz>`) set.
+
+Please note that `fromResult` is exempt from the cache key, so multiple requests
+with the same path and input, but different `fromResult` settings will lead to
+unforeseen results.
+This is on par with the regular (non-batching) caching, and we'll revisit this
+if it becomes a problem for users. Please create an issue on Github if it is
+problematic for you.
+
+Furthermore, batching queries are not wired up with `AbortController` like the
+non-batching equivalents are.

.github/workflows/pull-request.yaml

@@ -70,7 +70,7 @@ jobs:
         run: npx -w ${{ matrix.pkg }} attw --pack --ignore-rules no-resolution
         if: matrix.pkg == 'packages/opa'
       - name: jsr publish dry-run
-        run: npx -w ${{ matrix.pkg }} jsr publish --dry-run
+        run: npx -w ${{ matrix.pkg }} jsr publish --dry-run --allow-dirty
         if: matrix.pkg == 'packages/opa'
 
   success:

packages/opa-react/package.json

@@ -39,6 +39,7 @@
   "dependencies": {
     "@styra/opa": ">=1.1.3",
     "@tanstack/react-query": "^5.50.1",
+    "@yornaath/batshit": "^0.10.1",
     "lodash.merge": "^4.6.2",
     "zod": "^3.23.8"
   },

packages/opa-react/src/authz-provider.tsx

@@ -14,6 +14,48 @@ import {
   type BatchRequestOptions,
 } from "@styra/opa";
 import { type ServerError } from "@styra/opa/sdk/models/components";
+import { create, windowScheduler } from "@yornaath/batshit";
+
+type EvalQuery = {
+  path: string | undefined;
+  input: Input | undefined;
+  fromResult: ((_?: Result) => boolean) | undefined;
+};
+
+function key({ path, input }: EvalQuery): string {
+  return stringify({ path, input }); // Note: omit fromResult
+}
+
+const evals = (sdk: OPAClient) =>
+  create({
+    fetcher: async (evals: EvalQuery[]) => {
+      const evs = evals.map((x) => ({ ...x, k: key(x) }));
+      const groups = Object.groupBy(evs, ({ path }) => path); // group by path
+      return Promise.all(
+        Object.entries(groups).map(([path, inputs]) => {
+          const inps: Record<string, Input> = {};
+          const fromRs: Record<string, (_?: Result) => boolean> = {};
+          inputs.forEach(({ k, input, fromResult }) => {
+            inps[k] = input;
+            fromRs[k] = fromResult;
+          });
+          return sdk
+            .evaluateBatch(path, inps, { rejectMixed: true })
+            .then((res) =>
+              Object.fromEntries(
+                Object.entries(res).map(([k, res]) => [
+                  k,
+                  fromRs[k] ? fromRs[k](res) : res,
+                ]),
+              ),
+            );
+        }),
+      ).then((all: object[]) => all.reduce((acc, n) => ({ ...acc, ...n }), {})); // combine result arrays of objects
+    },
+    resolver: (results, query) => results[key(query)] ?? null,
+    scheduler: windowScheduler(10),
+    name: "@styra/opa-react",
+  });
 
 /** Abstracts the methods that are used from `OPAClient` of `@styra/opa`. */
 export interface OPAClient {
@@ -53,19 +95,21 @@ export type AuthzProviderContext = {
 
   /** The `@tanstack/react-query` client that's used for scheduling policy evaluation requests. */
   queryClient: QueryClient;
-
-  /** Whether or not policy evaluations should retry on transient failures. `false` means never; `true` means infinite retry; any number N means N retries. Defaults to 3. */
-  retry: boolean | number;
 };
 
 // Reference: https://reacttraining.com/blog/react-context-with-typescript
 export const AuthzContext = createContext<AuthzProviderContext | undefined>(
   undefined,
 );
 
-export type AuthzProviderProps = PropsWithChildren<
-  Omit<AuthzProviderContext, "queryClient">
->;
+export interface AuthzProviderProps
+  extends PropsWithChildren<Omit<AuthzProviderContext, "queryClient">> {
+  /** Whether or not policy evaluations should retry on transient failures. `false` means never; `true` means infinite retry; any number N means N retries. Defaults to 3. */
+  retry?: boolean | number;
+
+  /** Batch policy evaluation queries when possible, and supported by the backend. Defaults to `false`. */
+  batch?: boolean;
+}
 
 /**
  * Configures the authorization SDK, with default path/input of applicable.
@@ -86,25 +130,43 @@ export default function AuthzProvider({
   defaultPath,
   defaultInput,
   defaultFromResult,
-  retry = 3,
+  retry = 0, // Debugging
+  batch = false,
 }: AuthzProviderProps) {
+  const batcher = useMemo(
+    () => batch && opaClient && evals(opaClient),
+    [opaClient, batch],
+  );
   const defaultQueryFn = useCallback(
-    async ({ queryKey, meta = {}, signal }: QueryFunctionContext) => {
-      if (!opaClient) return;
-
+    async ({
+      queryKey,
+      meta = {},
+      signal,
+    }: QueryFunctionContext): Promise<Result> => {
       const [path, input] = queryKey as [string, Input];
-      const fromResult = meta["fromResult"] as (_?: Result) => boolean;
-      return path
-        ? opaClient.evaluate<Input, Result>(path, input, {
-            fromResult,
-            fetchOptions: { signal },
-          })
-        : opaClient.evaluateDefault<Input, Result>(input, {
-            fromResult,
-            fetchOptions: { signal },
-          });
+      const fromResult = meta["fromResult"] as
+        | ((_?: Result) => boolean)
+        | undefined;
+
+      if (!batcher) {
+        // use the default, unbatched queries backed by react-query
+        return path
+          ? opaClient.evaluate<Input, Result>(path, input, {
+              fromResult,
+              fetchOptions: { signal },
+            })
+          : opaClient.evaluateDefault<Input, Result>(input, {
+              fromResult,
+              fetchOptions: { signal },
+            });
+      }
+
+      if (!path)
+        throw new Error("batch requests need to have a defined query path");
+
+      return batcher.fetch({ path, input, fromResult });
     },
-    [opaClient],
+    [batcher, batch],
   );
   const queryClient = useMemo(
     () =>
@@ -116,7 +178,7 @@ export default function AuthzProvider({
           },
         },
       }),
-    [defaultQueryFn],
+    [defaultQueryFn, retry],
   );
 
   const context = useMemo<AuthzProviderContext>(
@@ -126,16 +188,8 @@ export default function AuthzProvider({
       defaultInput,
       defaultFromResult,
       queryClient: queryClient!,
-      retry,
     }),
-    [
-      opaClient,
-      defaultPath,
-      defaultInput,
-      defaultFromResult,
-      queryClient,
-      retry,
-    ],
+    [opaClient, defaultPath, defaultInput, defaultFromResult, queryClient],
   );
 
   if (!queryClient) return null;
@@ -144,3 +198,53 @@ export default function AuthzProvider({
     <AuthzContext.Provider value={context}>{children}</AuthzContext.Provider>
   );
 }
+
+// Taken from fast-json-stable-hash, MIT-licensed:
+// https://github.com/zkldi/fast-json-stable-hash/blob/31b3081e942c1ce491f9698fd0bf527847093036/index.js
+// That module was tricky to import because it's using `crypto` for hashing.
+// We only need a stable string.
+function stringify(obj: any) {
+  const type = typeof obj;
+  if (obj === undefined) return "_";
+
+  if (type === "string") {
+    return JSON.stringify(obj);
+  } else if (Array.isArray(obj)) {
+    let str = "[";
+
+    let al = obj.length - 1;
+
+    for (let i = 0; i < obj.length; i++) {
+      str += stringify(obj[i]);
+
+      if (i !== al) {
+        str += ",";
+      }
+    }
+
+    return `${str}]`;
+  } else if (type === "object" && obj !== null) {
+    let str = "{";
+    let keys = Object.keys(obj).sort();
+
+    let kl = keys.length - 1;
+
+    for (let i = 0; i < keys.length; i++) {
+      let key = keys[i];
+      str += `${JSON.stringify(key)}:${stringify(obj[key])}`;
+
+      if (i !== kl) {
+        str += ",";
+      }
+    }
+
+    return `${str}}`;
+  } else if (type === "number" || type === "boolean" || obj === null) {
+    // bool, num, null have correct auto-coercions
+    return `${obj}`;
+  } else {
+    throw new TypeError(
+      `Invalid JSON type of ${type}, value ${obj}. Can only hash JSON objects.`,
+    );
+  }
+}

packages/opa-react/src/package.json

@@ -0,0 +1,3 @@
+{
+  "type": "module"
+}

packages/opa-react/src/use-authz.ts

@@ -26,7 +26,13 @@ export default function useAuthz(
   if (context === undefined) {
     throw Error("Authz/useAuthz can only be used inside an AuthzProvider");
   }
-  const { defaultPath, defaultInput, defaultFromResult, queryClient } = context;
+  const {
+    defaultPath,
+    defaultInput,
+    defaultFromResult,
+    queryClient,
+    opaClient,
+  } = context;
   const p = path ?? defaultPath;
   const i = mergeInput(input, defaultInput);
   const fromR = fromResult ?? defaultFromResult;
@@ -40,6 +46,7 @@ export default function useAuthz(
     {
       queryKey: [p, i],
       meta: { fromResult: fromR },
+      enabled: !!opaClient,
     },
     queryClient,
   );