Skip to content

[4/n] Support status in selection syntax #29289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Apr 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion js_modules/dagster-ui/packages/ui-core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -600,7 +600,8 @@ describe('AssetLiveDataProvider', () => {

expect(resultFn).toHaveBeenCalled();
await waitFor(() => {
expect(hookResult.mock.calls[1][0]!).toEqual({
expect(hookResult.mock.calls[1][0]!).toEqual({});
expect(hookResult.mock.calls[2][0]!).toEqual({
['key1']: expect.any(Object),
['key2']: expect.any(Object),
});
Expand All @@ -615,8 +616,9 @@ describe('AssetLiveDataProvider', () => {
act(() => {
jest.runOnlyPendingTimers();
});
expect(hookResult.mock.calls[1][0]!).toEqual({});

expect(hookResult2.mock.calls[1][0]).toEqual({
expect(hookResult2.mock.calls[2][0]).toEqual({
['key1']: expect.any(Object),
});
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand Down
Original file line number Diff line number Diff line change
@@ -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';
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
import {GraphQueryItem} from '../app/GraphQueryImpl';
import {AssetNodeForGraphQueryFragment} from './types/useAssetGraphData.types';

export type AssetNode = AssetNodeForGraphQueryFragment;
export type AssetGraphQueryItem = GraphQueryItem & {
node: AssetNode;
};
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -47,10 +48,6 @@ export interface AssetGraphFetchScope {
useWorker?: boolean;
}

export type AssetGraphQueryItem = GraphQueryItem & {
node: AssetNode;
};

export function useFullAssetGraphData(
options: Omit<AssetGraphFetchScope, 'groupSelector' | 'pipelineSelector'>,
) {
Expand Down Expand Up @@ -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))),
Expand Down Expand Up @@ -274,8 +271,6 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop
};
}

type AssetNode = AssetNodeForGraphQueryFragment;

const computeGraphData = indexedDBAsyncMemoize<GraphDataState, typeof computeGraphDataWrapper>(
computeGraphDataWrapper,
(props) => {
Expand Down
Original file line number Diff line number Diff line change
@@ -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,
};
};
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -21,6 +28,7 @@ import {
OwnerAttributeExprContext,
ParenthesizedExpressionContext,
StartContext,
StatusAttributeExprContext,
TagAttributeExprContext,
TraversalAllowedExpressionContext,
UpAndDownTraversalExpressionContext,
Expand All @@ -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) {
Expand Down Expand Up @@ -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),
);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
import {CharStreams, CommonTokenStream} from 'antlr4ts';
import {AbstractParseTreeVisitor} from 'antlr4ts/tree/AbstractParseTreeVisitor';

import {AssetSelectionLexer} from './generated/AssetSelectionLexer';
import {AssetSelectionParser, StatusAttributeExprContext} from './generated/AssetSelectionParser';
import {AssetSelectionVisitor} from './generated/AssetSelectionVisitor';
import {getValue} from './util';

export type Filter = {field: 'status'; value: string};

export class AssetSelectionSupplementaryDataVisitor
extends AbstractParseTreeVisitor<void>
implements AssetSelectionVisitor<void>
{
public filters: Filter[] = [];

constructor() {
super();
}

protected defaultResult() {}

visitStatusAttributeExpr(ctx: StatusAttributeExprContext) {
const value: string = getValue(ctx.value());
this.filters.push({field: 'status', value});
}
}

export const parseExpression = (expression: string) => {
const inputStream = CharStreams.fromString(expression);
const lexer = new AssetSelectionLexer(inputStream);
const tokenStream = new CommonTokenStream(lexer);
const parser = new AssetSelectionParser(tokenStream);
const tree = parser.start();
const visitor = new AssetSelectionSupplementaryDataVisitor();
tree.accept(visitor);
return visitor.filters;
};
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/* eslint-disable jest/expect-expect */
import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
import {AssetGraphQueryItem} from '../../asset-graph/types';
import {
buildAssetNode,
buildDefinitionTag,
Expand Down
Original file line number Diff line number Diff line change
@@ -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(
Expand Down

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading