diff --git a/js_modules/dagster-ui/packages/ui-core/src/op-selection/AntlrOpSelection.ts b/js_modules/dagster-ui/packages/ui-core/src/op-selection/AntlrOpSelection.ts new file mode 100644 index 0000000000000..614bde690e347 --- /dev/null +++ b/js_modules/dagster-ui/packages/ui-core/src/op-selection/AntlrOpSelection.ts @@ -0,0 +1,42 @@ +import {CharStreams, CommonTokenStream} from 'antlr4ts'; + +import {AntlrOpSelectionVisitor} from './AntlrOpSelectionVisitor'; +import {GraphQueryItem} from '../app/GraphQueryImpl'; +import {AntlrInputErrorListener} from '../asset-selection/AntlrAssetSelection'; +import {OpSelectionLexer} from './generated/OpSelectionLexer'; +import {OpSelectionParser} from './generated/OpSelectionParser'; + +type OpSelectionQueryResult = { + all: GraphQueryItem[]; + focus: GraphQueryItem[]; +}; + +export const parseOpSelectionQuery = ( + all_ops: GraphQueryItem[], + query: string, +): OpSelectionQueryResult | Error => { + try { + const lexer = new OpSelectionLexer(CharStreams.fromString(query)); + lexer.removeErrorListeners(); + lexer.addErrorListener(new AntlrInputErrorListener()); + + const tokenStream = new CommonTokenStream(lexer); + + const parser = new OpSelectionParser(tokenStream); + parser.removeErrorListeners(); + parser.addErrorListener(new AntlrInputErrorListener()); + + const tree = parser.start(); + + const visitor = new AntlrOpSelectionVisitor(all_ops); + const all_selection = visitor.visit(tree); + const focus_selection = visitor.focus_ops; + + 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/op-selection/AntlrOpSelectionVisitor.ts b/js_modules/dagster-ui/packages/ui-core/src/op-selection/AntlrOpSelectionVisitor.ts new file mode 100644 index 0000000000000..7af548928dfe6 --- /dev/null +++ b/js_modules/dagster-ui/packages/ui-core/src/op-selection/AntlrOpSelectionVisitor.ts @@ -0,0 +1,123 @@ +import {AbstractParseTreeVisitor} from 'antlr4ts/tree/AbstractParseTreeVisitor'; + +import {GraphQueryItem, GraphTraverser} from '../app/GraphQueryImpl'; +import { + AllExpressionContext, + AndExpressionContext, + AttributeExpressionContext, + DownTraversalExpressionContext, + NameExprContext, + NameSubstringExprContext, + NotExpressionContext, + OrExpressionContext, + ParenthesizedExpressionContext, + StartContext, + TraversalAllowedExpressionContext, + UpAndDownTraversalExpressionContext, + UpTraversalExpressionContext, +} from './generated/OpSelectionParser'; +import {OpSelectionVisitor} from './generated/OpSelectionVisitor'; +import {getTraversalDepth, getValue} from '../asset-selection/AntlrAssetSelectionVisitor'; + +export class AntlrOpSelectionVisitor + extends AbstractParseTreeVisitor> + implements OpSelectionVisitor> +{ + all_ops: Set; + focus_ops: Set; + traverser: GraphTraverser; + + protected defaultResult() { + return new Set(); + } + + constructor(all_ops: GraphQueryItem[]) { + super(); + this.all_ops = new Set(all_ops); + this.focus_ops = new Set(); + this.traverser = new GraphTraverser(all_ops); + } + + 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_ops].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_ops; + } + + visitAttributeExpression(ctx: AttributeExpressionContext) { + return this.visit(ctx.attributeExpr()); + } + + visitParenthesizedExpression(ctx: ParenthesizedExpressionContext) { + return this.visit(ctx.expr()); + } + + visitNameExpr(ctx: NameExprContext) { + const value: string = getValue(ctx.value()); + const selection = [...this.all_ops].filter((i) => i.name === value); + selection.forEach((i) => this.focus_ops.add(i)); + return new Set(selection); + } + + visitNameSubstringExpr(ctx: NameSubstringExprContext) { + const value: string = getValue(ctx.value()); + const selection = [...this.all_ops].filter((i) => i.name.includes(value)); + selection.forEach((i) => this.focus_ops.add(i)); + return new Set(selection); + } +} diff --git a/js_modules/dagster-ui/packages/ui-core/src/op-selection/__tests__/AntlrOpSelection.test.ts b/js_modules/dagster-ui/packages/ui-core/src/op-selection/__tests__/AntlrOpSelection.test.ts new file mode 100644 index 0000000000000..8150daefa203d --- /dev/null +++ b/js_modules/dagster-ui/packages/ui-core/src/op-selection/__tests__/AntlrOpSelection.test.ts @@ -0,0 +1,118 @@ +/* eslint-disable jest/expect-expect */ + +import {GraphQueryItem} from '../../app/GraphQueryImpl'; +import {parseOpSelectionQuery} from '../AntlrOpSelection'; + +const TEST_GRAPH: GraphQueryItem[] = [ + // Top Layer + { + name: 'A', + inputs: [{dependsOn: []}], + outputs: [{dependedBy: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}], + }, + // Second Layer + { + name: 'B', + inputs: [{dependsOn: [{solid: {name: 'A'}}]}], + outputs: [{dependedBy: [{solid: {name: 'C'}}]}], + }, + { + name: 'B2', + inputs: [{dependsOn: [{solid: {name: 'A'}}]}], + outputs: [{dependedBy: [{solid: {name: 'C'}}]}], + }, + // Third Layer + { + name: 'C', + inputs: [{dependsOn: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}], + outputs: [{dependedBy: []}], + }, +]; + +function assertQueryResult(query: string, expectedNames: string[]) { + const result = parseOpSelectionQuery(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((op) => op.name))).toEqual(new Set(expectedNames)); +} + +// Most tests copied from AntlrAssetSelection.test.ts +describe('parseOpSelectionQuery', () => { + describe('invalid queries', () => { + it('should throw on invalid queries', () => { + expect(parseOpSelectionQuery(TEST_GRAPH, 'A')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'name:A name:B')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'not')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'and')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'name:A and')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'sinks(*)')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'roots(*)')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'notafunction()')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'tag:foo=')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(TEST_GRAPH, 'owner')).toBeInstanceOf(Error); + expect(parseOpSelectionQuery(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']); + }); + }); +});