feat(TM-source): use Tanstack query to cache requests

packages/source-iottwinmaker/package.json

@@ -58,6 +58,7 @@
     "@aws-sdk/client-secrets-manager": "3.353.0",
     "@aws-sdk/url-parser": "3.347.0",
     "@iot-app-kit/core": "6.2.0",
+    "@tanstack/query-core": "^4.29.11",
     "lodash": "^4.17.21",
     "rxjs": "^7.8.1"
   },

packages/source-iottwinmaker/src/initialize.ts

@@ -4,6 +4,7 @@ import { KinesisVideoClient } from '@aws-sdk/client-kinesis-video';
 import { KinesisVideoArchivedMediaClient } from '@aws-sdk/client-kinesis-video-archived-media';
 import { S3Client } from '@aws-sdk/client-s3';
 import { SecretsManagerClient } from '@aws-sdk/client-secrets-manager';
+import { QueryClient } from '@tanstack/query-core';
 
 import {
   kinesisVideoArchivedMediaSdk,
@@ -126,31 +127,42 @@ export const initialize = (
   const secretsManagerClient: SecretsManagerClient =
     authInput.secretsManagerClient ?? secretsManagersdk(inputWithCred.awsCredentials, inputWithCred.awsRegion);
 
-  const twinMakerMetadataModule = new TwinMakerMetadataModule(workspaceId, twinMakerClient);
+  // For caching TwinMaker API calls
+  const cachedQueryClient = new QueryClient({
+    defaultOptions: {
+      queries: {
+        staleTime: Infinity,
+      },
+    },
+  });
+
+  const twinMakerMetadataModule = new TwinMakerMetadataModule(workspaceId, twinMakerClient, cachedQueryClient);
   const twinMakerTimeSeriesModule = new TimeSeriesDataModule<TwinMakerDataStreamQuery>(
     createDataSource(twinMakerMetadataModule, twinMakerClient)
   );
 
+  const timeSeriesDataQuery: (query: TwinMakerQuery) => TimeSeriesDataQuery = (query: TwinMakerQuery) => ({
+    toQueryString: () =>
+      JSON.stringify({
+        source: SOURCE,
+        queryType: 'time-series-data',
+        query,
+      }),
+    build: (_sessionId: string, params: TimeSeriesDataRequest) =>
+      new TwinMakerTimeSeriesDataProvider(twinMakerMetadataModule, twinMakerTimeSeriesModule, {
+        queries: [
+          {
+            workspaceId,
+            ...query,
+          },
+        ],
+        request: params,
+      }),
+  });
+
   return {
     query: {
-      timeSeriesData: (query: TwinMakerQuery): TimeSeriesDataQuery => ({
-        toQueryString: () =>
-          JSON.stringify({
-            source: SOURCE,
-            queryType: 'time-series-data',
-            query,
-          }),
-        build: (_sessionId: string, params: TimeSeriesDataRequest) =>
-          new TwinMakerTimeSeriesDataProvider(twinMakerMetadataModule, twinMakerTimeSeriesModule, {
-            queries: [
-              {
-                workspaceId,
-                ...query,
-              },
-            ],
-            request: params,
-          }),
-      }),
+      timeSeriesData: (query: TwinMakerQuery): TimeSeriesDataQuery => timeSeriesDataQuery(query),
     },
     s3SceneLoader: (sceneId: string) => new S3SceneLoader({ workspaceId, sceneId, twinMakerClient, s3Client }),
     sceneMetadataModule: (sceneId: string) =>

packages/source-iottwinmaker/src/metadata-module/TwinMakerMetadataModule.spec.ts

@@ -1,14 +1,23 @@
 import { GetEntityResponse } from '@aws-sdk/client-iottwinmaker';
 import { ErrorDetails } from '@iot-app-kit/core';
 import { createMockTwinMakerSDK } from '../__mocks__/iottwinmakerSDK';
-import { TwinMakerMetadataCache } from './TwinMakerMetadataCache';
 import { TwinMakerMetadataModule } from './TwinMakerMetadataModule';
+import { QueryClient } from '@tanstack/query-core';
+
+const createCache = () =>
+  new QueryClient({
+    defaultOptions: {
+      queries: {
+        staleTime: Infinity,
+      },
+    },
+  });
 
 describe('TwinMakerMetadataModule', () => {
   const mockEntity1 = { entityId: 'test-1' } as GetEntityResponse;
   const mockEntity2 = { entityId: 'test-2' } as GetEntityResponse;
 
-  let cache: TwinMakerMetadataCache = new TwinMakerMetadataCache();
+  let cache: QueryClient = createCache();
   let module: TwinMakerMetadataModule;
 
   describe('fetchEntity', () => {
@@ -19,22 +28,18 @@ describe('TwinMakerMetadataModule', () => {
     beforeEach(() => {
       jest.clearAllMocks();
 
-      cache = new TwinMakerMetadataCache();
+      cache = createCache();
       module = new TwinMakerMetadataModule('ws-id', mockTMClient, cache);
     });
 
     it('should send request when the entity is not in the cache', async () => {
-      expect(cache.getEntity(mockEntity1.entityId!)).toBeUndefined();
-
       getEntity.mockResolvedValue(mockEntity1);
       const entity = await module.fetchEntity({ entityId: mockEntity1.entityId! });
       expect(getEntity).toBeCalledTimes(1);
       expect(entity).toEqual(mockEntity1);
     });
 
     it('should not send request when the request for the same entity is in process', async () => {
-      expect(cache.getEntity(mockEntity1.entityId!)).toBeUndefined();
-
       getEntity.mockResolvedValue(mockEntity1);
       const entities = await Promise.all([
         module.fetchEntity({ entityId: mockEntity1.entityId! }),
@@ -46,8 +51,6 @@ describe('TwinMakerMetadataModule', () => {
     });
 
     it('should send request when multiple calls for the different entities are triggered', async () => {
-      expect(cache.getEntity(mockEntity1.entityId!)).toBeUndefined();
-
       getEntity.mockResolvedValue(mockEntity1);
       await Promise.all([
         module.fetchEntity({ entityId: mockEntity1.entityId! }),
@@ -105,20 +108,16 @@ describe('TwinMakerMetadataModule', () => {
       getEntity.mockResolvedValueOnce(mockEntity1);
       getEntity.mockResolvedValueOnce(mockEntity2);
 
-      cache = new TwinMakerMetadataCache();
+      cache = createCache();
       module = new TwinMakerMetadataModule('ws-id', mockTMClient, cache);
     });
 
     it('should send request when the component type is not in the cache', async () => {
-      expect(cache.getEntitySummariesByComponentType(mockComponentTypeId)).toBeUndefined();
-
       const entities = await module.fetchEntitiesByComponentTypeId({ componentTypeId: mockComponentTypeId });
       expect(entities).toEqual([mockEntity1, mockEntity2]);
     });
 
     it('should not send request when the request for the same component type is in process', async () => {
-      expect(cache.getEntitySummariesByComponentType(mockComponentTypeId)).toBeUndefined();
-
       jest.spyOn(mockTMClient, 'send').mockResolvedValueOnce({} as never);
 
       await Promise.all([
@@ -129,8 +128,6 @@ describe('TwinMakerMetadataModule', () => {
     });
 
     it('should send request when multiple calls for the different component types are triggered', async () => {
-      expect(cache.getEntitySummariesByComponentType(mockComponentTypeId)).toBeUndefined();
-
       jest
         .spyOn(mockTMClient, 'send')
         .mockResolvedValueOnce({} as never)

packages/source-iottwinmaker/src/metadata-module/TwinMakerMetadataModule.ts

@@ -1,48 +1,39 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { GetEntityCommand, IoTTwinMakerClient, ListEntitiesCommand } from '@aws-sdk/client-iottwinmaker';
 import { isDefined } from '../time-series-data/utils/values';
-import { TwinMakerMetadataCache } from './TwinMakerMetadataCache';
 import type { EntitySummary, GetEntityResponse, ListEntitiesResponse } from '@aws-sdk/client-iottwinmaker';
 import type { ErrorDetails } from '@iot-app-kit/core';
+import { QueryClient } from '@tanstack/query-core';
 
 export class TwinMakerMetadataModule {
   private readonly workspaceId: string;
   private readonly tmClient: IoTTwinMakerClient;
-  private readonly cache: TwinMakerMetadataCache;
-
-  private readonly entityRequests: Record<string, Promise<GetEntityResponse>> = {};
-  private readonly listEntitiesByComponentTypeIdRequests: Record<string, Promise<EntitySummary[]>> = {};
+  private readonly cachedQueryClient: QueryClient;
 
   constructor(
     workspaceId: string,
     client: IoTTwinMakerClient,
-    cache: TwinMakerMetadataCache = new TwinMakerMetadataCache()
+    cachedQueryClient: QueryClient = new QueryClient({
+      defaultOptions: {
+        queries: {
+          staleTime: Infinity,
+        },
+      },
+    })
   ) {
     this.workspaceId = workspaceId;
     this.tmClient = client;
-    this.cache = cache;
+    this.cachedQueryClient = cachedQueryClient;
   }
 
   async fetchEntity({ entityId }: { entityId: string }): Promise<GetEntityResponse> {
-    const cachedEntity = this.cache.getEntity(entityId);
-    if (cachedEntity) {
-      return cachedEntity;
-    }
+    const request = new GetEntityCommand({ workspaceId: this.workspaceId, entityId });
 
     try {
-      if (this.entityRequests[entityId] === undefined) {
-        this.entityRequests[entityId] = this.tmClient.send(
-          new GetEntityCommand({
-            workspaceId: this.workspaceId,
-            entityId,
-          })
-        );
-      }
-
-      const response = await this.entityRequests[entityId];
-
-      this.cache.storeEntity(response);
-      delete this.entityRequests[entityId];
+      const response = await this.cachedQueryClient.fetchQuery({
+        queryKey: ['get-entity', this.workspaceId, entityId],
+        queryFn: () => this.tmClient.send(request),
+      });
 
       return response;
     } catch (err: any) {
@@ -53,19 +44,13 @@ export class TwinMakerMetadataModule {
   }
 
   async fetchEntitiesByComponentTypeId({ componentTypeId }: { componentTypeId: string }): Promise<GetEntityResponse[]> {
-    let summaries = this.cache.getEntitySummariesByComponentType(componentTypeId);
-
-    if (!summaries) {
-      if (this.listEntitiesByComponentTypeIdRequests[componentTypeId] === undefined) {
-        this.listEntitiesByComponentTypeIdRequests[componentTypeId] = this._requestEntitiesByComponentTypeId({
+    const summaries = await this.cachedQueryClient.fetchQuery({
+      queryKey: ['list-entities-by-component-type-id', this.workspaceId, componentTypeId],
+      queryFn: () =>
+        this._requestEntitiesByComponentTypeId({
           componentTypeId,
-        });
-      }
-
-      summaries = await this.listEntitiesByComponentTypeIdRequests[componentTypeId];
-
-      delete this.listEntitiesByComponentTypeIdRequests[componentTypeId];
-    }
+        }),
+    });
 
     const requests = summaries.map((summary) =>
       summary.entityId ? this.fetchEntity({ entityId: summary.entityId }) : undefined
@@ -94,8 +79,6 @@ export class TwinMakerMetadataModule {
         summaries.push(...(response.entitySummaries || []));
       } while (nextToken);
 
-      this.cache.storeEntitySummariesByComponentType(componentTypeId, summaries);
-
       return summaries;
     } catch (err: any) {
       const errorDetail: ErrorDetails = { msg: err.message, type: err.name, status: err.$metadata?.httpStatusCode };