diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts index 17722bca32662..d39434b2d0201 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts +++ b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts @@ -14,7 +14,7 @@ import {AssetSelectionParser} from './generated/AssetSelectionParser'; import {featureEnabled} from '../app/Flags'; import {filterByQuery} from '../app/GraphQueryImpl'; -class AntlrInputErrorListener implements ANTLRErrorListener { +export class AntlrInputErrorListener implements ANTLRErrorListener { syntaxError( recognizer: Recognizer, offendingSymbol: any, diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelectionVisitor.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelectionVisitor.ts index 1ce8cd6869cce..de82e2f95c28f 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelectionVisitor.ts +++ b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelectionVisitor.ts @@ -29,7 +29,7 @@ import {GraphTraverser} from '../app/GraphQueryImpl'; import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData'; import {buildRepoPathForHuman} from '../workspace/buildRepoAddress'; -function getTraversalDepth(ctx: TraversalContext): number { +export function getTraversalDepth(ctx: TraversalContext): number { if (ctx.STAR()) { return Number.MAX_SAFE_INTEGER; } @@ -39,7 +39,7 @@ function getTraversalDepth(ctx: TraversalContext): number { throw new Error('Invalid traversal'); } -function getFunctionName(ctx: FunctionNameContext): string { +export function getFunctionName(ctx: FunctionNameContext): string { if (ctx.SINKS()) { return 'sinks'; } @@ -49,7 +49,7 @@ function getFunctionName(ctx: FunctionNameContext): string { throw new Error('Invalid function name'); } -function getValue(ctx: ValueContext): string { +export function getValue(ctx: ValueContext): string { if (ctx.QUOTED_STRING()) { return ctx.text.slice(1, -1); } diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts index f27c4238980f0..bac998097a71f 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts +++ b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts @@ -87,7 +87,8 @@ describe('parseAssetSelectionQuery', () => { }); it('should parse key_substring query', () => { - assertQueryResult('key:A', ['A']); + assertQueryResult('key_substring:A', ['A']); + assertQueryResult('key_substring:B', ['B', 'B2']); }); it('should parse and query', () => { diff --git a/js_modules/dagster-ui/packages/ui-core/src/gantt/__tests__/toGraphQueryItems.test.ts b/js_modules/dagster-ui/packages/ui-core/src/gantt/__tests__/toGraphQueryItems.test.ts index 519bc5e6afa02..5969d4e7df38d 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/gantt/__tests__/toGraphQueryItems.test.ts +++ b/js_modules/dagster-ui/packages/ui-core/src/gantt/__tests__/toGraphQueryItems.test.ts @@ -48,9 +48,30 @@ describe('toGraphQueryItems', () => { inputs: [], name: 'a', outputs: [{dependedBy: [{solid: {name: 'b[1]'}}, {solid: {name: 'b[?]'}}]}], + metadata: { + attempts: [], + markers: [], + state: 'succeeded', + transitions: [], + }, + }, + { + inputs: [{dependsOn: [{solid: {name: 'a'}}]}], + name: 'b[1]', + outputs: [], + metadata: { + attempts: [], + markers: [], + state: 'succeeded', + transitions: [], + }, + }, + { + inputs: [{dependsOn: [{solid: {name: 'a'}}]}], + name: 'b[?]', + outputs: [], + metadata: undefined, }, - {inputs: [{dependsOn: [{solid: {name: 'a'}}]}], name: 'b[1]', outputs: []}, - {inputs: [{dependsOn: [{solid: {name: 'a'}}]}], name: 'b[?]', outputs: []}, ]); }); @@ -129,6 +150,12 @@ describe('toGraphQueryItems', () => { ], }, ], + metadata: { + attempts: [], + markers: [], + state: 'succeeded', + transitions: [], + }, }, { inputs: [ @@ -154,6 +181,12 @@ describe('toGraphQueryItems', () => { ], }, ], + metadata: { + attempts: [], + markers: [], + state: 'succeeded', + transitions: [], + }, }, { inputs: [ @@ -169,6 +202,12 @@ describe('toGraphQueryItems', () => { ], name: 'b[2]', outputs: [], + metadata: { + attempts: [], + markers: [], + state: 'succeeded', + transitions: [], + }, }, { inputs: [ @@ -194,6 +233,7 @@ describe('toGraphQueryItems', () => { ], }, ], + metadata: undefined, }, { inputs: [ @@ -209,6 +249,12 @@ describe('toGraphQueryItems', () => { ], name: 'c[1]', outputs: [], + metadata: { + attempts: [], + markers: [], + state: 'succeeded', + transitions: [], + }, }, { inputs: [ @@ -224,6 +270,7 @@ describe('toGraphQueryItems', () => { ], name: 'c[?]', outputs: [], + metadata: undefined, }, ]); }); diff --git a/js_modules/dagster-ui/packages/ui-core/src/gantt/toGraphQueryItems.tsx b/js_modules/dagster-ui/packages/ui-core/src/gantt/toGraphQueryItems.tsx index 031d5b595f619..7e98cf899f2bc 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/gantt/toGraphQueryItems.tsx +++ b/js_modules/dagster-ui/packages/ui-core/src/gantt/toGraphQueryItems.tsx @@ -5,6 +5,10 @@ import {GraphQueryItem} from '../app/GraphQueryImpl'; import {StepKind} from '../graphql/types'; import {IStepMetadata, IStepState} from '../runs/RunMetadataProvider'; +export type RunGraphQueryItem = GraphQueryItem & { + metadata: IStepMetadata | undefined; +}; + /** * Converts a Run execution plan into a tree of `GraphQueryItem` items that * can be used as the input to the "solid query" filtering algorithm or rendered @@ -18,7 +22,7 @@ import {IStepMetadata, IStepState} from '../runs/RunMetadataProvider'; export const toGraphQueryItems = ( plan: ExecutionPlanToGraphFragment, runtimeStepMetadata: {[key: string]: IStepMetadata}, -) => { +): RunGraphQueryItem[] => { // Step 1: Find unresolved steps in the initial plan and build a mapping // of their unresolved names to their resolved step keys, eg: // "multiply_input[*]" => ["multiply_input[1]", "multiply_input[2]"] @@ -47,7 +51,7 @@ export const toGraphQueryItems = ( } // Step 2: Create a graph node for each resolved step without any inputs or outputs. - const nodeTable: {[key: string]: GraphQueryItem} = {}; + const nodeTable: {[key: string]: RunGraphQueryItem} = {}; for (const step of plan.steps) { const stepRuntimeKeys = keyExpansionMap[step.key] || [step.key]; for (const key of stepRuntimeKeys) { @@ -55,6 +59,7 @@ export const toGraphQueryItems = ( name: key, inputs: [], outputs: [], + metadata: runtimeStepMetadata[key], }; } } diff --git a/js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelection.ts b/js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelection.ts new file mode 100644 index 0000000000000..400a4d5172aff --- /dev/null +++ b/js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelection.ts @@ -0,0 +1,42 @@ +import {CharStreams, CommonTokenStream} from 'antlr4ts'; + +import {AntlrRunSelectionVisitor} from './AntlrRunSelectionVisitor'; +import {AntlrInputErrorListener} from '../asset-selection/AntlrAssetSelection'; +import {RunGraphQueryItem} from '../gantt/toGraphQueryItems'; +import {RunSelectionLexer} from './generated/RunSelectionLexer'; +import {RunSelectionParser} from './generated/RunSelectionParser'; + +type RunSelectionQueryResult = { + all: RunGraphQueryItem[]; + focus: RunGraphQueryItem[]; +}; + +export const parseRunSelectionQuery = ( + all_runs: RunGraphQueryItem[], + query: string, +): RunSelectionQueryResult | Error => { + try { + const lexer = new RunSelectionLexer(CharStreams.fromString(query)); + lexer.removeErrorListeners(); + lexer.addErrorListener(new AntlrInputErrorListener()); + + const tokenStream = new CommonTokenStream(lexer); + + const parser = new RunSelectionParser(tokenStream); + parser.removeErrorListeners(); + parser.addErrorListener(new AntlrInputErrorListener()); + + const tree = parser.start(); + + const visitor = new AntlrRunSelectionVisitor(all_runs); + const all_selection = visitor.visit(tree); + const focus_selection = visitor.focus_runs; + + return { + all: Array.from(all_selection), + focus: Array.from(focus_selection), + }; + } catch (e) { + return e as Error; + } +}; diff --git a/js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelectionVisitor.ts b/js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelectionVisitor.ts new file mode 100644 index 0000000000000..c1f457693062a --- /dev/null +++ b/js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelectionVisitor.ts @@ -0,0 +1,165 @@ +import {AbstractParseTreeVisitor} from 'antlr4ts/tree/AbstractParseTreeVisitor'; + +import {GraphTraverser} from '../app/GraphQueryImpl'; +import {RunGraphQueryItem} from '../gantt/toGraphQueryItems'; +import { + AllExpressionContext, + AndExpressionContext, + AttributeExpressionContext, + DownTraversalExpressionContext, + FunctionCallExpressionContext, + NameExprContext, + NameSubstringExprContext, + NotExpressionContext, + OrExpressionContext, + ParenthesizedExpressionContext, + StartContext, + StatusAttributeExprContext, + TraversalAllowedExpressionContext, + UpAndDownTraversalExpressionContext, + UpTraversalExpressionContext, +} from './generated/RunSelectionParser'; +import {RunSelectionVisitor} from './generated/RunSelectionVisitor'; +import { + getFunctionName, + getTraversalDepth, + getValue, +} from '../asset-selection/AntlrAssetSelectionVisitor'; + +export class AntlrRunSelectionVisitor + extends AbstractParseTreeVisitor> + implements RunSelectionVisitor> +{ + all_runs: Set; + focus_runs: Set; + traverser: GraphTraverser; + + protected defaultResult() { + return new Set(); + } + + constructor(all_runs: RunGraphQueryItem[]) { + super(); + this.all_runs = new Set(all_runs); + this.focus_runs = new Set(); + this.traverser = new GraphTraverser(all_runs); + } + + visitStart(ctx: StartContext) { + return this.visit(ctx.expr()); + } + + visitTraversalAllowedExpression(ctx: TraversalAllowedExpressionContext) { + return this.visit(ctx.traversalAllowedExpr()); + } + + visitUpAndDownTraversalExpression(ctx: UpAndDownTraversalExpressionContext) { + const selection = this.visit(ctx.traversalAllowedExpr()); + const up_depth: number = getTraversalDepth(ctx.traversal(0)); + const down_depth: number = getTraversalDepth(ctx.traversal(1)); + const selection_copy = new Set(selection); + for (const item of selection_copy) { + this.traverser.fetchUpstream(item, up_depth).forEach((i) => selection.add(i)); + this.traverser.fetchDownstream(item, down_depth).forEach((i) => selection.add(i)); + } + return selection; + } + + visitUpTraversalExpression(ctx: UpTraversalExpressionContext) { + const selection = this.visit(ctx.traversalAllowedExpr()); + const traversal_depth: number = getTraversalDepth(ctx.traversal()); + const selection_copy = new Set(selection); + for (const item of selection_copy) { + this.traverser.fetchUpstream(item, traversal_depth).forEach((i) => selection.add(i)); + } + return selection; + } + + visitDownTraversalExpression(ctx: DownTraversalExpressionContext) { + const selection = this.visit(ctx.traversalAllowedExpr()); + const traversal_depth: number = getTraversalDepth(ctx.traversal()); + const selection_copy = new Set(selection); + for (const item of selection_copy) { + this.traverser.fetchDownstream(item, traversal_depth).forEach((i) => selection.add(i)); + } + return selection; + } + + visitNotExpression(ctx: NotExpressionContext) { + const selection = this.visit(ctx.expr()); + return new Set([...this.all_runs].filter((i) => !selection.has(i))); + } + + visitAndExpression(ctx: AndExpressionContext) { + const left = this.visit(ctx.expr(0)); + const right = this.visit(ctx.expr(1)); + return new Set([...left].filter((i) => right.has(i))); + } + + visitOrExpression(ctx: OrExpressionContext) { + const left = this.visit(ctx.expr(0)); + const right = this.visit(ctx.expr(1)); + return new Set([...left, ...right]); + } + + visitAllExpression(_ctx: AllExpressionContext) { + return this.all_runs; + } + + visitAttributeExpression(ctx: AttributeExpressionContext) { + return this.visit(ctx.attributeExpr()); + } + + visitFunctionCallExpression(ctx: FunctionCallExpressionContext) { + const function_name: string = getFunctionName(ctx.functionName()); + const selection = this.visit(ctx.expr()); + if (function_name === 'sinks') { + const sinks = new Set(); + for (const item of selection) { + const downstream = this.traverser + .fetchDownstream(item, Number.MAX_VALUE) + .filter((i) => selection.has(i)); + if (downstream.length === 0 || (downstream.length === 1 && downstream[0] === item)) { + sinks.add(item); + } + } + return sinks; + } + if (function_name === 'roots') { + const roots = new Set(); + for (const item of selection) { + const upstream = this.traverser + .fetchUpstream(item, Number.MAX_VALUE) + .filter((i) => selection.has(i)); + if (upstream.length === 0 || (upstream.length === 1 && upstream[0] === item)) { + roots.add(item); + } + } + return roots; + } + throw new Error(`Unknown function: ${function_name}`); + } + + visitParenthesizedExpression(ctx: ParenthesizedExpressionContext) { + return this.visit(ctx.expr()); + } + + visitNameExpr(ctx: NameExprContext) { + const value: string = getValue(ctx.value()); + const selection = [...this.all_runs].filter((i) => i.name === value); + selection.forEach((i) => this.focus_runs.add(i)); + return new Set(selection); + } + + visitNameSubstringExpr(ctx: NameSubstringExprContext) { + const value: string = getValue(ctx.value()); + const selection = [...this.all_runs].filter((i) => i.name.includes(value)); + selection.forEach((i) => this.focus_runs.add(i)); + return new Set(selection); + } + + visitStatusAttributeExpr(ctx: StatusAttributeExprContext) { + const state: string = getValue(ctx.value()).toLowerCase(); + return new Set([...this.all_runs].filter((i) => i.metadata?.state === state)); + } +} diff --git a/js_modules/dagster-ui/packages/ui-core/src/run-selection/__tests__/AntlrRunSelection.test.ts b/js_modules/dagster-ui/packages/ui-core/src/run-selection/__tests__/AntlrRunSelection.test.ts new file mode 100644 index 0000000000000..8f573934ae8ae --- /dev/null +++ b/js_modules/dagster-ui/packages/ui-core/src/run-selection/__tests__/AntlrRunSelection.test.ts @@ -0,0 +1,152 @@ +/* eslint-disable jest/expect-expect */ + +import {RunGraphQueryItem} from '../../gantt/toGraphQueryItems'; +import {IStepMetadata, IStepState} from '../../runs/RunMetadataProvider'; +import {parseRunSelectionQuery} from '../AntlrRunSelection'; + +function buildMetadata(state: IStepState): IStepMetadata { + return { + state, + attempts: [], + markers: [], + transitions: [], + }; +} + +const TEST_GRAPH: RunGraphQueryItem[] = [ + // Top Layer + { + name: 'A', + metadata: buildMetadata(IStepState.SUCCEEDED), + inputs: [{dependsOn: []}], + outputs: [{dependedBy: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}], + }, + // Second Layer + { + name: 'B', + metadata: buildMetadata(IStepState.FAILED), + inputs: [{dependsOn: [{solid: {name: 'A'}}]}], + outputs: [{dependedBy: [{solid: {name: 'C'}}]}], + }, + { + name: 'B2', + metadata: buildMetadata(IStepState.SKIPPED), + inputs: [{dependsOn: [{solid: {name: 'A'}}]}], + outputs: [{dependedBy: [{solid: {name: 'C'}}]}], + }, + // Third Layer + { + name: 'C', + metadata: buildMetadata(IStepState.RUNNING), + inputs: [{dependsOn: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}], + outputs: [{dependedBy: []}], + }, +]; + +function assertQueryResult(query: string, expectedNames: string[]) { + const result = parseRunSelectionQuery(TEST_GRAPH, query); + expect(result).not.toBeInstanceOf(Error); + if (result instanceof Error) { + throw result; + } + expect(result.all.length).toBe(expectedNames.length); + expect(new Set(result.all.map((run) => run.name))).toEqual(new Set(expectedNames)); +} + +// Most tests copied from AntlrAssetSelection.test.ts +describe('parseRunSelectionQuery', () => { + describe('invalid queries', () => { + it('should throw on invalid queries', () => { + expect(parseRunSelectionQuery(TEST_GRAPH, 'A')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'name:A name:B')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'not')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'and')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'name:A and')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'sinks')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'notafunction()')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'tag:foo=')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'owner')).toBeInstanceOf(Error); + expect(parseRunSelectionQuery(TEST_GRAPH, 'owner:owner@owner.com')).toBeInstanceOf(Error); + }); + }); + + describe('valid queries', () => { + it('should parse star query', () => { + assertQueryResult('*', ['A', 'B', 'B2', 'C']); + }); + + it('should parse name query', () => { + assertQueryResult('name:A', ['A']); + }); + + it('should parse name_substring query', () => { + assertQueryResult('name_substring:A', ['A']); + assertQueryResult('name_substring:B', ['B', 'B2']); + }); + + it('should parse and query', () => { + assertQueryResult('name:A and name:B', []); + assertQueryResult('name:A and name:B and name:C', []); + }); + + it('should parse or query', () => { + assertQueryResult('name:A or name:B', ['A', 'B']); + assertQueryResult('name:A or name:B or name:C', ['A', 'B', 'C']); + assertQueryResult('(name:A or name:B) and (name:B or name:C)', ['B']); + }); + + it('should parse upstream plus query', () => { + assertQueryResult('+name:A', ['A']); + assertQueryResult('+name:B', ['A', 'B']); + assertQueryResult('+name:C', ['B', 'B2', 'C']); + assertQueryResult('++name:C', ['A', 'B', 'B2', 'C']); + }); + + it('should parse downstream plus query', () => { + assertQueryResult('name:A+', ['A', 'B', 'B2']); + assertQueryResult('name:A++', ['A', 'B', 'B2', 'C']); + assertQueryResult('name:C+', ['C']); + assertQueryResult('name:B+', ['B', 'C']); + }); + + it('should parse upstream star query', () => { + assertQueryResult('*name:A', ['A']); + assertQueryResult('*name:B', ['A', 'B']); + assertQueryResult('*name:C', ['A', 'B', 'B2', 'C']); + }); + + it('should parse downstream star query', () => { + assertQueryResult('name:A*', ['A', 'B', 'B2', 'C']); + assertQueryResult('name:B*', ['B', 'C']); + assertQueryResult('name:C*', ['C']); + }); + + it('should parse up and down traversal queries', () => { + assertQueryResult('name:A* and *name:C', ['A', 'B', 'B2', 'C']); + assertQueryResult('*name:B*', ['A', 'B', 'C']); + assertQueryResult('name:A* and *name:C and *name:B*', ['A', 'B', 'C']); + assertQueryResult('name:A* and *name:B* and *name:C', ['A', 'B', 'C']); + }); + + it('should parse sinks query', () => { + assertQueryResult('sinks(*)', ['C']); + assertQueryResult('sinks(name:A)', ['A']); + assertQueryResult('sinks(name:A or name:B)', ['B']); + }); + + it('should parse roots query', () => { + assertQueryResult('roots(*)', ['A']); + assertQueryResult('roots(name:C)', ['C']); + assertQueryResult('roots(name:A or name:B)', ['A']); + }); + + it('should parse status query', () => { + assertQueryResult('status:succeeded', ['A']); + assertQueryResult('status:SUCCEEDED', ['A']); + assertQueryResult('status:failed', ['B']); + assertQueryResult('status:skipped', ['B2']); + assertQueryResult('status:"running"', ['C']); + assertQueryResult('status:not_a_status', []); + }); + }); +});