[Selection Input] Refactor auto-complete system.

js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetGraphExplorer.tsx

@@ -19,6 +19,7 @@ import {useCallback, useLayoutEffect, useMemo, useState} from 'react';
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 import {AssetGraphAssetSelectionInput} from 'shared/asset-graph/AssetGraphAssetSelectionInput.oss';
 import {useAssetGraphExplorerFilters} from 'shared/asset-graph/useAssetGraphExplorerFilters.oss';
+import {AssetSelectionInput} from 'shared/asset-selection/input/AssetSelectionInput.oss';
 import styled from 'styled-components';
 
 import {AssetEdges} from './AssetEdges';
@@ -53,7 +54,6 @@ import {
 import {AssetLocation, useFindAssetLocation} from './useFindAssetLocation';
 import {featureEnabled} from '../app/Flags';
 import {AssetLiveDataRefreshButton} from '../asset-data/AssetLiveDataProvider';
-import {AssetSelectionInput} from '../asset-selection/input/AssetSelectionInput';
 import {LaunchAssetExecutionButton} from '../assets/LaunchAssetExecutionButton';
 import {AssetKey} from '../assets/types';
 import {DEFAULT_MAX_ZOOM, SVGViewport} from '../graph/SVGViewport';

js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/AssetSelectionInput.tsx → js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/AssetSelectionInput.oss.tsx

@@ -1,12 +1,12 @@
 import type {Linter} from 'codemirror/addon/lint/lint';
-import {useMemo} from 'react';
 import {
   AssetSelectionLexer,
   AssetSelectionParser,
 } from 'shared/asset-selection/AssetSelectionAntlr.oss';
-import {createUseAssetSelectionAutoComplete as defaultCreateUseAssetSelectionAutoComplete} from 'shared/asset-selection/input/useAssetSelectionAutoComplete.oss';
+import {useAssetSelectionAutoCompleteProvider as defaultUseAssetSelectionAutoCompleteProvider} from 'shared/asset-selection/input/useAssetSelectionAutoCompleteProvider.oss';
 
 import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
+import {SelectionAutoCompleteProvider} from '../../selection/SelectionAutoCompleteProvider';
 import {SelectionAutoCompleteInput} from '../../selection/SelectionInput';
 import {createSelectionLinter} from '../../selection/createSelectionLinter';
 
@@ -15,7 +15,9 @@ interface AssetSelectionInputProps {
   value: string;
   onChange: (value: string) => void;
   linter?: Linter<any>;
-  createUseAssetSelectionAutoComplete?: typeof defaultCreateUseAssetSelectionAutoComplete;
+  useAssetSelectionAutoCompleteProvider?: (
+    assets: AssetGraphQueryItem[],
+  ) => SelectionAutoCompleteProvider;
 }
 
 const defaultLinter = createSelectionLinter({
@@ -28,17 +30,14 @@ export const AssetSelectionInput = ({
   onChange,
   assets,
   linter = defaultLinter,
-  createUseAssetSelectionAutoComplete = defaultCreateUseAssetSelectionAutoComplete,
+  useAssetSelectionAutoCompleteProvider = defaultUseAssetSelectionAutoCompleteProvider,
 }: AssetSelectionInputProps) => {
-  const useAssetSelectionAutoComplete = useMemo(
-    () => createUseAssetSelectionAutoComplete(assets),
-    [assets, createUseAssetSelectionAutoComplete],
-  );
+  const SelectionAutoCompleteProvider = useAssetSelectionAutoCompleteProvider(assets);
 
   return (
     <SelectionAutoCompleteInput
       id="asset-selection-input"
-      useAutoComplete={useAssetSelectionAutoComplete}
+      SelectionAutoCompleteProvider={SelectionAutoCompleteProvider}
       placeholder="Search and filter assets"
       linter={linter}
       value={value}

js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/useAssetSelectionAutoCompleteProvider.oss.tsx

@@ -0,0 +1,40 @@
+import {useMemo} from 'react';
+
+import {attributeToIcon, getAttributesMap} from './util';
+import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
+import {createSelectionAutoComplete} from '../../selection/SelectionAutoComplete';
+import {
+  SelectionAutoCompleteProvider,
+  createProvider,
+} from '../../selection/SelectionAutoCompleteProvider';
+
+export function useAssetSelectionAutoCompleteProvider(
+  assets: AssetGraphQueryItem[],
+): SelectionAutoCompleteProvider {
+  const attributesMap = useMemo(() => getAttributesMap(assets), [assets]);
+
+  const baseProvider = useMemo(
+    () =>
+      createProvider({
+        attributesMap,
+        primaryAttributeKey: 'key',
+        attributeToIcon,
+      }),
+    [attributesMap],
+  );
+  const selectionHint = useMemo(() => createSelectionAutoComplete(baseProvider), [baseProvider]);
+
+  return {
+    ...baseProvider,
+    useAutoComplete: ({line, cursorIndex}) => {
+      const autoCompleteResults = useMemo(
+        () => selectionHint(line, cursorIndex),
+        [line, cursorIndex],
+      );
+      return {
+        autoCompleteResults,
+        loading: false,
+      };
+    },
+  };
+}

js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/useAssetSelectionInput.tsx

@@ -1,9 +1,9 @@
 import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 import {AssetGraphAssetSelectionInput} from 'shared/asset-graph/AssetGraphAssetSelectionInput.oss';
+import {AssetSelectionInput} from 'shared/asset-selection/input/AssetSelectionInput.oss';
 import {useAssetSelectionState} from 'shared/asset-selection/useAssetSelectionState.oss';
 import {FilterableAssetDefinition} from 'shared/assets/useAssetDefinitionFilterState.oss';
 
-import {AssetSelectionInput} from './AssetSelectionInput';
 import {featureEnabled} from '../../app/Flags';
 import {useAssetSelectionFiltering} from '../useAssetSelectionFiltering';
 

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

@@ -1,10 +1,14 @@
+import {IconName} from '@dagster-io/ui-components';
+
 import {assertUnreachable} from '../../app/Util';
 import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
+import {isKindTag} from '../../graph/KindTags';
+import {weakMapMemoize} from '../../util/weakMapMemoize';
 import {buildRepoPathForHuman} from '../../workspace/buildRepoAddress';
 
 export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
   const assetNamesSet: Set<string> = new Set();
-  const tagNamesSet: Set<string> = new Set();
+  const tagNamesSet: Set<{key: string; value: string}> = new Set();
   const ownersSet: Set<string> = new Set();
   const groupsSet: Set<string> = new Set();
   const kindsSet: Set<string> = new Set();
@@ -13,13 +17,10 @@ export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
   assets.forEach((asset) => {
     assetNamesSet.add(asset.name);
     asset.node.tags.forEach((tag) => {
-      if (tag.key && tag.value) {
-        // We add quotes around the equal sign here because the auto-complete suggestion already wraps the entire value in quotes.
-        // So wer end up with tag:"key"="value" as the final suggestion
-        tagNamesSet.add(`${tag.key}"="${tag.value}`);
-      } else {
-        tagNamesSet.add(tag.key);
+      if (isKindTag(tag)) {
+        return;
       }
+      tagNamesSet.add(memoizedTag(tag.key, tag.value));
     });
     asset.node.owners.forEach((owner) => {
       switch (owner.__typename) {
@@ -62,3 +63,20 @@ export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
     code_location: codeLocations,
   };
 };
+
+const memoizedTag = weakMapMemoize((key: string, value: string) => ({
+  key,
+  value,
+}));
+
+export type Attribute = 'kind' | 'code_location' | 'group' | 'owner' | 'tag' | 'status' | 'key';
+
+export const attributeToIcon: Record<Attribute, IconName> = {
+  key: 'magnify_glass',
+  kind: 'compute_kind',
+  code_location: 'code_location',
+  group: 'asset_group',
+  owner: 'owner',
+  tag: 'tag',
+  status: 'status',
+};

js_modules/dagster-ui/packages/ui-core/src/gantt/GanttChartSelectionInput.tsx

@@ -1,13 +1,11 @@
-import {useMemo} from 'react';
 import styled from 'styled-components';
 
 import {RunGraphQueryItem} from './toGraphQueryItems';
-import {NO_STATE} from '../run-selection/AntlrRunSelectionVisitor';
+import {useGanttChartSelectionAutoCompleteProvider} from './useGanttChartSelectionAutoCompleteProvider';
 import {RunSelectionLexer} from '../run-selection/generated/RunSelectionLexer';
 import {RunSelectionParser} from '../run-selection/generated/RunSelectionParser';
 import {InputDiv, SelectionAutoCompleteInput} from '../selection/SelectionInput';
 import {createSelectionLinter} from '../selection/createSelectionLinter';
-import {useSelectionInputAutoComplete} from '../selection/useSelectionInputAutoComplete';
 import {weakMapMemoize} from '../util/weakMapMemoize';
 
 export const GanttChartSelectionInput = ({
@@ -19,12 +17,11 @@ export const GanttChartSelectionInput = ({
   value: string;
   onChange: (value: string) => void;
 }) => {
-  const useAutoComplete = useMemo(() => createAutoComplete(items), [items]);
   return (
     <Wrapper>
       <SelectionAutoCompleteInput
         id="run-gantt-chart"
-        useAutoComplete={useAutoComplete}
+        SelectionAutoCompleteProvider={useGanttChartSelectionAutoCompleteProvider(items)}
         placeholder="Search and filter steps"
         linter={getLinter()}
         value={value}
@@ -33,43 +30,10 @@ export const GanttChartSelectionInput = ({
     </Wrapper>
   );
 };
-
-function createAutoComplete(items: RunGraphQueryItem[]) {
-  return function useAutoComplete(value: string, cursor: number) {
-    const attributesMap = useMemo(() => {
-      const statuses = new Set<string>();
-      const names = new Set<string>();
-
-      items.forEach((item) => {
-        if (item.metadata?.state) {
-          statuses.add(item.metadata.state);
-        } else {
-          statuses.add(NO_STATE);
-        }
-        names.add(item.name);
-      });
-      return {name: Array.from(names), status: Array.from(statuses)};
-    }, []);
-
-    return {
-      autoCompleteResults: useSelectionInputAutoComplete({
-        nameBase: 'name',
-        attributesMap,
-        functions: FUNCTIONS,
-        value,
-        cursor,
-      }),
-      loading: false,
-    };
-  };
-}
-
 const getLinter = weakMapMemoize(() =>
   createSelectionLinter({Lexer: RunSelectionLexer, Parser: RunSelectionParser}),
 );
 
-const FUNCTIONS = ['sinks', 'roots'];
-
 const Wrapper = styled.div`
   ${InputDiv} {
     width: 24vw;

js_modules/dagster-ui/packages/ui-core/src/gantt/useGanttChartSelectionAutoCompleteProvider.tsx

@@ -0,0 +1,58 @@
+import {useMemo} from 'react';
+
+import {RunGraphQueryItem} from './toGraphQueryItems';
+import {NO_STATE} from '../run-selection/AntlrRunSelectionVisitor';
+import {createSelectionAutoComplete} from '../selection/SelectionAutoComplete';
+import {
+  SelectionAutoCompleteProvider,
+  createProvider,
+} from '../selection/SelectionAutoCompleteProvider';
+
+export function useGanttChartSelectionAutoCompleteProvider(
+  items: RunGraphQueryItem[],
+): SelectionAutoCompleteProvider {
+  const attributesMap = useMemo(() => {
+    const statuses = new Set<string>();
+    const names = new Set<string>();
+
+    items.forEach((item) => {
+      if (item.metadata?.state) {
+        statuses.add(item.metadata.state);
+      } else {
+        statuses.add(NO_STATE);
+      }
+      names.add(item.name);
+    });
+    return {name: Array.from(names), status: Array.from(statuses)};
+  }, [items]);
+
+  const baseProvider = useMemo(
+    () =>
+      createProvider({
+        attributesMap,
+        primaryAttributeKey: 'name',
+        attributeToIcon: iconMap,
+      }),
+    [attributesMap],
+  );
+  const selectionHint = useMemo(() => createSelectionAutoComplete(baseProvider), [baseProvider]);
+
+  return {
+    ...baseProvider,
+    useAutoComplete: ({line, cursorIndex}) => {
+      const autoCompleteResults = useMemo(
+        () => selectionHint(line, cursorIndex),
+        [line, cursorIndex],
+      );
+      return {
+        autoCompleteResults,
+        loading: false,
+      };
+    },
+  };
+}
+
+export const iconMap = {
+  name: 'magnify_glass',
+  status: 'status',
+} as const;

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

@@ -1,12 +1,11 @@
-import {useMemo} from 'react';
 import styled from 'styled-components';
 
+import {useOpGraphSelectionAutoCompleteProvider} from './useOpGraphSelectionAutoCompleteProvider';
 import {GraphQueryItem} from '../app/GraphQueryImpl';
 import {OpSelectionLexer} from '../op-selection/generated/OpSelectionLexer';
 import {OpSelectionParser} from '../op-selection/generated/OpSelectionParser';
 import {InputDiv, SelectionAutoCompleteInput} from '../selection/SelectionInput';
 import {createSelectionLinter} from '../selection/createSelectionLinter';
-import {useSelectionInputAutoComplete} from '../selection/useSelectionInputAutoComplete';
 import {weakMapMemoize} from '../util/weakMapMemoize';
 
 export const OpGraphSelectionInput = ({
@@ -18,12 +17,11 @@ export const OpGraphSelectionInput = ({
   value: string;
   onChange: (value: string) => void;
 }) => {
-  const useAutoComplete = useMemo(() => createAutoComplete(items), [items]);
   return (
     <Wrapper>
       <SelectionAutoCompleteInput
         id="op-graph"
-        useAutoComplete={useAutoComplete}
+        SelectionAutoCompleteProvider={useOpGraphSelectionAutoCompleteProvider(items)}
         placeholder="Search and filter ops"
         linter={getLinter()}
         value={value}
@@ -33,35 +31,10 @@ export const OpGraphSelectionInput = ({
   );
 };
 
-function createAutoComplete(items: GraphQueryItem[]) {
-  return function useAutoComplete(value: string, cursor: number) {
-    const attributesMap = useMemo(() => {
-      const names = new Set<string>();
-      items.forEach((item) => {
-        names.add(item.name);
-      });
-      return {name: Array.from(names)};
-    }, []);
-
-    return {
-      autoCompleteResults: useSelectionInputAutoComplete({
-        nameBase: 'name',
-        attributesMap,
-        functions: FUNCTIONS,
-        value,
-        cursor,
-      }),
-      loading: false,
-    };
-  };
-}
-
 const getLinter = weakMapMemoize(() =>
   createSelectionLinter({Lexer: OpSelectionLexer, Parser: OpSelectionParser}),
 );
 
-const FUNCTIONS = ['sinks', 'roots'];
-
 const Wrapper = styled.div`
   ${InputDiv} {
     width: 24vw;