[1/2] Asset Health Observe UIs phase 1

js_modules/dagster-ui/packages/ui-components/src/components/Icon.tsx

@@ -146,6 +146,7 @@ import expand_arrows from '../icon-svgs/expand_arrows.svg';
 import expand_less from '../icon-svgs/expand_less.svg';
 import expand_more from '../icon-svgs/expand_more.svg';
 import expectation from '../icon-svgs/expectation.svg';
+import failure_trend from '../icon-svgs/failure_trend.svg';
 import file_csv from '../icon-svgs/file_csv.svg';
 import file_json from '../icon-svgs/file_json.svg';
 import file_markdown from '../icon-svgs/file_markdown.svg';
@@ -351,6 +352,7 @@ import sticky_note from '../icon-svgs/sticky_note.svg';
 import storage_kind from '../icon-svgs/storage_kind.svg';
 import subtract from '../icon-svgs/subtract.svg';
 import success from '../icon-svgs/success.svg';
+import successful_trend from '../icon-svgs/successful_trend.svg';
 import sun from '../icon-svgs/sun.svg';
 import support from '../icon-svgs/support.svg';
 import sync from '../icon-svgs/sync.svg';
@@ -393,6 +395,7 @@ import visibility from '../icon-svgs/visibility.svg';
 import visibility_off from '../icon-svgs/visibility_off.svg';
 import warning from '../icon-svgs/warning.svg';
 import warning_outline from '../icon-svgs/warning_outline.svg';
+import warning_trend from '../icon-svgs/warning_trend.svg';
 import water from '../icon-svgs/water.svg';
 import waterfall_chart from '../icon-svgs/waterfall_chart.svg';
 import webhook from '../icon-svgs/webhook.svg';
@@ -559,6 +562,7 @@ export const Icons = {
   expand_less,
   expand_more,
   expectation,
+  failure_trend,
   file_csv,
   file_json,
   file_markdown,
@@ -760,6 +764,7 @@ export const Icons = {
   storage_kind,
   subtract,
   success,
+  successful_trend,
   sun,
   support,
   sync,
@@ -802,6 +807,7 @@ export const Icons = {
   visibility_off,
   warning,
   warning_outline,
+  warning_trend,
   water,
   waterfall_chart,
   webhook,

js_modules/dagster-ui/packages/ui-components/src/components/useDelayedState.tsx

@@ -1,7 +1,8 @@
-import {useEffect, useState} from 'react';
+import {useDeferredValue, useEffect, useState} from 'react';
 
 export const useDelayedState = (delayMsec: number) => {
-  const [ready, setReady] = useState(false);
+  const [_ready, setReady] = useState(false);
+  const ready = useDeferredValue(_ready);
 
   useEffect(() => {
     const timer = setTimeout(() => setReady(true), delayMsec);

js_modules/dagster-ui/packages/ui-core/src/asset-data/AssetHealthDataProvider.tsx

@@ -0,0 +1,82 @@
+import React from 'react';
+
+import {ApolloClient, gql, useApolloClient} from '../apollo-client';
+import {AssetBaseData} from './AssetBaseDataProvider';
+import {tokenForAssetKey, tokenToAssetKey} from '../asset-graph/Utils';
+import {AssetKeyInput} from '../graphql/types';
+import {liveDataFactory} from '../live-data-provider/Factory';
+import {LiveDataThreadID} from '../live-data-provider/LiveDataThread';
+import {useBlockTraceUntilTrue} from '../performance/TraceContext';
+import {AssetHealthQuery, AssetHealthQueryVariables} from './types/AssetHealthDataProvider.types';
+
+function init() {
+  return liveDataFactory(
+    () => {
+      return useApolloClient();
+    },
+    async (keys, client: ApolloClient<any>) => {
+      const assetKeys = keys.map(tokenToAssetKey);
+      const healthResponse = await client.query<AssetHealthQuery, AssetHealthQueryVariables>({
+        query: ASSETS_HEALTH_INFO_QUERY,
+        fetchPolicy: 'no-cache',
+        variables: {
+          assetKeys,
+        },
+      });
+
+      const {data} = healthResponse;
+
+      return Object.fromEntries(
+        data.assetNodes.map((node) => [tokenForAssetKey(node.assetKey), node]),
+      );
+    },
+  );
+}
+export const AssetHealthData = init();
+
+export function useAssetHealthData(assetKey: AssetKeyInput, thread: LiveDataThreadID = 'default') {
+  const result = AssetHealthData.useLiveDataSingle(tokenForAssetKey(assetKey), thread);
+  useBlockTraceUntilTrue('useAssetHealthData', !!result.liveData);
+  return result;
+}
+
+export function useAssetsHealthData(
+  assetKeys: AssetKeyInput[],
+  thread: LiveDataThreadID = 'AssetHealth', // Use AssetHealth to get 250 batch size
+) {
+  const keys = React.useMemo(() => assetKeys.map((key) => tokenForAssetKey(key)), [assetKeys]);
+  const result = AssetHealthData.useLiveData(keys, thread);
+  AssetBaseData.useLiveData(keys, thread);
+  useBlockTraceUntilTrue(
+    'useAssetsHealthData',
+    !!(Object.keys(result.liveDataByNode).length === assetKeys.length),
+  );
+  return result;
+}
+
+export const ASSETS_HEALTH_INFO_QUERY = gql`
+  query AssetHealthQuery($assetKeys: [AssetKeyInput!]!) {
+    assetNodes(assetKeys: $assetKeys) {
+      id
+      ...AssetHealthFragment
+    }
+  }
+
+  fragment AssetHealthFragment on AssetNode {
+    assetKey {
+      path
+    }
+
+    assetHealth {
+      assetHealth
+      materializationStatus
+      assetChecksStatus
+      freshnessStatus
+    }
+  }
+`;
+
+// For tests
+export function __resetForJest() {
+  Object.assign(AssetHealthData, init());
+}

js_modules/dagster-ui/packages/ui-core/src/asset-data/AssetLiveDataProvider.tsx

@@ -2,6 +2,7 @@ import uniq from 'lodash/uniq';
 import React, {useCallback, useMemo, useRef} from 'react';
 
 import {AssetBaseData, __resetForJest as __resetBaseData} from './AssetBaseDataProvider';
+import {AssetHealthData} from './AssetHealthDataProvider';
 import {
   AssetStaleStatusData,
   __resetForJest as __resetStaleData,
@@ -90,10 +91,17 @@ export const AssetLiveDataProvider = ({children}: {children: React.ReactNode}) =
 
   const staleKeysObserved = useRef<string[]>([]);
   const baseKeysObserved = useRef<string[]>([]);
+  const healthKeysObserved = useRef<string[]>([]);
 
   React.useEffect(() => {
     const onSubscriptionsChanged = () => {
-      const keys = Array.from(new Set(...staleKeysObserved.current, ...baseKeysObserved.current));
+      const keys = Array.from(
+        new Set([
+          ...staleKeysObserved.current,
+          ...baseKeysObserved.current,
+          ...healthKeysObserved.current,
+        ]),
+      );
       setAllObservedKeys(keys.map((key) => ({path: key.split('/')})));
     };
 
@@ -105,18 +113,24 @@ export const AssetLiveDataProvider = ({children}: {children: React.ReactNode}) =
       baseKeysObserved.current = keys;
       onSubscriptionsChanged();
     });
+    AssetHealthData.manager.setOnSubscriptionsChangedCallback((keys) => {
+      healthKeysObserved.current = keys;
+      onSubscriptionsChanged();
+    });
   }, []);
 
   const pollRate = React.useContext(LiveDataPollRateContext);
 
   React.useEffect(() => {
     AssetStaleStatusData.manager.setPollRate(pollRate);
     AssetBaseData.manager.setPollRate(pollRate);
+    AssetHealthData.manager.setPollRate(pollRate);
   }, [pollRate]);
 
   useDidLaunchEvent(() => {
     AssetStaleStatusData.manager.invalidateCache();
     AssetBaseData.manager.invalidateCache();
+    AssetHealthData.manager.invalidateCache();
   }, SUBSCRIPTION_MAX_POLL_RATE);
 
   React.useEffect(() => {
@@ -148,15 +162,18 @@ export const AssetLiveDataProvider = ({children}: {children: React.ReactNode}) =
       ) {
         AssetBaseData.manager.invalidateCache();
         AssetStaleStatusData.manager.invalidateCache();
+        AssetHealthData.manager.invalidateCache();
       }
     });
     return unobserve;
   }, [allObservedKeys]);
 
   return (
-    <AssetBaseData.LiveDataProvider>
-      <AssetStaleStatusData.LiveDataProvider>{children}</AssetStaleStatusData.LiveDataProvider>
-    </AssetBaseData.LiveDataProvider>
+    <AssetHealthData.LiveDataProvider>
+      <AssetBaseData.LiveDataProvider>
+        <AssetStaleStatusData.LiveDataProvider>{children}</AssetStaleStatusData.LiveDataProvider>
+      </AssetBaseData.LiveDataProvider>
+    </AssetHealthData.LiveDataProvider>
   );
 };
 

js_modules/dagster-ui/packages/ui-core/src/asset-data/__tests__/AssetLiveDataProvider.test.tsx

@@ -577,4 +577,47 @@ describe('AssetLiveDataProvider', () => {
       expect(resultFn4).toHaveBeenCalled();
     });
   });
+
+  it('should not return live data for keys that we unsubscribe from by changing the keys passed to the hook', async () => {
+    const assetKeys = [buildAssetKey({path: ['key1']}), buildAssetKey({path: ['key2']})];
+    const [mockedQuery, mockedFreshnessQuery] = buildMockedAssetGraphLiveQuery(assetKeys);
+
+    const resultFn = getMockResultFn(mockedQuery);
+
+    const hookResult = jest.fn();
+
+    const {rerender} = render(
+      <Test mocks={[mockedQuery, mockedFreshnessQuery]} hooks={[{keys: assetKeys, hookResult}]} />,
+    );
+
+    // Initially an empty object
+    expect(resultFn).toHaveBeenCalledTimes(0);
+    expect(hookResult.mock.calls[0]!.value).toEqual(undefined);
+
+    act(() => {
+      jest.runOnlyPendingTimers();
+    });
+
+    expect(resultFn).toHaveBeenCalled();
+    await waitFor(() => {
+      expect(hookResult.mock.calls[1][0]!).toEqual({
+        ['key1']: expect.any(Object),
+        ['key2']: expect.any(Object),
+      });
+    });
+
+    // Re-render with different asset keys (only asset key1 is in the new set)
+
+    const assetKeys2 = [buildAssetKey({path: ['key1']})];
+    const hookResult2 = jest.fn();
+    rerender(<Test mocks={[mockedQuery]} hooks={[{keys: assetKeys2, hookResult: hookResult2}]} />);
+
+    act(() => {
+      jest.runOnlyPendingTimers();
+    });
+
+    expect(hookResult2.mock.calls[1][0]).toEqual({
+      ['key1']: expect.any(Object),
+    });
+  });
 });

js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/useAssetSelectionInput.tsx

@@ -1,3 +1,4 @@
+import {useMemo} from 'react';
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 import {AssetGraphAssetSelectionInput} from 'shared/asset-graph/AssetGraphAssetSelectionInput.oss';
 import {AssetSelectionInput} from 'shared/asset-selection/input/AssetSelectionInput.oss';
@@ -30,26 +31,27 @@ export const useAssetSelectionInput = <
     loading: !!assetsLoading,
   });
 
-  let filterInput = (
-    <AssetGraphAssetSelectionInput
-      items={graphQueryItems}
-      value={assetSelection}
-      placeholder="Type an asset subset…"
-      onChange={setAssetSelection}
-      popoverPosition="bottom-left"
-    />
-  );
-
-  if (featureEnabled(FeatureFlag.flagSelectionSyntax)) {
-    filterInput = (
-      <AssetSelectionInput
+  const filterInput = useMemo(() => {
+    if (featureEnabled(FeatureFlag.flagSelectionSyntax)) {
+      return (
+        <AssetSelectionInput
+          value={assetSelection}
+          onChange={setAssetSelection}
+          assets={graphQueryItems}
+          onErrorStateChange={onErrorStateChange}
+        />
+      );
+    }
+    return (
+      <AssetGraphAssetSelectionInput
+        items={graphQueryItems}
         value={assetSelection}
+        placeholder="Type an asset subset…"
         onChange={setAssetSelection}
-        assets={graphQueryItems}
-        onErrorStateChange={onErrorStateChange}
+        popoverPosition="bottom-left"
       />
     );
-  }
+  }, [assetSelection, graphQueryItems, onErrorStateChange, setAssetSelection]);
 
   return {filterInput, loading, filtered: filtered as T[], assetSelection, setAssetSelection};
 };