Skip to content

Commit b57bf66

Browse files
salazarmsalazarm
committed
share cache manager
share cache manager share cache manager share cache manager share cache manager share cache manager
1 parent 26108aa commit b57bf66

File tree

3 files changed

+99
-80
lines changed

3 files changed

+99
-80
lines changed

‎js_modules/dagster-ui/packages/ui-core/src/assets/AssetsCatalogTable.tsx

+84-69
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ import {AssetGroupSelector} from '../graphql/types';
3636
import {useUpdatingRef} from '../hooks/useUpdatingRef';
3737
import {useBlockTraceUntilTrue} from '../performance/TraceContext';
3838
import {fetchPaginatedData} from '../runs/fetchPaginatedBucketData';
39-
import {CacheManager} from '../search/useIndexedDBCachedQuery';
39+
import {getCacheManager} from '../search/useIndexedDBCachedQuery';
4040
import {SyntaxError} from '../selection/CustomErrorListener';
4141
import {LoadingSpinner} from '../ui/Loading';
4242

@@ -54,7 +54,7 @@ export function useCachedAssets({
5454
}) {
5555
const {localCacheIdPrefix} = useContext(AppContext);
5656
const cacheManager = useMemo(
57-
() => new CacheManager<AssetTableFragment[]>(`${localCacheIdPrefix}/allAssetNodes`),
57+
() => getCacheManager<AssetTableFragment[]>(`${localCacheIdPrefix}/allAssetNodes`),
5858
[localCacheIdPrefix],
5959
);
6060

@@ -69,6 +69,11 @@ export function useCachedAssets({
6969
return {cacheManager};
7070
}
7171

72+
// Module-level cache variables
73+
let globalAssetsPromise: Promise<Asset[]> | null = null;
74+
let cachedAllAssets: Asset[] | null = null;
75+
let cachedAssetsFetchTime: number = 0;
76+
7277
export function useAllAssets({
7378
batchLimit = DEFAULT_BATCH_LIMIT,
7479
groupSelector,
@@ -77,14 +82,16 @@ export function useAllAssets({
7782
const [{error, assets}, setErrorAndAssets] = useState<{
7883
error: PythonErrorFragment | undefined;
7984
assets: Asset[] | undefined;
80-
}>({error: undefined, assets: undefined});
85+
}>({error: undefined, assets: cachedAllAssets || undefined});
8186

8287
const assetsRef = useUpdatingRef(assets);
8388

8489
const {cacheManager} = useCachedAssets({
8590
onAssetsLoaded: useCallback(
8691
(data) => {
92+
console.log('onAssetsLoaded', data);
8793
if (!assetsRef.current) {
94+
console.log('onAssetsLoaded', data);
8895
setErrorAndAssets({
8996
error: undefined,
9097
assets: data,
@@ -95,69 +102,57 @@ export function useAllAssets({
95102
),
96103
});
97104

98-
const allAssetsQuery = useCallback(async () => {
99-
if (groupSelector) {
100-
return;
105+
// Query function for all assets
106+
const fetchAllAssets = useCallback(async () => {
107+
if (cachedAllAssets && Date.now() - cachedAssetsFetchTime < 6000) {
108+
return cachedAllAssets;
101109
}
102-
try {
103-
const data = await fetchPaginatedData({
104-
async fetchData(cursor: string | null | undefined) {
105-
const {data} = await client.query<
106-
AssetCatalogTableQuery,
107-
AssetCatalogTableQueryVariables
108-
>({
109-
query: ASSET_CATALOG_TABLE_QUERY,
110-
fetchPolicy: 'no-cache',
111-
variables: {
112-
cursor,
113-
limit: batchLimit,
114-
},
115-
});
116-
117-
if (data.assetsOrError.__typename === 'PythonError') {
118-
return {
119-
data: [],
120-
cursor: undefined,
121-
hasMore: false,
122-
error: data.assetsOrError,
123-
};
124-
}
125-
const assets = data.assetsOrError.nodes;
126-
const hasMoreData = assets.length === batchLimit;
127-
const nextCursor = data.assetsOrError.cursor;
128-
return {
129-
data: assets,
130-
cursor: nextCursor,
131-
hasMore: hasMoreData,
132-
error: undefined,
133-
};
134-
},
135-
});
136-
cacheManager.set(data, AssetCatalogTableQueryVersion);
137-
setErrorAndAssets({error: undefined, assets: data});
138-
} catch (e: any) {
139-
if (e.__typename === 'PythonError') {
140-
setErrorAndAssets(({assets}) => ({
141-
error: e,
142-
assets,
143-
}));
144-
}
110+
if (!globalAssetsPromise) {
111+
globalAssetsPromise = (async () => {
112+
const allAssets = await fetchPaginatedData({
113+
async fetchData(cursor: string | null | undefined) {
114+
const {data} = await client.query<
115+
AssetCatalogTableQuery,
116+
AssetCatalogTableQueryVariables
117+
>({
118+
query: ASSET_CATALOG_TABLE_QUERY,
119+
fetchPolicy: 'no-cache',
120+
variables: {cursor, limit: batchLimit},
121+
});
122+
if (data.assetsOrError.__typename === 'PythonError') {
123+
return {
124+
data: [],
125+
cursor: undefined,
126+
hasMore: false,
127+
error: data.assetsOrError,
128+
};
129+
}
130+
const assets = data.assetsOrError.nodes;
131+
const hasMoreData = assets.length === batchLimit;
132+
const nextCursor = data.assetsOrError.cursor;
133+
return {data: assets, cursor: nextCursor, hasMore: hasMoreData, error: undefined};
134+
},
135+
});
136+
cachedAssetsFetchTime = Date.now();
137+
cachedAllAssets = allAssets;
138+
cacheManager.set(allAssets, AssetCatalogTableQueryVersion);
139+
globalAssetsPromise = null;
140+
return allAssets;
141+
})();
145142
}
146-
}, [batchLimit, cacheManager, client, groupSelector]);
143+
return globalAssetsPromise;
144+
}, [batchLimit, cacheManager, client]);
147145

146+
// Query function for group assets
148147
const groupQuery = useCallback(async () => {
149-
if (!groupSelector) {
150-
return;
151-
}
152-
function onData(queryData: typeof data) {
148+
const cacheKey = JSON.stringify(groupSelector);
149+
if (groupTableCache.has(cacheKey)) {
150+
const cachedData = groupTableCache.get(cacheKey);
153151
setErrorAndAssets({
154152
error: undefined,
155-
assets: queryData.assetNodes?.map(definitionToAssetTableFragment),
153+
assets: cachedData.assetNodes?.map(definitionToAssetTableFragment),
156154
});
157-
}
158-
const cacheKey = JSON.stringify(groupSelector);
159-
if (groupTableCache.has(cacheKey)) {
160-
onData(groupTableCache.get(cacheKey));
155+
return;
161156
}
162157
const {data} = await client.query<
163158
AssetCatalogGroupTableQuery,
@@ -168,23 +163,43 @@ export function useAllAssets({
168163
fetchPolicy: 'no-cache',
169164
});
170165
groupTableCache.set(cacheKey, data);
171-
onData(data);
172-
}, [groupSelector, client]);
166+
setErrorAndAssets({
167+
error: undefined,
168+
assets: data.assetNodes?.map(definitionToAssetTableFragment),
169+
});
170+
}, [client, groupSelector]);
173171

174-
const query = groupSelector ? groupQuery : allAssetsQuery;
172+
const query = groupSelector ? groupQuery : fetchAllAssets;
175173

176174
useEffect(() => {
177-
query();
178-
}, [query]);
175+
if (groupSelector) {
176+
groupQuery();
177+
} else {
178+
fetchAllAssets()
179+
.then((allAssets) => setErrorAndAssets({error: undefined, assets: allAssets}))
180+
.catch((e: any) => {
181+
if (e.__typename === 'PythonError') {
182+
setErrorAndAssets((prev) => ({error: e, assets: prev.assets}));
183+
}
184+
});
185+
}
186+
}, [fetchAllAssets, groupQuery, groupSelector]);
179187

180-
return useMemo(() => {
181-
return {
188+
const assetsByAssetKey = useMemo(
189+
() => Object.fromEntries(assets?.map((asset) => [tokenForAssetKey(asset.key), asset]) ?? []),
190+
[assets],
191+
);
192+
193+
return useMemo(
194+
() => ({
182195
assets,
196+
assetsByAssetKey,
183197
error,
184198
loading: !assets && !error,
185199
query,
186-
};
187-
}, [assets, error, query]);
200+
}),
201+
[assets, assetsByAssetKey, error, query],
202+
);
188203
}
189204

190205
interface AssetCatalogTableProps {
@@ -253,7 +268,7 @@ export const AssetsCatalogTable = ({
253268
[filtered, prefixPath, view],
254269
);
255270

256-
const refreshState = useRefreshAtInterval({
271+
const refreshState = useRefreshAtInterval<any>({
257272
refresh: query,
258273
intervalMs: 4 * FIFTEEN_SECONDS,
259274
leading: true,

‎js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/useAllAssets.test.tsx

-3
Original file line numberDiff line numberDiff line change
@@ -112,9 +112,6 @@ describe('useAllAssets', () => {
112112
},
113113
});
114114

115-
await waitFor(() => {
116-
expect(result.current.assets?.length).toBe(1);
117-
});
118115
await waitFor(() => {
119116
expect(result.current.assets?.length).toBe(5);
120117
});

‎js_modules/dagster-ui/packages/ui-core/src/search/useIndexedDBCachedQuery.tsx

+15-8
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import memoize from 'lodash/memoize';
21
import React, {useCallback, useEffect, useMemo} from 'react';
32

43
import {
@@ -12,6 +11,7 @@ import {usePreviousDistinctValue} from '../hooks/usePrevious';
1211
import {useUpdatingRef} from '../hooks/useUpdatingRef';
1312
import {CompletionType, useBlockTraceUntilTrue} from '../performance/TraceContext';
1413
import {cache} from '../util/idb-lru-cache';
14+
import {weakMapMemoize} from '../util/weakMapMemoize';
1515

1616
type CacheData<TQuery> = {
1717
data: TQuery;
@@ -20,7 +20,7 @@ type CacheData<TQuery> = {
2020

2121
export const KEY_PREFIX = 'indexdbQueryCache:';
2222

23-
export class CacheManager<TQuery> {
23+
class CacheManager<TQuery> {
2424
private cache: ReturnType<typeof cache<CacheData<TQuery>>> | undefined;
2525
private key: string;
2626
private current?: CacheData<TQuery>;
@@ -63,8 +63,7 @@ export class CacheManager<TQuery> {
6363
async set(data: TQuery, version: number | string): Promise<void> {
6464
if (
6565
this.current?.data === data ||
66-
(JSON.stringify(this.current?.data) === JSON.stringify(data) &&
67-
this.current?.version === version)
66+
(stringified(this.current?.data) === stringified(data) && this.current?.version === version)
6867
) {
6968
return;
7069
}
@@ -249,7 +248,7 @@ export function useIndexedDBCachedQuery<TQuery, TVariables extends OperationVari
249248
data &&
250249
// Work around a weird jest issue where it returns an empty object if no mocks are found...
251250
Object.keys(data).length &&
252-
(!dataRef.current || JSON.stringify(dataRef.current) !== JSON.stringify(data))
251+
(!dataRef.current || stringified(dataRef.current) !== stringified(data))
253252
) {
254253
setData(data);
255254
}
@@ -348,7 +347,7 @@ export function useGetData() {
348347

349348
export function useGetCachedData() {
350349
return useCallback(async <TQuery,>({key, version}: {key: string; version: number | string}) => {
351-
const cacheManager = new CacheManager<TQuery>(key);
350+
const cacheManager = getCacheManager<TQuery>(key);
352351
return await cacheManager.get(version);
353352
}, []);
354353
}
@@ -359,13 +358,21 @@ export function useClearCachedData() {
359358
}, []);
360359
}
361360

362-
let getCacheManager = memoize(<TQuery,>(key: string) => {
361+
export let getCacheManager = weakMapMemoize(<TQuery,>(key: string) => {
363362
return new CacheManager<TQuery>(key);
364363
});
365364

366365
export const __resetForJest = () => {
367366
Object.keys(globalFetchStates).forEach((key) => delete globalFetchStates[key]);
368-
getCacheManager = memoize(<TQuery,>(key: string) => {
367+
getCacheManager = weakMapMemoize(<TQuery,>(key: string) => {
369368
return new CacheManager<TQuery>(key);
370369
});
371370
};
371+
372+
const stringified = weakMapMemoize((data: any) => {
373+
try {
374+
return JSON.stringify(data);
375+
} catch {
376+
return '';
377+
}
378+
});

0 commit comments

Comments
 (0)