Implement AntlrOpSelectionVisitor on frontend

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;
+  }
+};

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<Set<GraphQueryItem>>
+  implements OpSelectionVisitor<Set<GraphQueryItem>>
+{
+  all_ops: Set<GraphQueryItem>;
+  focus_ops: Set<GraphQueryItem>;
+  traverser: GraphTraverser<GraphQueryItem>;
+
+  protected defaultResult() {
+    return new Set<GraphQueryItem>();
+  }
+
+  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);
+  }
+}

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:[email protected]')).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']);
+    });
+  });
+});