fix(react-components): asset search hook to support more than 1000 results

react-components/src/components/RuleBasedOutputs/RuleBasedOutputsSelector.tsx

@@ -4,14 +4,13 @@
 import { useEffect, type ReactElement, useState, useMemo } from 'react';
 
 import { CogniteCadModel } from '@cognite/reveal';
-import { type ModelMappingsWithAssets, useAllMappedEquipmentAssetMappings } from '../..';
+import { useAllMappedEquipmentAssetMappings } from '../..';
 import { type RuleOutputSet, type AssetStylingGroupAndStyleIndex } from './types';
 import { generateRuleBasedOutputs, traverseExpressionToGetTimeseries } from './utils';
 import { use3dModels } from '../../hooks/use3dModels';
 import { EMPTY_ARRAY } from '../../utilities/constants';
-import { type Datapoints, type Asset } from '@cognite/sdk';
+import { type Datapoints, type Asset, type AssetMapping3D } from '@cognite/sdk';
 import { isDefined } from '../../utilities/isDefined';
-import { type InfiniteData } from '@tanstack/react-query';
 import { type AssetIdsAndTimeseries } from '../../utilities/types';
 import { useAssetsAndTimeseriesLinkageDataQuery } from '../../query/useAssetsAndTimeseriesLinkageDataQuery';
 
@@ -31,7 +30,7 @@ export function RuleBasedOutputsSelector({
   const [stylingGroups, setStylingsGroups] = useState<AssetStylingGroupAndStyleIndex[]>();
 
   const {
-    data: assetMappings,
+    data: modelAssetPage,
     isFetching,
     hasNextPage,
     fetchNextPage
@@ -45,12 +44,24 @@ export function RuleBasedOutputsSelector({
 
   const contextualizedAssetNodes = useMemo(() => {
     return (
-      assetMappings?.pages
+      modelAssetPage?.pages
         .flat()
-        .flatMap((item) => item.assets)
+        .flatMap((modelsAssetPage) =>
+          modelsAssetPage.modelsAssets.flatMap((modelsAsset) => modelsAsset.assets)
+        )
         .map(convertAssetMetadataKeysToLowerCase) ?? []
     );
-  }, [assetMappings]);
+  }, [modelAssetPage]);
+
+  const assetMappings = useMemo(() => {
+    return (
+      modelAssetPage?.pages
+        .flat()
+        .flatMap((modelAssetPage) => modelAssetPage.modelsAssets)
+        .flatMap((node) => node.mappings)
+        .filter(isDefined) ?? []
+    );
+  }, [modelAssetPage]);
 
   const timeseriesExternalIds = useMemo(() => {
     const expressions = ruleSet?.rulesWithOutputs
@@ -70,7 +81,7 @@ export function RuleBasedOutputsSelector({
   }, [stylingGroups]);
 
   useEffect(() => {
-    if (assetMappings === undefined || models === undefined || isFetching) return;
+    if (modelAssetPage === undefined || models === undefined || isFetching) return;
     if (timeseriesExternalIds.length > 0 && isLoadingAssetIdsAndTimeseriesData) return;
 
     setStylingsGroups(EMPTY_ARRAY);
@@ -81,10 +92,11 @@ export function RuleBasedOutputsSelector({
       if (!(model instanceof CogniteCadModel)) {
         return;
       }
+      const assetMappingsData = assetMappings.map((response) => response.items).flat();
       setStylingsGroups(
         await initializeRuleBasedOutputs({
           model,
-          assetMappings,
+          assetMappings: assetMappingsData,
           contextualizedAssetNodes,
           ruleSet,
           assetIdsAndTimeseries: assetIdsWithTimeseriesData?.assetIdsWithTimeseries ?? [],
@@ -106,19 +118,16 @@ async function initializeRuleBasedOutputs({
   timeseriesDatapoints
 }: {
   model: CogniteCadModel;
-  assetMappings: InfiniteData<ModelMappingsWithAssets[]>;
+  assetMappings: AssetMapping3D[];
   contextualizedAssetNodes: Asset[];
   ruleSet: RuleOutputSet;
   assetIdsAndTimeseries: AssetIdsAndTimeseries[];
   timeseriesDatapoints: Datapoints[] | undefined;
 }): Promise<AssetStylingGroupAndStyleIndex[]> {
-  const flatAssetsMappingsList = assetMappings.pages.flat().flatMap((item) => item.mappings);
-  const flatMappings = flatAssetsMappingsList.flatMap((node) => node.items);
-
   const collectionStylings = await generateRuleBasedOutputs({
     model,
     contextualizedAssetNodes,
-    assetMappings: flatMappings,
+    assetMappings,
     ruleSet,
     assetIdsAndTimeseries,
     timeseriesDatapoints

react-components/src/hooks/network/getAssetsList.ts

@@ -0,0 +1,44 @@
+/*!
+ * Copyright 2024 Cognite AS
+ */
+import { type CogniteClient, type CursorResponse, type Asset } from '@cognite/sdk';
+import { buildFilter } from '../../utilities/buildFilter';
+
+export const getAssetsList = async (
+  sdk: CogniteClient,
+  {
+    query,
+    cursor,
+    limit = 1000,
+    sort = [{ property: ['_score_'] }],
+    aggregatedProperties = ['path']
+  }: {
+    query: string;
+    cursor?: string;
+    limit?: number;
+    sort?: Array<{ property: string[] }>;
+    aggregatedProperties?: string[];
+  }
+): Promise<{ items: Asset[]; nextCursor: string | undefined }> => {
+  const advancedFilter = buildFilter(query);
+
+  return await sdk
+    .post<CursorResponse<Asset[]>>(`/api/v1/projects/${sdk.project}/assets/list`, {
+      headers: {
+        'cdf-version': 'alpha'
+      },
+      data: {
+        limit,
+        sort,
+        advancedFilter,
+        aggregatedProperties,
+        cursor
+      }
+    })
+    .then(({ data }) => {
+      return {
+        items: data.items,
+        nextCursor: data.nextCursor
+      };
+    });
+};

react-components/src/query/useSearchMappedEquipmentAssetMappings.tsx

@@ -1,6 +1,7 @@
 /*!
  * Copyright 2023 Cognite AS
  */
+import { useRef } from 'react';
 import { type AddModelOptions } from '@cognite/reveal';
 import {
   type Asset,
@@ -10,13 +11,13 @@ import {
 } from '@cognite/sdk';
 import {
   type UseInfiniteQueryResult,
-  type UseQueryResult,
   useInfiniteQuery,
-  useQuery,
   type InfiniteData
 } from '@tanstack/react-query';
 import { useSDK } from '../components/RevealCanvas/SDKProvider';
-import { chunk } from 'lodash';
+import { getAssetsList } from '../hooks/network/getAssetsList';
+import { useAssetMappedNodesForRevisions } from '../components/CacheProvider/AssetMappingCacheProvider';
+import { isDefined } from '../utilities/isDefined';
 
 export type ModelMappings = {
   model: AddModelOptions;
@@ -27,53 +28,92 @@ export type ModelMappingsWithAssets = ModelMappings & {
   assets: Asset[];
 };
 
+export type AssetPage = {
+  assets: Asset[];
+  nextCursor: string | undefined;
+};
+
+export type ModelAssetPage = {
+  modelsAssets: ModelMappingsWithAssets[];
+  nextCursor: string | undefined;
+};
+
 export const useSearchMappedEquipmentAssetMappings = (
   query: string,
   models: AddModelOptions[],
   limit: number = 100,
   userSdk?: CogniteClient
-): UseQueryResult<Asset[]> => {
+): UseInfiniteQueryResult<InfiniteData<AssetPage[]>, Error> => {
   const sdk = useSDK(userSdk);
-  const modelsKey = models.map((model) => [model.modelId, model.revisionId]);
-  const { data: assetMappings, isFetched } = useAllMappedEquipmentAssetMappings(models, sdk);
+  const { data: assetMappingList, isFetched } = useAssetMappedNodesForRevisions(
+    models.map((model) => ({ ...model, type: 'cad' }))
+  );
+  const allAssetMappings = useAllMappedEquipmentAssetMappings(models, sdk);
 
-  return useQuery({
-    queryKey: ['reveal', 'react-components', 'search-mapped-asset-mappings', query, modelsKey],
-    queryFn: async () => {
+  return useInfiniteQuery({
+    queryKey: [
+      'reveal',
+      'react-components',
+      'search-mapped-asset-mappings',
+      query,
+      ...models.map((model) => [model.modelId, model.revisionId])
+    ],
+    queryFn: async ({ pageParam }: { pageParam: string | undefined }) => {
+      if (allAssetMappings.data === undefined) {
+        return { assets: [], nextCursor: undefined };
+      }
       if (query === '') {
-        const mappedAssets =
-          assetMappings?.pages
+        const assets = allAssetMappings.data?.pages.flatMap((modelWithAssets) =>
+          modelWithAssets
+            .map((modelWithAsset) =>
+              modelWithAsset.modelsAssets.flatMap((modelsAsset) => modelsAsset.assets)
+            )
             .flat()
-            .map((item) => item.assets)
-            .flat() ?? [];
-        return mappedAssets;
+        );
+        return { assets, nextCursor: undefined };
       }
-
-      const searchedAssets = await sdk.assets.search({ search: { query }, limit: 1000 });
-      const assetMappingsWithSearch = await getAssetMappingsByModels(
-        sdk,
-        models,
+      if (assetMappingList === undefined) {
+        return { assets: [], nextCursor: undefined };
+      }
+      const assetsResponse = await getAssetsList(sdk, {
+        query,
         limit,
-        searchedAssets.map((asset) => asset.id)
-      );
+        cursor: pageParam
+      });
+
+      const assets = assetsResponse.items.filter((asset) => asset !== undefined);
 
-      const assetMappingsSet = createAssetMappingsSet(assetMappingsWithSearch);
-      const filteredSearchedAssets = searchedAssets.filter((asset) =>
-        assetMappingsSet.has(asset.id)
-      );
+      const assetMappingsWithSearch = assetMappingList.flatMap((mapping) => {
+        return mapping.assetMappings.filter((assetMapping) =>
+          assets.map((asset) => asset.id).includes(assetMapping.assetId)
+        );
+      });
 
-      return filteredSearchedAssets;
+      const assetMappingsSet = new Set(assetMappingsWithSearch.map((mapping) => mapping.assetId));
+      const filteredSearchedAssets = assets.filter((asset) => assetMappingsSet.has(asset.id));
+
+      return {
+        assets: filteredSearchedAssets,
+        nextCursor: assetsResponse.nextCursor
+      };
     },
+    initialPageParam: undefined,
     staleTime: Infinity,
-    enabled: isFetched && assetMappings !== undefined
+    getNextPageParam: (_lastPage, allPages) => {
+      const lastPageData = allPages[allPages.length - 1];
+      return lastPageData.nextCursor;
+    },
+    enabled: isFetched && assetMappingList !== undefined && assetMappingList.length > 0
   });
 };
 
 export const useAllMappedEquipmentAssetMappings = (
   models: AddModelOptions[],
-  userSdk?: CogniteClient
-): UseInfiniteQueryResult<InfiniteData<ModelMappingsWithAssets[]>, Error> => {
+  userSdk?: CogniteClient,
+  limit: number = 1000
+): UseInfiniteQueryResult<InfiniteData<ModelAssetPage[]>, Error> => {
   const sdk = useSDK(userSdk);
+  const usedCursors = useRef(new Set());
 
   return useInfiniteQuery({
     queryKey: [
@@ -101,21 +141,46 @@ export const useAllMappedEquipmentAssetMappings = (
 
         const mappings = await sdk.assetMappings3D.filter(model.modelId, model.revisionId, {
           cursor: nextCursor === 'start' ? undefined : nextCursor,
-          limit: 1000
+          limit
         });
 
+        usedCursors.current.add(nextCursor);
+
         return { mappings, model };
       });
 
       const currentPagesOfAssetMappings = await Promise.all(currentPagesOfAssetMappingsPromises);
 
       const modelsAssets = await getAssetsFromAssetMappings(sdk, currentPagesOfAssetMappings);
+      const nextCursors = currentPagesOfAssetMappings
+        .map(({ mappings }) => mappings.nextCursor)
+        .filter(isDefined);
 
-      return modelsAssets;
+      return await Promise.resolve({
+        modelsAssets,
+        nextCursors
+      });
     },
     initialPageParam: models.map((model) => ({ cursor: 'start', model })),
     staleTime: Infinity,
-    getNextPageParam
+    getNextPageParam: (lastPage: {
+      modelsAssets: ModelMappingsWithAssets[];
+      nextCursors: string[];
+    }): Array<{ cursor: string | undefined; model: AddModelOptions }> | undefined => {
+      const nextCursors = lastPage.nextCursors
+        .map((cursor, index) => ({ cursor, model: lastPage.modelsAssets[index].model }))
+        .filter((mappingModel) => {
+          if (mappingModel.cursor === undefined || usedCursors.current.has(mappingModel.cursor)) {
+            return false;
+          }
+          usedCursors.current.add(mappingModel.cursor);
+          return true;
+        });
+      if (nextCursors.length === 0) {
+        return undefined;
+      }
+      return nextCursors;
+    }
   });
 };
 
@@ -183,39 +248,6 @@ function getNextPageParam(
   return nextCursors;
 }
 
-async function getAssetMappingsByModels(
-  sdk: CogniteClient,
-  models: AddModelOptions[],
-  limit: number = 1000,
-  assetIdsFilter?: number[]
-): Promise<ModelMappings[]> {
-  const mappedEquipmentPromises = models.map(async (model) => {
-    if (assetIdsFilter === undefined) {
-      const mappings = await sdk.assetMappings3D.filter(model.modelId, model.revisionId, {
-        limit
-      });
-      return [{ mappings, model }];
-    }
-
-    const deduplicatedAssetIds = Array.from(new Set(assetIdsFilter));
-    const chunkedFilter = chunk(deduplicatedAssetIds, 100);
-
-    const chunkedPromises = chunkedFilter.map(async (chunk) => {
-      const mappings = await sdk.assetMappings3D.filter(model.modelId, model.revisionId, {
-        filter: { assetIds: chunk },
-        limit
-      });
-      return { mappings, model };
-    });
-
-    return await Promise.all(chunkedPromises);
-  });
-
-  const mappedEquipment = await Promise.all(mappedEquipmentPromises);
-
-  return mappedEquipment.flat();
-}
-
 async function getAssetsFromAssetMappings(
   sdk: CogniteClient,
   modelsMappings: Array<{ model: AddModelOptions; mappings: ListResponse<AssetMapping3D[]> }>
@@ -239,11 +271,3 @@ async function getAssetsFromAssetMappings(
 
   return mappingsWithAssets;
 }
-
-function createAssetMappingsSet(
-  assetMappings: Array<{ model: AddModelOptions; mappings: ListResponse<AssetMapping3D[]> }>
-): Set<number> {
-  return new Set(
-    assetMappings.map(({ mappings }) => mappings.items.map((item) => item.assetId)).flat()
-  );
-}