[INTERNAL_BRANCH=status-selection-syntax]

Author: salazarm

Diff for: js_modules/dagster-ui/packages/ui-core/package.json

@@ -19,7 +19,8 @@
     "generate-run-selection": "ts-node -O '{\"module\": \"commonjs\"}' ./src/scripts/generateRunSelection.ts && eslint src/run-selection/generated/ --fix -c .eslintrc.js",
     "generate-op-selection": "ts-node -O '{\"module\": \"commonjs\"}' ./src/scripts/generateOpSelection.ts && eslint src/op-selection/generated/ --fix -c .eslintrc.js",
     "storybook": "storybook dev -p 6006",
-    "build-storybook": "storybook build"
+    "build-storybook": "storybook build",
+    "find-circular-dependencies": "ts-node -O '{\"module\": \"commonjs\"}' ./src/scripts/findCircularDependencies.ts"
   },
   "peerDependencies": {
     "@apollo/client": "3.9.11",

Diff for: js_modules/dagster-ui/packages/ui-core/src/app/UserSettingsDialog/UserSettingsDialog.tsx

@@ -176,7 +176,7 @@ const UserSettingsDialogContent = ({onClose, visibleFlags}: DialogContentProps)
             }}
           >
             <Box flex={{direction: 'row', gap: 4, alignItems: 'center'}}>
-              Reset IndexedDB cache
+              Reset IndexedDB cache (Close all other tabs first!)
               <Tooltip content="If you're seeing stale definitions or experiencing client side bugs then this may fix it">
                 <Icon name="info" />
               </Tooltip>

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetGraphExplorer.tsx

@@ -42,13 +42,9 @@ import {
 import {assetKeyTokensInRange} from './assetKeyTokensInRange';
 import {AssetGraphLayout, GroupLayout} from './layout';
 import {AssetGraphExplorerSidebar} from './sidebar/Sidebar';
+import {AssetGraphQueryItem} from './types';
 import {AssetNodeForGraphQueryFragment} from './types/useAssetGraphData.types';
-import {
-  AssetGraphFetchScope,
-  AssetGraphQueryItem,
-  useAssetGraphData,
-  useFullAssetGraphData,
-} from './useAssetGraphData';
+import {AssetGraphFetchScope, useAssetGraphData, useFullAssetGraphData} from './useAssetGraphData';
 import {AssetLocation, useFindAssetLocation} from './useFindAssetLocation';
 import {AssetLiveDataRefreshButton} from '../asset-data/AssetLiveDataProvider';
 import {LaunchAssetExecutionButton} from '../assets/LaunchAssetExecutionButton';

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-graph/ComputeGraphData.types.ts

@@ -1,6 +1,7 @@
+import {AssetGraphQueryItem} from './types';
 import {AssetKey} from '../assets/types';
 import {AssetNodeForGraphQueryFragment} from './types/useAssetGraphData.types';
-import {AssetGraphFetchScope, AssetGraphQueryItem} from './useAssetGraphData';
+import {AssetGraphFetchScope} from './useAssetGraphData';
 
 type BaseType = {
   id: number;

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-graph/layout.ts

@@ -1,7 +1,7 @@
 import * as dagre from 'dagre';
 
 import {GraphData, GraphId, GraphNode, groupIdForNode, isGroupId} from './Utils';
-import {IBounds, IPoint} from '../graph/common';
+import type {IBounds, IPoint} from '../graph/common';
 import {ChangeReason} from '../graphql/types';
 
 export type AssetLayoutDirection = 'vertical' | 'horizontal';

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx

@@ -10,6 +10,7 @@ import {GraphData, buildGraphData as buildGraphDataImpl, tokenForAssetKey} from
 import {gql} from '../apollo-client';
 import {computeGraphData as computeGraphDataImpl} from './ComputeGraphData';
 import {BuildGraphDataMessageType, ComputeGraphDataMessageType} from './ComputeGraphData.types';
+import {AssetGraphQueryItem, AssetNode} from './types';
 import {featureEnabled} from '../app/Flags';
 import {
   AssetGraphQuery,
@@ -47,10 +48,6 @@ export interface AssetGraphFetchScope {
   useWorker?: boolean;
 }
 
-export type AssetGraphQueryItem = GraphQueryItem & {
-  node: AssetNode;
-};
-
 export function useFullAssetGraphData(
   options: Omit<AssetGraphFetchScope, 'groupSelector' | 'pipelineSelector'>,
 ) {
@@ -194,7 +191,7 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop
   const currentRequestRef = useRef(0);
 
   const {loading: supplementaryDataLoading, data: supplementaryData} =
-    useAssetGraphSupplementaryData(opsQuery);
+    useAssetGraphSupplementaryData(opsQuery, allNodes);
 
   const spawnComputeGraphDataWorker = useMemo(
     () => workerSpawner(() => new Worker(new URL('./ComputeGraphData.worker', import.meta.url))),
@@ -274,8 +271,6 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop
   };
 }
 
-type AssetNode = AssetNodeForGraphQueryFragment;
-
 const computeGraphData = indexedDBAsyncMemoize<GraphDataState, typeof computeGraphDataWrapper>(
   computeGraphDataWrapper,
   (props) => {

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphSupplementaryData.oss.tsx

@@ -1,13 +1,46 @@
-import {AssetKey} from '../assets/types';
+import {useMemo} from 'react';
 
-export type SupplementaryInformation = Record<string, AssetKey[]> | null | undefined;
+import {useAssetsHealthData} from '../asset-data/AssetHealthDataProvider';
+import {parseExpression} from '../asset-selection/AssetSelectionSupplementaryDataVisitor';
+import {getSupplementaryDataKey} from '../asset-selection/util';
+import {AssetKey} from '../assets/types';
+import {AssetNodeForGraphQueryFragment} from './types/useAssetGraphData.types';
+import {SupplementaryInformation} from '../asset-selection/types';
 
-// Stub for cloud to supply extended selection syntax filtering capabilities
+const emptyObject = {} as SupplementaryInformation;
 export const useAssetGraphSupplementaryData = (
-  _: string,
-): {loading: boolean; data: SupplementaryInformation | null} => {
+  selection: string,
+  nodes: AssetNodeForGraphQueryFragment[],
+): {loading: boolean; data: SupplementaryInformation} => {
+  const {liveDataByNode} = useAssetsHealthData(
+    useMemo(() => nodes.map((node) => node.assetKey), [nodes]),
+  );
+
+  const loading = Object.keys(liveDataByNode).length !== nodes.length;
+
+  const assetsByStatus = useMemo(() => {
+    return Object.values(liveDataByNode).reduce(
+      (acc, liveData) => {
+        const status = liveData.assetHealth?.assetHealth ?? 'UNKNOWN';
+        const supplementaryDataKey = getSupplementaryDataKey({
+          field: 'status',
+          value: status,
+        });
+        acc[supplementaryDataKey] = acc[supplementaryDataKey] || [];
+        acc[supplementaryDataKey].push(liveData.assetKey);
+        return acc;
+      },
+      {} as Record<string, AssetKey[]>,
+    );
+  }, [liveDataByNode]);
+
+  const needsAssetHealthData = useMemo(() => {
+    const filters = parseExpression(selection);
+    return filters.some((filter) => filter.field === 'status');
+  }, [selection]);
+
   return {
-    loading: false,
-    data: null,
+    loading: needsAssetHealthData && loading,
+    data: loading ? emptyObject : assetsByStatus,
   };
 };

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelectionVisitor.oss.ts

@@ -1,10 +1,17 @@
 import {AbstractParseTreeVisitor} from 'antlr4ts/tree/AbstractParseTreeVisitor';
 import escapeRegExp from 'lodash/escapeRegExp';
-import {SupplementaryInformation} from 'shared/asset-graph/useAssetGraphSupplementaryData.oss';
 
-import {getFunctionName, getTraversalDepth, getValue} from './util';
+import {SupplementaryInformation} from './types';
+import {
+  getAssetsByKey,
+  getFunctionName,
+  getSupplementaryDataKey,
+  getTraversalDepth,
+  getValue,
+} from './util';
 import {GraphTraverser} from '../app/GraphQueryImpl';
-import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData';
+import {tokenForAssetKey} from '../asset-graph/Utils';
+import {AssetGraphQueryItem} from '../asset-graph/types';
 import {buildRepoPathForHuman} from '../workspace/buildRepoAddress';
 import {
   AllExpressionContext,
@@ -21,6 +28,7 @@ import {
   OwnerAttributeExprContext,
   ParenthesizedExpressionContext,
   StartContext,
+  StatusAttributeExprContext,
   TagAttributeExprContext,
   TraversalAllowedExpressionContext,
   UpAndDownTraversalExpressionContext,
@@ -32,20 +40,24 @@ export class AntlrAssetSelectionVisitor
   extends AbstractParseTreeVisitor<Set<AssetGraphQueryItem>>
   implements AssetSelectionVisitor<Set<AssetGraphQueryItem>>
 {
-  all_assets: Set<AssetGraphQueryItem>;
-  focus_assets: Set<AssetGraphQueryItem>;
-  traverser: GraphTraverser<AssetGraphQueryItem>;
+  protected all_assets: Set<AssetGraphQueryItem>;
+  public focus_assets: Set<AssetGraphQueryItem>;
+  protected traverser: GraphTraverser<AssetGraphQueryItem>;
+  protected supplementaryData: SupplementaryInformation;
+  protected allAssetsByKey: Map<string, AssetGraphQueryItem>;
 
   protected defaultResult() {
     return new Set<AssetGraphQueryItem>();
   }
 
   // Supplementary data is not used in oss
-  constructor(all_assets: AssetGraphQueryItem[], _supplementaryData?: SupplementaryInformation) {
+  constructor(all_assets: AssetGraphQueryItem[], supplementaryData: SupplementaryInformation) {
     super();
     this.all_assets = new Set(all_assets);
     this.focus_assets = new Set();
     this.traverser = new GraphTraverser(all_assets);
+    this.supplementaryData = supplementaryData;
+    this.allAssetsByKey = getAssetsByKey(all_assets);
   }
 
   visitStart(ctx: StartContext) {
@@ -208,4 +220,21 @@ export class AntlrAssetSelectionVisitor
     }
     return selection;
   }
+
+  visitStatusAttributeExpr(ctx: StatusAttributeExprContext) {
+    const statusName: string = getValue(ctx.value());
+    const supplementaryDataKey = getSupplementaryDataKey({
+      field: 'status',
+      value: statusName,
+    });
+    const matchingAssetKeys = this.supplementaryData?.[supplementaryDataKey];
+    if (!matchingAssetKeys) {
+      return new Set<AssetGraphQueryItem>();
+    }
+    return new Set(
+      matchingAssetKeys
+        .map((key) => this.allAssetsByKey.get(tokenForAssetKey(key))!)
+        .filter(Boolean),
+    );
+  }
 }

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts

@@ -1,5 +1,5 @@
 /* eslint-disable jest/expect-expect */
-import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
+import {AssetGraphQueryItem} from '../../asset-graph/types';
 import {
   buildAssetNode,
   buildDefinitionTag,

Diff for: js_modules/dagster-ui/packages/ui-core/src/asset-selection/filterAssetSelectionByQuery.oss.ts

@@ -1,10 +1,10 @@
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
-import {SupplementaryInformation} from 'shared/asset-graph/useAssetGraphSupplementaryData.oss';
 
 import {AssetSelectionQueryResult, parseAssetSelectionQuery} from './parseAssetSelectionQuery';
+import {SupplementaryInformation} from './types';
 import {featureEnabled} from '../app/Flags';
 import {filterByQuery} from '../app/GraphQueryImpl';
-import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData';
+import {AssetGraphQueryItem} from '../asset-graph/types';
 import {weakMapMemoize} from '../util/weakMapMemoize';
 
 export const filterAssetSelectionByQuery = weakMapMemoize(