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 d79957eb17e91..1ce8cd6869cce 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 @@ -18,6 +18,7 @@ import { ParenthesizedExpressionContext, StartContext, TagAttributeExprContext, + TraversalAllowedExpressionContext, TraversalContext, UpAndDownTraversalExpressionContext, UpTraversalExpressionContext, @@ -77,12 +78,16 @@ 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()); + } + + 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)); const selection_copy = new Set(selection); @@ -94,7 +99,7 @@ export class AntlrAssetSelectionVisitor } visitUpTraversalExpression(ctx: UpTraversalExpressionContext) { - const selection = this.visit(ctx.expr()); + 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) { @@ -104,7 +109,7 @@ export class AntlrAssetSelectionVisitor } visitDownTraversalExpression(ctx: DownTraversalExpressionContext) { - const selection = this.visit(ctx.expr()); + 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) { @@ -130,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()); @@ -164,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); @@ -234,8 +243,4 @@ export class AntlrAssetSelectionVisitor } return selection; } - - visitStart(ctx: StartContext) { - return this.visit(ctx.expr()); - } } 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 a67a0c998e4c7..f27c4238980f0 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 @@ -130,7 +130,8 @@ describe('parseAssetSelectionQuery', () => { 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:C and *key:B*', ['A', 'B', 'C']); + assertQueryResult('key:A* and *key:B* and *key:C', ['A', 'B', 'C']); }); it('should parse sinks query', () => {