ANTLR only allow traversal on certain expressions (#26068)

## Summary & Motivation
We discovered some ambiguity caused by traversal expressions that may cause parser to return an unexpected asset selection. For example, currently the parser reads `*key:a and key:b*` as `*(key:a and key:b)*`. We believe this is unintuitive, so we've decided to force the traversal tokens to only apply to the expression they are immediately next to.

js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts

@@ -21,7 +21,7 @@ class AntlrInputErrorListener implements ANTLRErrorListener<any> {
     line: number,
     charPositionInLine: number,
     msg: string,
-    e: RecognitionException | undefined,
+    _e: RecognitionException | undefined,
   ): void {
     if (offendingSymbol) {
       throw new Error(`Syntax error caused by "${offendingSymbol.text}": ${msg}`);

js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelectionVisitor.ts

@@ -18,6 +18,7 @@ import {
   ParenthesizedExpressionContext,
   StartContext,
   TagAttributeExprContext,
+  TraversalAllowedExpressionContext,
   TraversalContext,
   UpAndDownTraversalExpressionContext,
   UpTraversalExpressionContext,
@@ -77,34 +78,41 @@ export class AntlrAssetSelectionVisitor
     this.traverser = new GraphTraverser(all_assets);
   }
 
-  visitAttributeExpression(ctx: AttributeExpressionContext) {
-    return this.visit(ctx.attributeExpr());
+  visitStart(ctx: StartContext) {
+    return this.visit(ctx.expr());
   }
 
-  visitUpTraversalExpression(ctx: UpTraversalExpressionContext) {
-    const selection = this.visit(ctx.expr());
-    const traversal_depth: number = getTraversalDepth(ctx.traversal());
-    for (const item of selection) {
-      this.traverser.fetchUpstream(item, traversal_depth).forEach((i) => selection.add(i));
-    }
-    return selection;
+  visitTraversalAllowedExpression(ctx: TraversalAllowedExpressionContext) {
+    return this.visit(ctx.traversalAllowedExpr());
   }
 
   visitUpAndDownTraversalExpression(ctx: UpAndDownTraversalExpressionContext) {
-    const selection = this.visit(ctx.expr());
+    const selection = this.visit(ctx.traversalAllowedExpr());
     const up_depth: number = getTraversalDepth(ctx.traversal(0));
     const down_depth: number = getTraversalDepth(ctx.traversal(1));
-    for (const item of selection) {
+    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.expr());
+    const selection = this.visit(ctx.traversalAllowedExpr());
     const traversal_depth: number = getTraversalDepth(ctx.traversal());
-    for (const item of selection) {
+    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;
@@ -127,6 +135,14 @@ export class AntlrAssetSelectionVisitor
     return new Set([...left, ...right]);
   }
 
+  visitAllExpression(_ctx: AllExpressionContext) {
+    return this.all_assets;
+  }
+
+  visitAttributeExpression(ctx: AttributeExpressionContext) {
+    return this.visit(ctx.attributeExpr());
+  }
+
   visitFunctionCallExpression(ctx: FunctionCallExpressionContext) {
     const function_name: string = getFunctionName(ctx.functionName());
     const selection = this.visit(ctx.expr());
@@ -161,10 +177,6 @@ export class AntlrAssetSelectionVisitor
     return this.visit(ctx.expr());
   }
 
-  visitAllExpression(_ctx: AllExpressionContext) {
-    return this.all_assets;
-  }
-
   visitKeyExpr(ctx: KeyExprContext) {
     const value: string = getValue(ctx.value());
     const selection = [...this.all_assets].filter((i) => i.name === value);
@@ -231,8 +243,4 @@ export class AntlrAssetSelectionVisitor
     }
     return selection;
   }
-
-  visitStart(ctx: StartContext) {
-    return this.visit(ctx.expr());
-  }
 }

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

@@ -18,7 +18,7 @@ const TEST_GRAPH: AssetGraphQueryItem[] = [
       owners: [buildUserAssetOwner({email: '[email protected]'})],
     }),
     inputs: [{dependsOn: []}],
-    outputs: [{dependedBy: [{solid: {name: 'B'}}]}],
+    outputs: [{dependedBy: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}],
   },
   // Second Layer
   {
@@ -30,7 +30,12 @@ const TEST_GRAPH: AssetGraphQueryItem[] = [
     inputs: [{dependsOn: [{solid: {name: 'A'}}]}],
     outputs: [{dependedBy: [{solid: {name: 'C'}}]}],
   },
-
+  {
+    name: 'B2',
+    node: buildAssetNode(),
+    inputs: [{dependsOn: [{solid: {name: 'A'}}]}],
+    outputs: [{dependedBy: [{solid: {name: 'C'}}]}],
+  },
   // Third Layer
   {
     name: 'C',
@@ -41,7 +46,7 @@ const TEST_GRAPH: AssetGraphQueryItem[] = [
         location: buildRepositoryLocation({name: 'my_location'}),
       }),
     }),
-    inputs: [{dependsOn: [{solid: {name: 'B'}}]}],
+    inputs: [{dependsOn: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}],
     outputs: [{dependedBy: []}],
   },
 ];
@@ -74,7 +79,7 @@ describe('parseAssetSelectionQuery', () => {
 
   describe('valid queries', () => {
     it('should parse star query', () => {
-      assertQueryResult('*', ['A', 'B', 'C']);
+      assertQueryResult('*', ['A', 'B', 'B2', 'C']);
     });
 
     it('should parse key query', () => {
@@ -99,25 +104,34 @@ describe('parseAssetSelectionQuery', () => {
     it('should parse upstream plus query', () => {
       assertQueryResult('+key:A', ['A']);
       assertQueryResult('+key:B', ['A', 'B']);
-      assertQueryResult('++key:C', ['A', 'B', 'C']);
+      assertQueryResult('+key:C', ['B', 'B2', 'C']);
+      assertQueryResult('++key:C', ['A', 'B', 'B2', 'C']);
     });
 
     it('should parse downstream plus query', () => {
+      assertQueryResult('key:A+', ['A', 'B', 'B2']);
+      assertQueryResult('key:A++', ['A', 'B', 'B2', 'C']);
       assertQueryResult('key:C+', ['C']);
       assertQueryResult('key:B+', ['B', 'C']);
-      assertQueryResult('key:A++', ['A', 'B', 'C']);
     });
 
     it('should parse upstream star query', () => {
       assertQueryResult('*key:A', ['A']);
       assertQueryResult('*key:B', ['A', 'B']);
-      assertQueryResult('**key:C', ['A', 'B', 'C']);
+      assertQueryResult('*key:C', ['A', 'B', 'B2', 'C']);
     });
 
     it('should parse downstream star query', () => {
-      assertQueryResult('key:C*', ['C']);
+      assertQueryResult('key:A*', ['A', 'B', 'B2', 'C']);
       assertQueryResult('key:B*', ['B', 'C']);
-      assertQueryResult('key:A**', ['A', 'B', 'C']);
+      assertQueryResult('key:C*', ['C']);
+    });
+
+    it('should parse up and down traversal queries', () => {
+      assertQueryResult('key:A* and *key:C', ['A', 'B', 'B2', 'C']);
+      assertQueryResult('*key:B*', ['A', 'B', 'C']);
+      assertQueryResult('key:A* and *key:C and *key:B*', ['A', 'B', 'C']);
+      assertQueryResult('key:A* and *key:B* and *key:C', ['A', 'B', 'C']);
     });
 
     it('should parse sinks query', () => {