diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx index e24833c98444e..b30668c72bc5a 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx +++ b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/useAssetGraphData.tsx @@ -13,7 +13,8 @@ import { AssetNodeForGraphQueryFragment, } from './types/useAssetGraphData.types'; import {usePrefixedCacheKey} from '../app/AppProvider'; -import {GraphQueryItem, filterByQuery} from '../app/GraphQueryImpl'; +import {GraphQueryItem} from '../app/GraphQueryImpl'; +import {filterAssetSelectionByQuery} from '../asset-selection/AntlrAssetSelection'; import {AssetKey} from '../assets/types'; import {AssetGroupSelector, PipelineSelector} from '../graphql/types'; import {useIndexedDBCachedQuery} from '../search/useIndexedDBCachedQuery'; @@ -117,7 +118,7 @@ export function useAssetGraphData(opsQuery: string, options: AssetGraphFetchScop // In the future it might be ideal to move this server-side, but we currently // get to leverage the useQuery cache almost 100% of the time above, making this // super fast after the first load vs a network fetch on every page view. - const {all: allFilteredByOpQuery} = filterByQuery(graphQueryItems, opsQuery); + const {all: allFilteredByOpQuery} = filterAssetSelectionByQuery(graphQueryItems, opsQuery); const kinds = options.kinds?.map((c) => c.toLowerCase()); const all = kinds?.length ? allFilteredByOpQuery.filter( 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 fb52dc47984aa..16a506158cfd5 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 @@ -5,11 +5,14 @@ import { RecognitionException, Recognizer, } from 'antlr4ts'; +import {FeatureFlag} from 'shared/app/FeatureFlags.oss'; import {AntlrAssetSelectionVisitor} from './AntlrAssetSelectionVisitor'; import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData'; import {AssetSelectionLexer} from './generated/AssetSelectionLexer'; import {AssetSelectionParser} from './generated/AssetSelectionParser'; +import {featureEnabled} from '../app/Flags'; +import {filterByQuery} from '../app/GraphQueryImpl'; class AntlrInputErrorListener implements ANTLRErrorListener { syntaxError( @@ -27,10 +30,15 @@ class AntlrInputErrorListener implements ANTLRErrorListener { } } +type AssetSelectionQueryResult = { + all: AssetGraphQueryItem[]; + focus: AssetGraphQueryItem[]; +}; + export const parseAssetSelectionQuery = ( all_assets: AssetGraphQueryItem[], query: string, -): AssetGraphQueryItem[] | Error => { +): AssetSelectionQueryResult | Error => { try { const lexer = new AssetSelectionLexer(CharStreams.fromString(query)); lexer.removeErrorListeners(); @@ -45,8 +53,29 @@ export const parseAssetSelectionQuery = ( const tree = parser.start(); const visitor = new AntlrAssetSelectionVisitor(all_assets); - return [...visitor.visit(tree)]; + const all_selection = visitor.visit(tree); + const focus_selection = visitor.focus_assets; + + return { + all: Array.from(all_selection), + focus: Array.from(focus_selection), + }; } catch (e) { return e as Error; } }; + +export const filterAssetSelectionByQuery = ( + all_assets: AssetGraphQueryItem[], + query: string, +): AssetSelectionQueryResult => { + if (featureEnabled(FeatureFlag.flagAssetSelectionSyntax)) { + const result = parseAssetSelectionQuery(all_assets, query); + if (result instanceof Error) { + // fall back to old behavior + return filterByQuery(all_assets, query); + } + return result; + } + return filterByQuery(all_assets, query); +}; 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 756fda6b6807a..b639d61f2f9a3 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 @@ -63,6 +63,7 @@ export class AntlrAssetSelectionVisitor implements AssetSelectionVisitor> { all_assets: Set; + focus_assets: Set; traverser: GraphTraverser; protected defaultResult() { @@ -72,6 +73,7 @@ export class AntlrAssetSelectionVisitor constructor(all_assets: AssetGraphQueryItem[]) { super(); this.all_assets = new Set(all_assets); + this.focus_assets = new Set(); this.traverser = new GraphTraverser(all_assets); } @@ -165,12 +167,16 @@ export class AntlrAssetSelectionVisitor visitKeyExpr(ctx: KeyExprContext) { const value: string = getValue(ctx.value()); - return new Set([...this.all_assets].filter((i) => i.name === value)); + const selection = [...this.all_assets].filter((i) => i.name === value); + selection.forEach((i) => this.focus_assets.add(i)); + return new Set(selection); } visitKeySubstringExpr(ctx: KeySubstringExprContext) { const value: string = getValue(ctx.value()); - return new Set([...this.all_assets].filter((i) => i.name.includes(value))); + const selection = [...this.all_assets].filter((i) => i.name.includes(value)); + selection.forEach((i) => this.focus_assets.add(i)); + return new Set(selection); } visitTagAttributeExpr(ctx: TagAttributeExprContext) { 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 f2b2c6eed0810..6ba6a43624c0a 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 @@ -52,8 +52,8 @@ function assertQueryResult(query: string, expectedNames: string[]) { if (result instanceof Error) { throw result; } - expect(result.length).toBe(expectedNames.length); - expect(new Set(result.map((asset) => asset.name))).toEqual(new Set(expectedNames)); + expect(result.all.length).toBe(expectedNames.length); + expect(new Set(result.all.map((asset) => asset.name))).toEqual(new Set(expectedNames)); } describe('parseAssetSelectionQuery', () => {