feat(base-data-service): Add struct option for validation

packages/base-data-service/package.json

@@ -58,6 +58,7 @@
   "dependencies": {
     "@metamask/messenger": "^2.0.0",
     "@metamask/storage-service": "^1.0.2",
+    "@metamask/superstruct": "^3.1.0",
     "@metamask/utils": "^11.11.0",
     "@tanstack/query-core": "^4.43.0",
     "cockatiel": "^3.1.2",

packages/base-data-service/src/BaseDataService.test.ts

@@ -222,6 +222,35 @@ describe('BaseDataService', () => {
     expect(publishSpy).toHaveBeenCalledTimes(8);
   });
 
+  describe('validation', () => {
+    beforeAll(() => {
+      jest.useRealTimers();
+    });
+
+    afterAll(() => {
+      jest.useFakeTimers({ doNotFake: ['nextTick', 'setImmediate'] });
+    });
+
+    beforeEach(() => {
+      cleanAll();
+    });
+
+    it('throws when fetchQuery response fails struct validation', async () => {
+      const messenger = new Messenger({ namespace: serviceName });
+      const service = new ExampleDataService(messenger);
+
+      mockAssets({ status: 200, body: { foo: 'bar' } });
+      mockAssets({ status: 200, body: { foo: 'bar' } });
+      mockAssets({ status: 200, body: { foo: 'bar' } });
+
+      await expect(service.getAssets(MOCK_ASSETS)).rejects.toThrow(
+        /ExampleDataService:getAssets.*returned an invalid response:/u,
+      );
+
+      service.destroy();
+    });
+  });
+
   describe('service policy', () => {
     beforeAll(() => {
       jest.useRealTimers();

packages/base-data-service/src/BaseDataService.ts

@@ -8,6 +8,7 @@ import type {
   StorageServiceRemoveItemAction,
   StorageServiceSetItemAction,
 } from '@metamask/storage-service';
+import { Struct } from '@metamask/superstruct';
 import { Duration, inMilliseconds } from '@metamask/utils';
 import type { Json } from '@metamask/utils';
 import {
@@ -33,6 +34,7 @@ import {
   CreateServicePolicyOptions,
   ServicePolicy,
 } from './createServicePolicy';
+import { processQueryResponse } from './utils';
 
 // Data service queries use the following format: ['ServiceActionName', ...params]
 export type QueryKey = [string, ...Json[]];
@@ -81,6 +83,10 @@ type DataServiceEvents<ServiceName extends string> =
   | DataServiceCacheUpdatedEvent<ServiceName>
   | DataServiceGranularCacheUpdatedEvent<ServiceName>;
 
+type AdditionalQueryOptions<QueryFnData> = {
+  struct?: Struct<QueryFnData>;
+};
+
 // Defaults to apply to all data service queries if no default option specified
 const QUERY_CLIENT_DEFAULTS: DefaultOptions = {
   queries: {
@@ -240,26 +246,33 @@ export class BaseDataService<
    *
    * @param options - The options defining the query. Keep in mind that `queryKey` and `queryFn` are required when using data services.
    * Additionally `retry` and `retryDelay` are not available, retries can be customized using the `servicePolicyOptions`.
+   * @param options.struct - An optional struct for validating the response of the query function.
    * @returns The query results.
    */
   protected async fetchQuery<
     TQueryFnData extends Json,
     TError = unknown,
     TData = TQueryFnData,
     TQueryKey extends QueryKey = QueryKey,
-  >(
-    options: WithRequired<
-      OmitKeyof<
-        FetchQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
-        'retry' | 'retryDelay'
-      >,
-      'queryKey' | 'queryFn'
+  >({
+    struct,
+    ...options
+  }: WithRequired<
+    OmitKeyof<
+      FetchQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
+      'retry' | 'retryDelay'
     >,
-  ): Promise<TData> {
+    'queryKey' | 'queryFn'
+  > &
+    AdditionalQueryOptions<TQueryFnData>): Promise<TData> {
     return this.#queryClient.fetchQuery({
       ...options,
       queryFn: (context) =>
-        this.#policy.execute(() => options.queryFn(context)),
+        // TODO: Consider if the validation should happen outside the policy executor
+        this.#policy.execute(async () => {
+          const response = await options.queryFn(context);
+          return processQueryResponse(options.queryKey, response, struct);
+        }),
     });
   }
 
@@ -268,6 +281,7 @@ export class BaseDataService<
    *
    * @param options - The options defining the query. Keep in mind that `queryKey` and `queryFn` are required when using data services.
    * Additionally `retry` and `retryDelay` are not available, retries can be customized using the `servicePolicyOptions`.
+   * @param options.struct - An optional struct for validating the response of the query function.
    * @param pageParam - An optional page parameter.
    * @returns The query result, exclusively the requested page is returned.
    */
@@ -278,13 +292,17 @@ export class BaseDataService<
     TQueryKey extends QueryKey = QueryKey,
     TPageParam extends Json = Json,
   >(
-    options: WithRequired<
+    {
+      struct,
+      ...options
+    }: WithRequired<
       OmitKeyof<
         FetchInfiniteQueryOptions<TQueryFnData, TError, TData, TQueryKey>,
         'retry' | 'retryDelay'
       >,
       'queryKey' | 'queryFn'
-    >,
+    > &
+      AdditionalQueryOptions<TQueryFnData>,
     pageParam?: TPageParam,
   ): Promise<TData> {
     const cache = this.#queryClient.getQueryCache();
@@ -297,12 +315,14 @@ export class BaseDataService<
       const result = await this.#queryClient.fetchInfiniteQuery({
         ...options,
         queryFn: (context) =>
-          this.#policy.execute(() =>
-            options.queryFn({
+          this.#policy.execute(async () => {
+            const response = await options.queryFn({
               ...context,
               pageParam: context.pageParam ?? pageParam,
-            }),
-          ),
+            });
+
+            return processQueryResponse(options.queryKey, response, struct);
+          }),
       });
 
       return result.pages[0];

packages/base-data-service/src/utils.ts

@@ -0,0 +1,32 @@
+import { Struct, validate } from '@metamask/superstruct';
+
+import { QueryKey } from './BaseDataService';
+
+/**
+ * Process query responses, validating them using Superstruct if a struct is defined.
+ *
+ * @param queryKey - The query key.
+ * @param response - The query response
+ * @param struct - The struct defining the schema for the query response.
+ * @returns The query response, coerced by Superstruct if needed.
+ * @throws If the query response does not match the struct.
+ */
+export function processQueryResponse<Response>(
+  queryKey: QueryKey,
+  response: Response,
+  struct?: Struct<Response>,
+): Response {
+  if (!struct) {
+    return response;
+  }
+
+  const [error, result] = validate(response, struct);
+
+  if (error) {
+    throw new Error(
+      `${queryKey[0]} returned an invalid response: ${error.message}.`,
+    );
+  }
+
+  return result;
+}

packages/base-data-service/tests/ExampleDataService.ts

@@ -1,5 +1,12 @@
 import { Messenger } from '@metamask/messenger';
-import { CaipAssetId, Duration, inMilliseconds, Json } from '@metamask/utils';
+import { object, number, string, array } from '@metamask/superstruct';
+import {
+  CaipAssetType,
+  CaipAssetTypeStruct,
+  Duration,
+  inMilliseconds,
+  Json,
+} from '@metamask/utils';
 import { ConstantBackoff } from 'cockatiel';
 
 import {
@@ -28,11 +35,20 @@ export type ExampleMessenger = Messenger<
 >;
 
 export type GetAssetsResponse = {
-  assetId: CaipAssetId;
+  assetId: CaipAssetType;
   decimals: number;
   name: string;
   symbol: string;
-};
+}[];
+
+const GetAssetsResponseStruct = array(
+  object({
+    assetId: CaipAssetTypeStruct,
+    decimals: number(),
+    name: string(),
+    symbol: string(),
+  }),
+);
 
 export type GetActivityResponse = {
   data: Json[];
@@ -102,6 +118,7 @@ export class ExampleDataService extends BaseDataService<
       },
       staleTime: inMilliseconds(1, Duration.Day),
       cacheTime: inMilliseconds(1, Duration.Day),
+      struct: GetAssetsResponseStruct,
     });
   }
 

yarn.lock

@@ -5893,6 +5893,7 @@ __metadata:
     "@metamask/auto-changelog": "npm:^6.1.0"
     "@metamask/messenger": "npm:^2.0.0"
     "@metamask/storage-service": "npm:^1.0.2"
+    "@metamask/superstruct": "npm:^3.1.0"
     "@metamask/utils": "npm:^11.11.0"
     "@tanstack/query-core": "npm:^4.43.0"
     "@ts-bridge/cli": "npm:^0.6.4"