Hook up ANTLR in useAssetGraphData.tsx (#26018)

## Summary & Motivation
We want to hook up the new asset selection syntax on the frontend which lives in the `useAssetGraphData` hook.

## How I Tested These Changes
`AntlrAssetSelection.test.ts`

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(

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<any> {
   syntaxError(
@@ -27,10 +30,15 @@ class AntlrInputErrorListener implements ANTLRErrorListener<any> {
   }
 }
 
+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);
+};

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

@@ -63,6 +63,7 @@ export class AntlrAssetSelectionVisitor
   implements AssetSelectionVisitor<Set<AssetGraphQueryItem>>
 {
   all_assets: Set<AssetGraphQueryItem>;
+  focus_assets: Set<AssetGraphQueryItem>;
   traverser: GraphTraverser<AssetGraphQueryItem>;
 
   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) {

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', () => {