tag fix

dagster-io · Jan 26, 2025 · 794d7e3 · 794d7e3
1 parent 2480c65
commit 794d7e3
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 22 deletions.
diff --git a/.../packages/ui-core/src/asset-selection/input/useAssetSelectionAutoCompleteProvider.oss.tsx b/.../packages/ui-core/src/asset-selection/input/useAssetSelectionAutoCompleteProvider.oss.tsx
@@ -3,7 +3,6 @@ import {useMemo} from 'react';
 
 import {getAttributesMap} from './util';
 import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
-import {useUpdatingRef} from '../../hooks/useUpdatingRef';
 import {createSelectionAutoComplete} from '../../selection/SelectionAutoComplete';
 import {
   BaseSuggestion,
@@ -20,7 +19,14 @@ type Suggestion =
       text: string;
       displayText: string;
       type: 'function' | 'attribute-value' | 'attribute-with-value';
-      attributeName?: string;
+      attributeName?: Omit<Attribute, 'tag'>;
+    }
+  | {
+      text: string;
+      displayText: string;
+      type: 'attribute-value' | 'attribute-with-value';
+      attributeName: 'tag';
+      tag: {key: string; value: string};
     }
   | {
       text: string;
@@ -37,19 +43,25 @@ type Suggestion =
 export function useAssetSelectionAutoCompleteProvider(
   assets: AssetGraphQueryItem[],
 ): SelectionAutoCompleteProvider<Suggestion> {
-  const attributesMapRef = useUpdatingRef(getAttributesMap(assets));
+  const attributesMap = useMemo(() => getAttributesMap(assets), [assets]);
 
   const baseProvider = useMemo(
     () =>
-      createSelectionAutoCompleteProviderFromAttributeMap<
-        typeof attributesMapRef.current,
-        'key',
-        Suggestion
-      >({
+      createSelectionAutoCompleteProviderFromAttributeMap<typeof attributesMap, 'key', Suggestion>({
         nameBase: 'key',
-        attributesMapRef,
+        attributesMap,
         functions: FUNCTIONS,
-        doesValueIncludeQuery: (_attribute, value, query) => value.includes(query),
+        doesValueIncludeQuery: (_attribute, value, query) => {
+          if (typeof value !== 'string') {
+            // This is a tag
+            return (
+              value.key.includes(query) ||
+              value.value.includes(query) ||
+              `${value.key}=${value.value}`.includes(query)
+            );
+          }
+          return value.includes(query);
+        },
         createAttributeSuggestion: (attribute, textCallback) => {
           const text = `${attribute}:`;
           return {
@@ -61,10 +73,23 @@ export function useAssetSelectionAutoCompleteProvider(
           };
         },
         createAttributeValueSuggestion: (attribute, value, textCallback) => {
-          const text = `"${value}"`;
+          let text;
+          let displayText;
+          if (typeof value !== 'string') {
+            if (value.key && value.value) {
+              text = `"${value.key}"="${value.value}"`;
+              displayText = `${value.key}=${value.value}`;
+            } else {
+              text = `"${value.key}"`;
+              displayText = value.key;
+            }
+          } else {
+            text = `"${value}"`;
+            displayText = value;
+          }
           return {
             text: textCallback ? textCallback(text) : text,
-            displayText: value,
+            displayText,
             type: 'attribute-value',
             attributeName: attribute,
           };
@@ -95,7 +120,7 @@ export function useAssetSelectionAutoCompleteProvider(
           };
         },
       }),
-    [attributesMapRef],
+    [attributesMap],
   );
   const selectionHint = useMemo(() => createSelectionAutoComplete(baseProvider), [baseProvider]);
 
@@ -124,7 +149,7 @@ const attributeToIcon: Record<Attribute, IconName> = {
   status: 'status',
 };
 
-export const SuggestionItem = ({suggestion}: {suggestion: Suggestion | BaseSuggestion}) => {
+const SuggestionItem = ({suggestion}: {suggestion: Suggestion | BaseSuggestion}) => {
   let label;
   let icon: IconName | null = null;
   let value: string | null = 'displayText' in suggestion ? suggestion.displayText : null;

diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/util.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/util.ts
@@ -1,11 +1,12 @@
 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();
@@ -17,13 +18,7 @@ export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
       if (isKindTag(tag)) {
         return;
       }
-      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);
-      }
+      tagNamesSet.add(memoizedTag(tag.key, tag.value));
     });
     asset.node.owners.forEach((owner) => {
       switch (owner.__typename) {
@@ -66,3 +61,8 @@ export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
     code_location: codeLocations,
   };
 };
+
+const memoizedTag = weakMapMemoize((key: string, value: string) => ({
+  key,
+  value,
+}));