[2/4] Add asset selection filtering to asset catalog table. (#25970)

## Summary & Motivation

Add asset selection filtering to asset catalog table.

## How I Tested These Changes

Manually tested
<img width="1718" alt="Screenshot 2024-11-16 at 7 54 23 PM"
src="https://github.com/user-attachments/assets/df535d76-2ce5-431b-a89c-34f59dc79173">

## Changelog

[ui] The asset catalog now supports filtering using the asset selection
syntax.

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

@@ -718,7 +718,6 @@ const AssetGraphExplorerWithData = ({
             <TopbarWrapper>
               <Box flex={{direction: 'column'}} style={{width: '100%'}}>
                 <Box
-                  border={filterBar ? 'bottom' : undefined}
                   flex={{gap: 12, alignItems: 'center'}}
                   padding={{left: showSidebar ? 12 : 24, vertical: 12, right: 12}}
                 >

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

@@ -24,6 +24,7 @@ export const AssetGraphFilterBar = ({
     <Box
       flex={{direction: 'row', justifyContent: 'space-between', gap: 12, alignItems: 'center'}}
       padding={{vertical: 8, horizontal: 12}}
+      border="top"
     >
       <Box flex={{gap: 12, alignItems: 'center', direction: 'row', grow: 1}}>
         {activeFiltersJsx}

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

@@ -306,7 +306,6 @@ export const AssetGraphExplorerSidebar = React.memo(
             padding: '12px 24px',
             paddingRight: 12,
           }}
-          border="bottom"
         >
           <SearchFilter
             values={React.useMemo(() => {
@@ -321,7 +320,7 @@ export const AssetGraphExplorerSidebar = React.memo(
             <Button icon={<Icon name="panel_show_right" />} onClick={hideSidebar} />
           </Tooltip>
         </Box>
-        <div>
+        <Box border="top">
           {loading ? (
             <Box flex={{direction: 'column', gap: 9}} padding={12}>
               <Skeleton $height={21} $width="50%" />
@@ -407,7 +406,7 @@ export const AssetGraphExplorerSidebar = React.memo(
               </Inner>
             </Container>
           )}
-        </div>
+        </Box>
       </div>
     );
   },

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

@@ -0,0 +1,55 @@
+import {useMemo} from 'react';
+import {FilterableAssetDefinition} from 'shared/assets/useAssetDefinitionFilterState.oss';
+
+import {COMMON_COLLATOR} from '../app/Util';
+import {tokenForAssetKey} from '../asset-graph/Utils';
+import {AssetNodeForGraphQueryFragment} from '../asset-graph/types/useAssetGraphData.types';
+import {useAssetGraphData} from '../asset-graph/useAssetGraphData';
+
+export const useAssetSelectionFiltering = <
+  T extends {
+    id: string;
+    key: {path: Array<string>};
+    definition?: FilterableAssetDefinition | null;
+  },
+>({
+  assetSelection,
+  assets,
+}: {
+  assetSelection: string;
+
+  assets: T[];
+}) => {
+  const assetsByKey = useMemo(
+    () => Object.fromEntries(assets.map((asset) => [tokenForAssetKey(asset.key), asset])),
+    [assets],
+  );
+
+  const {fetchResult, graphQueryItems, graphAssetKeys} = useAssetGraphData(
+    assetSelection,
+    useMemo(
+      () => ({
+        hideEdgesToNodesOutsideQuery: true,
+        hideNodesMatching: (node: AssetNodeForGraphQueryFragment) => {
+          return !assetsByKey[tokenForAssetKey(node.assetKey)];
+        },
+      }),
+      [assetsByKey],
+    ),
+  );
+
+  const filtered = useMemo(() => {
+    return (
+      graphAssetKeys
+        .map((key) => assetsByKey[tokenForAssetKey(key)]!)
+        .sort((a, b) => COMMON_COLLATOR.compare(a.key.path.join(''), b.key.path.join(''))) ?? []
+    );
+  }, [graphAssetKeys, assetsByKey]);
+
+  const filteredByKey = useMemo(
+    () => Object.fromEntries(filtered.map((asset) => [tokenForAssetKey(asset.key), asset])),
+    [filtered],
+  );
+
+  return {filtered, filteredByKey, fetchResult, graphAssetKeys, graphQueryItems};
+};

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

@@ -0,0 +1,34 @@
+import {AssetGraphAssetSelectionInput} from 'shared/asset-graph/AssetGraphAssetSelectionInput.oss';
+import {useAssetSelectionState} from 'shared/asset-selection/useAssetSelectionState.oss';
+import {FilterableAssetDefinition} from 'shared/assets/useAssetDefinitionFilterState.oss';
+
+import {useAssetSelectionFiltering} from './useAssetSelectionFiltering';
+
+export const useAssetSelectionInput = <
+  T extends {
+    id: string;
+    key: {path: Array<string>};
+    definition?: FilterableAssetDefinition | null;
+  },
+>(
+  assets: T[],
+) => {
+  const [assetSelection, setAssetSelection] = useAssetSelectionState();
+
+  const {graphQueryItems, fetchResult, filtered} = useAssetSelectionFiltering({
+    assetSelection,
+    assets,
+  });
+
+  const filterInput = (
+    <AssetGraphAssetSelectionInput
+      items={graphQueryItems}
+      value={assetSelection}
+      placeholder="Type an asset subset…"
+      onChange={setAssetSelection}
+      popoverPosition="bottom-left"
+    />
+  );
+
+  return {filterInput, fetchResult, filtered, assetSelection, setAssetSelection};
+};

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

@@ -0,0 +1,8 @@
+import {useQueryPersistedState} from '../hooks/useQueryPersistedState';
+
+export function useAssetSelectionState() {
+  return useQueryPersistedState<string>({
+    queryKey: 'asset-selection',
+    defaults: {['asset-selection']: ''},
+  });
+}

js_modules/dagster-ui/packages/ui-core/src/assets/AssetTable.tsx

@@ -39,7 +39,7 @@ interface Props {
   belowActionBarComponents: React.ReactNode;
   prefixPath: string[];
   displayPathForAsset: (asset: Asset) => string[];
-  searchPath: string;
+  assetSelection: string;
   isFiltered: boolean;
   kindFilter?: StaticSetFilter<string>;
   isLoading: boolean;
@@ -52,7 +52,7 @@ export const AssetTable = ({
   refreshState,
   prefixPath,
   displayPathForAsset,
-  searchPath,
+  assetSelection,
   isFiltered,
   view,
   kindFilter,
@@ -81,7 +81,7 @@ export const AssetTable = ({
 
   const content = () => {
     if (!assets.length) {
-      if (searchPath) {
+      if (assetSelection) {
         return (
           <Box padding={{top: 64}}>
             <NonIdealState
@@ -90,12 +90,12 @@ export const AssetTable = ({
               description={
                 isFiltered ? (
                   <div>
-                    No assets matching <strong>{searchPath}</strong> were found in the selected
+                    No assets matching <strong>{assetSelection}</strong> were found in the selected
                     filters
                   </div>
                 ) : (
                   <div>
-                    No assets matching <strong>{searchPath}</strong> were found
+                    No assets matching <strong>{assetSelection}</strong> were found
                   </div>
                 )
               }

js_modules/dagster-ui/packages/ui-core/src/assets/AssetsCatalogTable.tsx

@@ -3,7 +3,7 @@ import * as React from 'react';
 import {useCallback, useContext, useEffect, useLayoutEffect, useMemo, useState} from 'react';
 import {useRouteMatch} from 'react-router-dom';
 import {useSetRecoilState} from 'recoil';
-import {AssetCatalogTableBottomActionBar} from 'shared/assets/AssetCatalogTableBottomActionBar.oss';
+import {AssetGraphFilterBar} from 'shared/asset-graph/AssetGraphFilterBar.oss';
 import {useAssetCatalogFiltering} from 'shared/assets/useAssetCatalogFiltering.oss';
 
 import {AssetTable} from './AssetTable';
@@ -19,14 +19,14 @@ import {
   AssetCatalogTableQueryVersion,
 } from './types/AssetsCatalogTable.types';
 import {AssetViewType, useAssetView} from './useAssetView';
-import {useBasicAssetSearchInput} from './useBasicAssetSearchInput';
 import {gql, useApolloClient} from '../apollo-client';
 import {AppContext} from '../app/AppContext';
 import {PYTHON_ERROR_FRAGMENT} from '../app/PythonErrorFragment';
 import {PythonErrorInfo} from '../app/PythonErrorInfo';
 import {FIFTEEN_SECONDS, useRefreshAtInterval} from '../app/QueryRefresh';
 import {currentPageAtom} from '../app/analytics';
 import {PythonErrorFragment} from '../app/types/PythonErrorFragment.types';
+import {useAssetSelectionInput} from '../asset-selection/useAssetSelectionInput';
 import {AssetGroupSelector} from '../graphql/types';
 import {useUpdatingRef} from '../hooks/useUpdatingRef';
 import {useBlockTraceUntilTrue} from '../performance/TraceContext';
@@ -210,13 +210,10 @@ export const AssetsCatalogTable = ({
     activeFiltersJsx,
     kindFilter,
   } = useAssetCatalogFiltering({assets});
+  const {filterInput, filtered, fetchResult, assetSelection, setAssetSelection} =
+    useAssetSelectionInput(partiallyFiltered);
 
-  const {searchPath, filterInput, filtered} = useBasicAssetSearchInput(
-    partiallyFiltered,
-    prefixPath,
-  );
-
-  useBlockTraceUntilTrue('useAllAssets', !!assets?.length);
+  useBlockTraceUntilTrue('useAllAssets', !!assets?.length && !fetchResult.loading);
 
   const {displayPathForAsset, displayed} = useMemo(
     () =>
@@ -280,11 +277,15 @@ export const AssetsCatalogTable = ({
         </>
       }
       belowActionBarComponents={
-        <AssetCatalogTableBottomActionBar activeFiltersJsx={activeFiltersJsx} />
+        <AssetGraphFilterBar
+          activeFiltersJsx={activeFiltersJsx}
+          assetSelection={assetSelection}
+          setAssetSelection={setAssetSelection}
+        />
       }
       refreshState={refreshState}
       prefixPath={prefixPath || emptyArray}
-      searchPath={searchPath}
+      assetSelection={assetSelection}
       displayPathForAsset={displayPathForAsset}
       kindFilter={kindFilter}
     />

js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/AssetTables.test.tsx

@@ -32,6 +32,20 @@ const MOCKS = [
 // This file must be mocked because Jest can't handle `import.meta.url`.
 jest.mock('../../graph/asyncGraphLayout', () => ({}));
 
+jest.mock('shared/asset-selection/useAssetSelectionInput', () => {
+  return {
+    useAssetSelectionInput: (assets: any) => {
+      return {
+        filterInput: <div />,
+        fetchResult: {loading: false},
+        filtered: assets,
+        assetSelection: '',
+        setAssetSelection: () => {},
+      };
+    },
+  };
+});
+
 describe('AssetTable', () => {
   beforeAll(() => {
     mockViewportClientRect();

js_modules/dagster-ui/packages/ui-core/src/search/types.ts

@@ -9,7 +9,7 @@ import {
   SearchScheduleFragment,
   SearchSensorFragment,
 } from './types/useGlobalSearch.types';
-import {DefinitionTag} from '../graphql/types';
+import {AssetKey, DefinitionTag} from '../graphql/types';
 
 export enum SearchResultType {
   AssetGroup,
@@ -54,6 +54,7 @@ export function isAssetFilterSearchResultType(
 }
 
 export type SearchResult = {
+  key?: AssetKey;
   label: string;
   description: string;
   href: string;