Add new selection syntax to job op graphs (#26836)

## Summary & Motivation

1. Consolidate `flagOpSelectionSyntax` into `flagSelectionSyntax`
2. Add a selection syntax input for the op graph and use it if the
feature flag is enabled

## How I Tested These Changes
<img width="1202" alt="Screenshot 2025-01-04 at 12 31 34 AM"
src="https://github.com/user-attachments/assets/433613ee-ff4e-4d3b-a76f-2d2093144caf"
/>

js_modules/dagster-ui/packages/ui-core/src/app/FeatureFlags.oss.tsx

@@ -6,7 +6,6 @@ export enum FeatureFlag {
   flagLegacyRunsPage = 'flagLegacyRunsPage',
   flagSelectionSyntax = 'flagSelectionSyntax',
   flagAssetSelectionWorker = 'flagAssetSelectionWorker',
-  flagOpSelectionSyntax = 'flagOpSelectionSyntax',
 
   // Flags for tests
   __TestFlagDefaultNone = '__TestFlagDefaultNone',

js_modules/dagster-ui/packages/ui-core/src/app/useVisibleFeatureFlagRows.oss.tsx

@@ -33,4 +33,21 @@ export const useVisibleFeatureFlagRows = () => [
       </>
     ),
   },
+  {
+    key: 'Enable new selection syntax for Asset/Op/Run graphs',
+    flagType: FeatureFlag.flagSelectionSyntax,
+    label: (
+      <>
+        Enable new selection syntax for Asset/Op/Run graphs (
+        <a
+          href="https://github.com/dagster-io/dagster/discussions/26849"
+          target="_blank"
+          rel="noreferrer"
+        >
+          Learn more
+        </a>
+        )
+      </>
+    ),
+  },
 ];

js_modules/dagster-ui/packages/ui-core/src/launchpad/OpSelector.tsx

@@ -7,7 +7,6 @@ import {OpSelectorQuery, OpSelectorQueryVariables} from './types/OpSelector.type
 import {filterByQuery} from '../app/GraphQueryImpl';
 import {PYTHON_ERROR_FRAGMENT} from '../app/PythonErrorFragment';
 import {ShortcutHandler} from '../app/ShortcutHandler';
-import {filterOpSelectionByQuery} from '../op-selection/AntlrOpSelection';
 import {explodeCompositesInHandleGraph} from '../pipelines/CompositeSupport';
 import {GRAPH_EXPLORER_SOLID_HANDLE_FRAGMENT} from '../pipelines/GraphExplorer';
 import {GraphQueryInput} from '../ui/GraphQueryInput';
@@ -86,7 +85,7 @@ export const OpSelector = (props: IOpSelectorProps) => {
   const opsFetchError =
     (data?.pipelineOrError.__typename !== 'Pipeline' && data?.pipelineOrError.message) || null;
 
-  const queryResultOps = filterOpSelectionByQuery(ops, query).all;
+  const queryResultOps = filterByQuery(ops, query).all;
   const invalidOpSelection = !loading && queryResultOps.length === 0;
 
   const errorMessage = invalidOpSelection

js_modules/dagster-ui/packages/ui-core/src/op-selection/AntlrOpSelection.ts

@@ -8,15 +8,15 @@ import {OpSelectionLexer} from './generated/OpSelectionLexer';
 import {OpSelectionParser} from './generated/OpSelectionParser';
 import {featureEnabled} from '../app/Flags';
 
-type OpSelectionQueryResult = {
-  all: GraphQueryItem[];
-  focus: GraphQueryItem[];
+type OpSelectionQueryResult<T extends GraphQueryItem> = {
+  all: T[];
+  focus: T[];
 };
 
-export const parseOpSelectionQuery = (
-  all_ops: GraphQueryItem[],
+export const parseOpSelectionQuery = <T extends GraphQueryItem>(
+  all_ops: T[],
   query: string,
-): OpSelectionQueryResult | Error => {
+): OpSelectionQueryResult<T> | Error => {
   try {
     const lexer = new OpSelectionLexer(CharStreams.fromString(query));
     lexer.removeErrorListeners();
@@ -43,11 +43,11 @@ export const parseOpSelectionQuery = (
   }
 };
 
-export const filterOpSelectionByQuery = (
-  all_ops: GraphQueryItem[],
+export const filterOpSelectionByQuery = <T extends GraphQueryItem>(
+  all_ops: T[],
   query: string,
-): OpSelectionQueryResult => {
-  if (featureEnabled(FeatureFlag.flagOpSelectionSyntax)) {
+): OpSelectionQueryResult<T> => {
+  if (featureEnabled(FeatureFlag.flagSelectionSyntax)) {
     const result = parseOpSelectionQuery(all_ops, query);
     if (result instanceof Error) {
       // fall back to old behavior

js_modules/dagster-ui/packages/ui-core/src/op-selection/AntlrOpSelectionVisitor.ts

@@ -19,19 +19,19 @@ import {
 import {OpSelectionVisitor} from './generated/OpSelectionVisitor';
 import {getTraversalDepth, getValue} from '../asset-selection/AntlrAssetSelectionVisitor';
 
-export class AntlrOpSelectionVisitor
-  extends AbstractParseTreeVisitor<Set<GraphQueryItem>>
-  implements OpSelectionVisitor<Set<GraphQueryItem>>
+export class AntlrOpSelectionVisitor<T extends GraphQueryItem>
+  extends AbstractParseTreeVisitor<Set<T>>
+  implements OpSelectionVisitor<Set<T>>
 {
-  all_ops: Set<GraphQueryItem>;
-  focus_ops: Set<GraphQueryItem>;
-  traverser: GraphTraverser<GraphQueryItem>;
+  all_ops: Set<T>;
+  focus_ops: Set<T>;
+  traverser: GraphTraverser<T>;
 
   protected defaultResult() {
-    return new Set<GraphQueryItem>();
+    return new Set<T>();
   }
 
-  constructor(all_ops: GraphQueryItem[]) {
+  constructor(all_ops: T[]) {
     super();
     this.all_ops = new Set(all_ops);
     this.focus_ops = new Set();

js_modules/dagster-ui/packages/ui-core/src/pipelines/GraphExplorer.tsx

@@ -9,19 +9,23 @@ import {
 } from '@dagster-io/ui-components';
 import qs from 'qs';
 import {useEffect, useMemo, useState} from 'react';
+import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 import styled from 'styled-components';
 
 import {EmptyDAGNotice, EntirelyFilteredDAGNotice, LoadingNotice} from './GraphNotices';
 import {ExplorerPath} from './PipelinePathUtils';
 import {SIDEBAR_ROOT_CONTAINER_FRAGMENT} from './SidebarContainerOverview';
 import {SidebarRoot} from './SidebarRoot';
 import {gql} from '../apollo-client';
+import {OpGraphSelectionInput} from './OpGraphSelectionInput';
+import {featureEnabled} from '../app/Flags';
 import {GraphExplorerFragment, GraphExplorerSolidHandleFragment} from './types/GraphExplorer.types';
 import {filterByQuery} from '../app/GraphQueryImpl';
 import {Route} from '../app/Route';
 import {isHiddenAssetGroupJob} from '../asset-graph/Utils';
 import {OP_GRAPH_OP_FRAGMENT, OpGraph} from '../graph/OpGraph';
 import {useOpLayout} from '../graph/asyncGraphLayout';
+import {filterOpSelectionByQuery} from '../op-selection/AntlrOpSelection';
 import {OpNameOrPath} from '../ops/OpNameOrPath';
 import {GraphQueryInput} from '../ui/GraphQueryInput';
 import {RepoAddress} from '../workspace/types';
@@ -156,10 +160,15 @@ export const GraphExplorer = (props: GraphExplorerProps) => {
     (options.explodeComposites ||
       solids.some((f) => f.definition.__typename === 'CompositeSolidDefinition'));
 
-  const queryResultOps = useMemo(
-    () => (solidsQueryEnabled ? filterByQuery(solids, opsQuery) : {all: solids, focus: []}),
-    [opsQuery, solids, solidsQueryEnabled],
-  );
+  const queryResultOps = useMemo(() => {
+    if (solidsQueryEnabled) {
+      if (featureEnabled(FeatureFlag.flagSelectionSyntax)) {
+        return filterOpSelectionByQuery(solids, opsQuery);
+      }
+      return filterByQuery(solids, opsQuery);
+    }
+    return {all: solids, focus: []};
+  }, [opsQuery, solids, solidsQueryEnabled]);
 
   const highlightedOps = useMemo(
     () => queryResultOps.all.filter((s) => s.name.toLowerCase().includes(nameMatch.toLowerCase())),
@@ -193,13 +202,21 @@ export const GraphExplorer = (props: GraphExplorerProps) => {
         <ErrorBoundary region="op graph">
           {solidsQueryEnabled ? (
             <QueryOverlay>
-              <GraphQueryInput
-                items={solids}
-                value={explorerPath.opsQuery}
-                placeholder="Type an op subset…"
-                popoverPosition="bottom-left"
-                onChange={handleQueryChange}
-              />
+              {featureEnabled(FeatureFlag.flagSelectionSyntax) ? (
+                <OpGraphSelectionInput
+                  items={solids}
+                  value={explorerPath.opsQuery}
+                  onChange={handleQueryChange}
+                />
+              ) : (
+                <GraphQueryInput
+                  items={solids}
+                  value={explorerPath.opsQuery}
+                  placeholder="Type an op subset…"
+                  popoverPosition="bottom-left"
+                  onChange={handleQueryChange}
+                />
+              )}
             </QueryOverlay>
           ) : breadcrumbs.length > 1 ? (
             <BreadcrumbsOverlay>

js_modules/dagster-ui/packages/ui-core/src/pipelines/OpGraphSelectionInput.tsx

@@ -0,0 +1,53 @@
+import {useMemo} from 'react';
+import styled from 'styled-components';
+
+import {GraphQueryItem} from '../app/GraphQueryImpl';
+import {OpSelectionLexer} from '../op-selection/generated/OpSelectionLexer';
+import {OpSelectionParser} from '../op-selection/generated/OpSelectionParser';
+import {InputDiv, SelectionAutoCompleteInput} from '../selection/SelectionAutoCompleteInput';
+import {createSelectionLinter} from '../selection/createSelectionLinter';
+import {weakMapMemoize} from '../util/weakMapMemoize';
+
+export const OpGraphSelectionInput = ({
+  items,
+  value,
+  onChange,
+}: {
+  items: GraphQueryItem[];
+  value: string;
+  onChange: (value: string) => void;
+}) => {
+  const attributesMap = useMemo(() => {
+    const names = new Set<string>();
+    items.forEach((item) => {
+      names.add(item.name);
+    });
+    return {name: Array.from(names)};
+  }, [items]);
+
+  return (
+    <Wrapper>
+      <SelectionAutoCompleteInput
+        nameBase="name"
+        attributesMap={attributesMap}
+        placeholder="Type an op subset"
+        functions={FUNCTIONS}
+        linter={getLinter()}
+        value={value}
+        onChange={onChange}
+      />
+    </Wrapper>
+  );
+};
+
+const getLinter = weakMapMemoize(() =>
+  createSelectionLinter({Lexer: OpSelectionLexer, Parser: OpSelectionParser}),
+);
+
+const FUNCTIONS = ['sinks', 'roots'];
+
+const Wrapper = styled.div`
+  ${InputDiv} {
+    width: 24vw;
+  }
+`;