feat: support search on collection tree (#917)

Signed-off-by: ryjiang <jiangruiyi@gmail.com>

client/src/i18n/cn/collection.ts

@@ -2,6 +2,7 @@ const collectionTrans = {
   noLoadData: '没有加载的 Collection',
   noData: '没有 Collection',
   collectionId: 'Collection ID',
+  searchCollections: '搜索 Collection',
 
   rowCount: 'Entity 数量（大约）',
   count: 'Entity 数量',

client/src/i18n/en/collection.ts

@@ -2,6 +2,7 @@ const collectionTrans = {
   noLoadData: 'No Loaded Collection',
   noData: 'No Collection',
   collectionId: 'Collection ID',
+  searchCollections: 'Search collections',
 
   rowCount: 'Approx Count',
   count: 'Entity Count',

client/src/pages/databases/tree/CollectionNode.tsx

@@ -0,0 +1,80 @@
+import React from 'react';
+import { Tooltip, Typography, Box, IconButton } from '@mui/material';
+import { CollectionObject } from '@server/types';
+import { formatNumber } from '@/utils';
+import { TreeNodeBox, Count, StatusDot } from './StyledComponents';
+import HighlightText from './HighlightText';
+import icons from '@/components/icons/Icons';
+
+interface CollectionNodeProps {
+  data: CollectionObject;
+  highlight?: string;
+  isSelected?: boolean;
+  isContextMenuTarget?: boolean;
+  hasChildren?: boolean;
+  isExpanded?: boolean;
+  depth?: number;
+  onToggleExpand?: (e: React.MouseEvent) => void;
+  onClick?: () => void;
+  onContextMenu?: (e: React.MouseEvent) => void;
+}
+
+const CollectionNode: React.FC<CollectionNodeProps> = ({
+  data,
+  highlight = '',
+  isSelected = false,
+  isContextMenuTarget = false,
+  hasChildren = false,
+  isExpanded = false,
+  depth = 0,
+  onToggleExpand,
+  onClick,
+  onContextMenu,
+}) => {
+  const ExpandIcon = icons.rightArrow;
+
+  const handleToggleExpand = (e: React.MouseEvent) => {
+    e.stopPropagation();
+    onToggleExpand?.(e);
+  };
+
+  return (
+    <TreeNodeBox
+      isSelected={isSelected}
+      isContextMenuTarget={isContextMenuTarget}
+      isCollectionNoSchema={!data.schema}
+      depth={depth}
+      onClick={onClick}
+      onContextMenu={onContextMenu}
+    >
+      {hasChildren && (
+        <IconButton size="small" onClick={handleToggleExpand} sx={{ mr: 0 }}>
+          {isExpanded ? (
+            <ExpandIcon sx={{ transform: `rotate(90deg)` }} />
+          ) : (
+            <ExpandIcon />
+          )}
+        </IconButton>
+      )}
+
+      <Box sx={{ mr: 1, display: 'flex', alignItems: 'center' }}>
+        <StatusDot status={getStatus(data)} />
+      </Box>
+
+      <Tooltip title={data.collection_name} placement="top">
+        <Typography noWrap className="collectionName">
+          <HighlightText text={data.collection_name} highlight={highlight} />
+        </Typography>
+      </Tooltip>
+      <Count>({formatNumber(data.rowCount || 0)})</Count>
+    </TreeNodeBox>
+  );
+};
+
+const getStatus = (data: CollectionObject) => {
+  if (!data.schema || !data.schema.hasVectorIndex) return 'noIndex';
+  if (data.status === 'loading') return 'loading';
+  return data.loaded ? 'loaded' : 'unloaded';
+};
+
+export default CollectionNode;

client/src/pages/databases/tree/DatabaseNode.tsx

@@ -0,0 +1,73 @@
+import React from 'react';
+import { Box, IconButton, Tooltip, Typography } from '@mui/material';
+import icons from '@/components/icons/Icons';
+import { NodeContent, CollectionCount, TreeNodeBox } from './StyledComponents';
+
+interface DatabaseNodeProps {
+  database: string;
+  collectionCount: number;
+  isSelected: boolean;
+  isContextMenuTarget: boolean;
+  onNodeClick: () => void;
+  onContextMenu: (event: React.MouseEvent) => void;
+  onSearchClick: (event: React.MouseEvent) => void;
+}
+
+const DatabaseNode: React.FC<DatabaseNodeProps> = ({
+  database,
+  collectionCount,
+  isSelected,
+  isContextMenuTarget,
+  onNodeClick,
+  onContextMenu,
+  onSearchClick,
+}) => {
+  const DatabaseIcon = icons.database;
+  const SearchIcon = icons.search;
+
+  return (
+    <TreeNodeBox
+      isSelected={isSelected}
+      isContextMenuTarget={isContextMenuTarget}
+      isCollectionNoSchema={false}
+      depth={0}
+      sx={{
+        position: 'sticky',
+        top: 0,
+        zIndex: 2,
+      }}
+      onClick={onNodeClick}
+      onContextMenu={onContextMenu}
+    >
+      <Box sx={{ width: 0 }} />
+      <Box sx={{ mr: 1, display: 'flex', alignItems: 'center' }}>
+        <DatabaseIcon />
+      </Box>
+      <NodeContent>
+        <Tooltip title={database} placement="top">
+          <Typography noWrap sx={{ flex: 1 }}>
+            {database}
+            <CollectionCount>({collectionCount})</CollectionCount>
+          </Typography>
+        </Tooltip>
+        <IconButton
+          size="small"
+          onClick={onSearchClick}
+          sx={{
+            padding: '4px',
+            ml: 0.5,
+            position: 'relative',
+            right: '-8px',
+            '&:hover': {
+              backgroundColor: 'transparent',
+            },
+          }}
+        >
+          <SearchIcon sx={{ fontSize: '16px', color: 'text.secondary' }} />
+        </IconButton>
+      </NodeContent>
+    </TreeNodeBox>
+  );
+};
+
+export default DatabaseNode;

client/src/pages/databases/tree/HighlightText.tsx

@@ -0,0 +1,43 @@
+import React from 'react';
+import { useTheme } from '@mui/material/styles';
+
+interface HighlightTextProps {
+  text: string;
+  highlight: string;
+}
+
+const HighlightText: React.FC<HighlightTextProps> = ({ text, highlight }) => {
+  const theme = useTheme();
+  const isLight = theme.palette.mode === 'light';
+
+  if (!highlight.trim()) {
+    return <>{text}</>;
+  }
+  const regex = new RegExp(`(${highlight})`, 'gi');
+  const parts = text.split(regex);
+
+  return (
+    <>
+      {parts.map((part, i) =>
+        regex.test(part) ? (
+          <span
+            key={i}
+            style={{
+              backgroundColor: isLight
+                ? theme.palette.highlight.light
+                : theme.palette.highlight.dark,
+              borderRadius: '2px',
+              padding: '0 2px',
+            }}
+          >
+            {part}
+          </span>
+        ) : (
+          part
+        )
+      )}
+    </>
+  );
+};
+
+export default HighlightText;

client/src/pages/databases/tree/SearchBox.tsx

@@ -0,0 +1,101 @@
+import React from 'react';
+import { useTranslation } from 'react-i18next';
+import { Box, TextField, InputAdornment, IconButton } from '@mui/material';
+import icons from '@/components/icons/Icons';
+
+interface SearchBoxProps {
+  searchQuery: string;
+  onSearchChange: (value: string) => void;
+  onClose: () => void;
+}
+
+const SearchBox: React.FC<SearchBoxProps> = ({
+  searchQuery,
+  onSearchChange,
+  onClose,
+}) => {
+  const { t: collectionTrans } = useTranslation('collection');
+  const SearchIcon = icons.search;
+  const ClearIcon = icons.clear;
+
+  return (
+    <Box
+      sx={{
+        width: '100%',
+        display: 'flex',
+        alignItems: 'center',
+        backgroundColor: 'background.grey',
+      }}
+    >
+      <TextField
+        size="small"
+        placeholder={collectionTrans('searchCollections')}
+        value={searchQuery}
+        onChange={e => onSearchChange(e.target.value)}
+        autoComplete="off"
+        autoFocus
+        fullWidth
+        InputProps={{
+          sx: {
+            paddingLeft: '8px',
+            paddingRight: '8px',
+            height: '30px',
+          },
+          startAdornment: (
+            <InputAdornment position="start">
+              <SearchIcon sx={{ fontSize: '16px', color: 'text.secondary' }} />
+            </InputAdornment>
+          ),
+          endAdornment: (
+            <InputAdornment position="end">
+              <IconButton
+                size="small"
+                onClick={onClose}
+                edge="end"
+                sx={{
+                  padding: '4px',
+                  '&:hover': {
+                    backgroundColor: 'transparent',
+                  },
+                }}
+              >
+                <ClearIcon
+                  sx={{
+                    fontSize: '16px',
+                    color: 'text.secondary',
+                  }}
+                />
+              </IconButton>
+            </InputAdornment>
+          ),
+        }}
+        sx={{
+          '& .MuiOutlinedInput-root': {
+            backgroundColor: 'transparent',
+            fontSize: '15px',
+            '& fieldset': {
+              borderColor: 'transparent',
+            },
+            '&:hover fieldset': {
+              borderColor: 'transparent',
+            },
+            '&.Mui-focused fieldset': {
+              borderColor: 'transparent',
+            },
+          },
+          '& .MuiInputBase-input': {
+            padding: '4px 0',
+            fontSize: '13px',
+            '&::placeholder': {
+              fontSize: '13px',
+              opacity: 1,
+              color: 'text.disabled',
+            },
+          },
+        }}
+      />
+    </Box>
+  );
+};
+
+export default SearchBox;

client/src/pages/databases/tree/StyledComponents.ts

@@ -1,24 +1,8 @@
 import { styled } from '@mui/material/styles';
 import { Box, MenuItem, Divider } from '@mui/material';
 
-export const CollectionNodeWrapper = styled(Box)({
-  minHeight: '24px',
-  lineHeight: '24px',
-  display: 'flex',
-  minWidth: 190,
-});
-
-export const CollectionNameWrapper = styled(Box)({
-  display: 'flex',
-  alignItems: 'center',
-  width: 'calc(100% - 45px)',
-  '& .collectionName': {
-    minWidth: 36,
-  },
-});
-
 export const Count = styled('span')(({ theme }) => ({
-  fontSize: '13px',
+  fontSize: '11px',
   fontWeight: 500,
   marginLeft: theme.spacing(0.5),
   color: theme.palette.text.secondary,
@@ -29,13 +13,9 @@ export const StatusDot = styled(Box, {
   shouldForwardProp: prop => prop !== 'status',
 })<{ status: 'loaded' | 'unloaded' | 'loading' | 'noIndex' }>(
   ({ theme, status }) => ({
-    width: '8px',
-    height: '8px',
+    width: '6px',
+    height: '6px',
     borderRadius: '50%',
-    position: 'absolute',
-    right: 6,
-    top: 10,
-    zIndex: 1,
     ...(status === 'loaded' && {
       border: `1px solid ${theme.palette.primary.main}`,
       backgroundColor: theme.palette.primary.main,
@@ -63,3 +43,123 @@ export const StyledMenuItem = styled(MenuItem)(({ theme }) => ({
 export const StyledDivider = styled(Divider)(({ theme }) => ({
   margin: 0,
 }));
+
+// New styled components for tree view
+export const TreeContainer = styled(Box)(({ theme }) => ({
+  height: 'calc(100vh - 64px)',
+  overflow: 'auto',
+  fontSize: '15px',
+  color: theme.palette.text.primary,
+  '& .MuiSvgIcon-root': {
+    fontSize: '14px',
+    color: theme.palette.text.primary,
+  },
+  '&::-webkit-scrollbar': {
+    width: '6px',
+    height: '6px',
+  },
+  '&::-webkit-scrollbar-track': {
+    background: 'transparent',
+  },
+  '&::-webkit-scrollbar-thumb': {
+    background:
+      theme.palette.mode === 'dark'
+        ? 'rgba(255, 255, 255, 0.15)'
+        : 'rgba(0, 0, 0, 0.15)',
+    borderRadius: '3px',
+    '&:hover': {
+      background:
+        theme.palette.mode === 'dark'
+          ? 'rgba(255, 255, 255, 0.25)'
+          : 'rgba(0, 0, 0, 0.25)',
+    },
+  },
+  '& > div': {
+    width: '100%',
+  },
+}));
+
+export const TreeContent = styled(Box)({
+  height: '100%',
+  width: '100%',
+  position: 'relative',
+  overflow: 'hidden',
+});
+
+export const NoResultsBox = styled(Box)(({ theme }) => ({
+  display: 'flex',
+  justifyContent: 'center',
+  alignItems: 'center',
+  height: '100px',
+  color: theme.palette.text.secondary,
+}));
+
+export const TreeNodeBox = styled(Box, {
+  shouldForwardProp: prop =>
+    ![
+      'isSelected',
+      'isContextMenuTarget',
+      'isCollectionNoSchema',
+      'depth',
+    ].includes(prop as string),
+})<{
+  isSelected: boolean;
+  isContextMenuTarget: boolean;
+  isCollectionNoSchema: boolean;
+  depth: number;
+}>(({ theme, isSelected, isCollectionNoSchema, depth }) => ({
+  position: 'absolute',
+  top: 0,
+  left: 0,
+  width: '100%',
+  height: '30px',
+  transform: 'translateY(0)',
+  display: 'flex',
+  alignItems: 'center',
+  cursor: 'pointer',
+  paddingLeft: `${depth * 16 === 0 ? 8 : depth * 12}px`,
+  paddingRight: '8px',
+  backgroundColor: isSelected
+    ? theme.palette.primary.light
+    : theme.palette.background.default,
+  '&:hover': {
+    backgroundColor: theme.palette.primary.light,
+  },
+  opacity: isCollectionNoSchema ? 0.5 : 1,
+  color: isCollectionNoSchema ? theme.palette.text.disabled : 'inherit',
+  pointerEvents: isCollectionNoSchema ? 'none' : 'auto',
+  marginBottom: depth === 0 ? '0' : '0',
+  boxSizing: 'border-box',
+}));
+
+export const SearchBoxContainer = styled(Box)({
+  position: 'absolute',
+  top: 0,
+  left: 0,
+  width: '100%',
+  height: '30px',
+  transform: 'translateY(0)',
+  boxSizing: 'border-box',
+});
+
+export const NodeContent = styled(Box)({
+  display: 'flex',
+  alignItems: 'center',
+  flex: 1,
+  minWidth: 0,
+  paddingRight: '4px',
+});
+
+export const SearchButton = styled(Box)(({ theme }) => ({
+  padding: '4px',
+  marginLeft: theme.spacing(1),
+  '&:hover': {
+    backgroundColor: 'transparent',
+  },
+}));
+
+export const CollectionCount = styled('span')({
+  marginLeft: 8,
+  color: '#888',
+  fontSize: 12,
+});

client/src/pages/databases/tree/index.tsx

@@ -7,111 +7,84 @@ import React, {
   useContext,
 } from 'react';
 import { useTranslation } from 'react-i18next';
-import { useParams } from 'react-router-dom';
+import { useParams, useNavigate } from 'react-router-dom';
 import icons from '@/components/icons/Icons';
-import {
-  Tooltip,
-  Typography,
-  Grow,
-  Popover,
-  Box,
-  IconButton,
-} from '@mui/material'; // Added Box, IconButton
-import { useNavigate } from 'react-router-dom';
+import { Grow, Popover, Box, IconButton, Fade } from '@mui/material';
 import { CollectionObject } from '@server/types';
-import { formatNumber } from '@/utils';
 import {
-  DatabaseTreeItem as OriginalDatabaseTreeItem, // Rename original type
+  DatabaseTreeItem as OriginalDatabaseTreeItem,
   TreeNodeType,
   DatabaseTreeProps,
   ContextMenu,
   TreeNodeObject,
 } from './types';
 import { TreeContextMenu } from './TreeContextMenu';
-import { useVirtualizer } from '@tanstack/react-virtual'; // Import virtualizer
+import { useVirtualizer } from '@tanstack/react-virtual';
 import { dataContext } from '@/context';
+import CollectionNode from './CollectionNode';
+import DatabaseNode from './DatabaseNode';
+import SearchBox from './SearchBox';
 import {
-  CollectionNodeWrapper,
-  CollectionNameWrapper,
-  Count,
-  StatusDot,
+  TreeContainer,
+  TreeContent,
+  NoResultsBox,
+  TreeNodeBox,
+  SearchBoxContainer,
 } from './StyledComponents';
 
-// Define a type for the flattened list item
 interface FlatTreeItem {
   id: string;
   name: string;
   depth: number;
   type: TreeNodeType;
-  data: TreeNodeObject | null; // Allow null for DB node
-  isExpanded: boolean; // Track if the node itself is expanded
-  hasChildren: boolean; // Does it have children?
-  originalNode: OriginalDatabaseTreeItem; // Keep original node ref if needed
+  data: TreeNodeObject | null;
+  isExpanded: boolean;
+  hasChildren: boolean;
+  originalNode: OriginalDatabaseTreeItem;
 }
 
-const CollectionNode: React.FC<{ data: CollectionObject }> = ({ data }) => {
-  const { t: commonTrans } = useTranslation();
-  const hasIndex = data.schema && data.schema.hasVectorIndex;
-  const loadStatus = hasIndex
-    ? data.loaded
-      ? commonTrans('status.loaded')
-      : commonTrans('status.unloaded')
-    : commonTrans('status.noVectorIndex');
-
-  const getStatus = () => {
-    if (!data.schema || !data.schema.hasVectorIndex) return 'noIndex';
-    if (data.status === 'loading') return 'loading';
-    return data.loaded ? 'loaded' : 'unloaded';
-  };
-
-  return (
-    <CollectionNodeWrapper>
-      <CollectionNameWrapper>
-        <Tooltip title={data.collection_name} placement="top">
-          <Typography noWrap className="collectionName">
-            {data.collection_name}
-          </Typography>
-        </Tooltip>
-        <Count>({formatNumber(data.rowCount || 0)})</Count>
-      </CollectionNameWrapper>
-      <Tooltip title={loadStatus} placement="top">
-        <StatusDot status={getStatus()} />
-      </Tooltip>
-    </CollectionNodeWrapper>
-  );
-};
-
 const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
-  const {
-    database,
-    collections,
-    params,
-    treeHeight = 'calc(100vh - 64px)',
-  } = props;
-
-  // context
+  const { database, collections, params } = props;
+
   const navigate = useNavigate();
   const { collectionName = '' } = useParams<{ collectionName: string }>();
   const { batchRefreshCollections } = useContext(dataContext);
+  const { t } = useTranslation();
 
-  // State
   const [contextMenu, setContextMenu] = useState<ContextMenu | null>(null);
   const [expandedItems, setExpandedItems] = useState<Set<string>>(
     new Set([database])
-  ); // Start with DB expanded
+  );
   const [selectedItemId, setSelectedItemId] = useState<string | null>(
     params.collectionName ? `c_${params.collectionName}` : database
   );
+  const [searchQuery, setSearchQuery] = useState<string>('');
+  const [debouncedSearchQuery, setDebouncedSearchQuery] = useState<string>('');
+  const [isSearchOpen, setIsSearchOpen] = useState(false);
+  const [showScrollTop, setShowScrollTop] = useState(false);
 
-  // Ref for the scrollable element
   const parentRef = useRef<HTMLDivElement>(null);
 
-  // Icons
-  const DatabaseIcon = icons.database;
-  const CollectionIcon = icons.navCollection;
   const ExpandIcon = icons.rightArrow;
 
-  // --- Tree Flattening Logic ---
+  useEffect(() => {
+    const timer = setTimeout(() => {
+      setDebouncedSearchQuery(searchQuery);
+    }, 300);
+
+    return () => {
+      clearTimeout(timer);
+    };
+  }, [searchQuery]);
+
+  const filteredCollections = useMemo(() => {
+    if (!debouncedSearchQuery) return collections;
+    const query = debouncedSearchQuery.toLowerCase();
+    return collections.filter(c =>
+      c.collection_name.toLowerCase().includes(query)
+    );
+  }, [collections, debouncedSearchQuery]);
+
   const flattenTree = useCallback(
     (
       node: OriginalDatabaseTreeItem,
@@ -144,39 +117,39 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
     []
   );
 
-  // --- Memoized Flattened Data ---
   const flattenedNodes = useMemo(() => {
-    // Rebuild the tree structure in the format needed for flattenTree
-    const children = collections.map(c => ({
+    const children = filteredCollections.map(c => ({
       id: `c_${c.collection_name}`,
       name: c.collection_name,
       type: 'collection' as TreeNodeType,
       data: c,
-      children: [], // Collections don't have children in this structure
-      expanded: false, // Not relevant for leaf nodes here
+      children: [],
+      expanded: false,
     }));
 
     const tree: OriginalDatabaseTreeItem = {
       id: database,
       name: database,
-      expanded: expandedItems.has(database), // Use state here
+      expanded: expandedItems.has(database),
       type: 'db',
       children: children,
-      data: undefined, // DB node has no specific data object here
+      data: undefined,
     };
 
-    return flattenTree(tree, 0, expandedItems);
-  }, [database, collections, expandedItems, flattenTree]); // Dependencies
+    const allNodes = flattenTree(tree, 0, expandedItems);
+
+    return allNodes.filter(
+      node => node.type !== 'db' && node.type !== 'search'
+    );
+  }, [database, filteredCollections, expandedItems, flattenTree, isSearchOpen]);
 
-  // --- Virtualizer Instance ---
   const rowVirtualizer = useVirtualizer({
     count: flattenedNodes.length,
     getScrollElement: () => parentRef.current,
-    estimateSize: () => 30, // Adjust this based on your average item height
-    overscan: 5, // Render items slightly outside the viewport
+    estimateSize: () => 30,
+    overscan: 5,
   });
 
-  // --- Event Handlers ---
   const handleToggleExpand = (nodeId: string) => {
     setExpandedItems(prev => {
       const newSet = new Set(prev);
@@ -187,15 +160,12 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
       }
       return newSet;
     });
-    // Close context menu on expand/collapse
     setContextMenu(null);
   };
 
-  // Scroll to the selected item
   const skipNextScrollRef = useRef(false);
 
   const handleNodeClick = (node: FlatTreeItem) => {
-    // flag to skip the next scroll
     skipNextScrollRef.current = true;
 
     setSelectedItemId(node.id);
@@ -209,10 +179,7 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
     setContextMenu(null);
   };
 
-  const handleContextMenu = (
-    event: React.MouseEvent, // Use React.MouseEvent
-    node: FlatTreeItem
-  ) => {
+  const handleContextMenu = (event: React.MouseEvent, node: FlatTreeItem) => {
     event.preventDefault();
     event.stopPropagation();
     setContextMenu({
@@ -220,7 +187,7 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
       mouseY: event.clientY - 4,
       nodeId: node.id,
       nodeType: node.type,
-      object: node.data, // Pass the data associated with the flat node
+      object: node.data,
     });
   };
 
@@ -228,7 +195,6 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
     setContextMenu(null);
   };
 
-  // Effect for closing context menu on outside click
   useEffect(() => {
     document.addEventListener('click', handleCloseContextMenu);
     return () => {
@@ -236,19 +202,15 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
     };
   }, []);
 
-  // Track visible items for refreshing
   useEffect(() => {
     if (!collections.length) return;
 
-    // Track whether we're currently scrolling
     let isScrolling = false;
     let scrollTimeoutId: NodeJS.Timeout | null = null;
 
-    // Save references to stable values to avoid dependency changes
     const currentFlattenedNodes = flattenedNodes;
 
     const refreshVisibleCollections = () => {
-      // Early return if the component is unmounted
       if (!parentRef.current) return;
 
       const visibleItems = rowVirtualizer.getVirtualItems();
@@ -268,41 +230,33 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
       }
     };
 
-    // This will be called when scrolling starts
     const handleScrollStart = () => {
       if (!isScrolling) {
         isScrolling = true;
-        // Execute on scroll start
         refreshVisibleCollections();
       }
 
-      // Clear any existing timeout
       if (scrollTimeoutId) {
         clearTimeout(scrollTimeoutId);
       }
 
-      // Set a new timeout for scroll end detection
       scrollTimeoutId = setTimeout(() => {
         isScrolling = false;
-        // Execute on scroll end
         refreshVisibleCollections();
         scrollTimeoutId = null;
-      }, 500); // Wait for scrolling to stop for 300ms
+      }, 500);
     };
 
-    // Initial refresh when component mounts - with delay to ensure UI is ready
     const initialRefreshTimeout = setTimeout(() => {
       refreshVisibleCollections();
     }, 100);
 
-    // Setup scroll listener
     const scrollElement = parentRef.current;
     if (scrollElement) {
       scrollElement.addEventListener('scroll', handleScrollStart);
 
       return () => {
         scrollElement.removeEventListener('scroll', handleScrollStart);
-        // Clear timeout on cleanup
         if (scrollTimeoutId) {
           clearTimeout(scrollTimeoutId);
         }
@@ -316,7 +270,12 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
       }
       clearTimeout(initialRefreshTimeout);
     };
-  }, [collections.length, batchRefreshCollections, rowVirtualizer]);
+  }, [
+    collections.length,
+    batchRefreshCollections,
+    rowVirtualizer,
+    debouncedSearchQuery,
+  ]);
 
   useEffect(() => {
     if (skipNextScrollRef.current) {
@@ -339,123 +298,198 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
     }
   }, [collectionName]);
 
-  // --- Rendering ---
+  // Add scroll handler for scroll-to-top button
+  useEffect(() => {
+    const handleScroll = () => {
+      if (!parentRef.current) return;
+      const scrollTop = parentRef.current.scrollTop;
+      setShowScrollTop(scrollTop > 200); // Show button when scrolled more than 200px
+    };
+
+    const scrollElement = parentRef.current;
+    if (scrollElement) {
+      scrollElement.addEventListener('scroll', handleScroll);
+      return () => scrollElement.removeEventListener('scroll', handleScroll);
+    }
+  }, []);
+
+  const handleScrollToTop = () => {
+    if (parentRef.current) {
+      parentRef.current.scrollTo({
+        top: 0,
+        // behavior: 'smooth',
+      });
+    }
+  };
+
   return (
     <>
-      <Box
-        ref={parentRef}
-        sx={{
-          height: treeHeight,
-          overflow: 'auto',
-          fontSize: '15px',
-          color: theme => theme.palette.text.primary,
-          '& .MuiSvgIcon-root': {
-            fontSize: '14px',
-            color: theme => theme.palette.text.primary,
-          },
-        }}
-      >
-        <Box
+      <TreeContainer ref={parentRef}>
+        {isSearchOpen ? (
+          <SearchBoxContainer
+            sx={{
+              position: 'sticky',
+              top: 0,
+              zIndex: 2,
+              backgroundColor: 'inherit',
+            }}
+          >
+            <SearchBox
+              searchQuery={searchQuery}
+              onSearchChange={setSearchQuery}
+              onClose={() => {
+                setSearchQuery('');
+                setIsSearchOpen(false);
+              }}
+            />
+          </SearchBoxContainer>
+        ) : (
+          <DatabaseNode
+            database={database}
+            collectionCount={collections.length}
+            isSelected={selectedItemId === database}
+            isContextMenuTarget={contextMenu?.nodeId === database}
+            onNodeClick={() =>
+              handleNodeClick({
+                id: database,
+                name: database,
+                depth: 0,
+                type: 'db',
+                data: null,
+                isExpanded: false,
+                hasChildren: false,
+                originalNode: {
+                  id: database,
+                  name: database,
+                  type: 'db',
+                  children: [],
+                  expanded: false,
+                },
+              })
+            }
+            onContextMenu={e =>
+              handleContextMenu(e, {
+                id: database,
+                name: database,
+                depth: 0,
+                type: 'db',
+                data: null,
+                isExpanded: false,
+                hasChildren: false,
+                originalNode: {
+                  id: database,
+                  name: database,
+                  type: 'db',
+                  children: [],
+                  expanded: false,
+                },
+              })
+            }
+            onSearchClick={e => {
+              e.stopPropagation();
+              setIsSearchOpen(!isSearchOpen);
+            }}
+          />
+        )}
+
+        <TreeContent
           sx={{
             height: `${rowVirtualizer.getTotalSize()}px`,
-            width: '100%',
-            position: 'relative',
-            overflow: 'hidden',
+            marginTop: 0,
           }}
         >
-          {rowVirtualizer.getVirtualItems().map(virtualItem => {
-            const node = flattenedNodes[virtualItem.index];
-            if (!node) return null; // Should not happen
-
-            const isSelected = node.id === selectedItemId;
-            const isContextMenuTarget = contextMenu?.nodeId === node.id;
-            const isCollectionNoSchema =
-              node.type === 'collection' &&
-              (!node.data || !(node.data as CollectionObject).schema);
-
-            return (
-              <Box
-                key={node.id}
-                sx={{
-                  position: 'absolute',
-                  top: 0,
-                  left: 0,
-                  width: `calc(100% - ${node.depth * 20}px)`, // Adjust for padding
-                  height: `${virtualItem.size}px`,
-                  transform: `translateY(${virtualItem.start}px)`,
-                  display: 'flex',
-                  flex: 1,
-                  alignItems: 'center',
-                  cursor: 'pointer',
-                  paddingLeft: `${node.depth * 20}px`, // Indent based on depth
-                  // Apply selection/hover styles based on the original classes
-                  backgroundColor: isSelected
-                    ? 'rgba(10, 206, 130, 0.28)'
-                    : isContextMenuTarget
-                      ? 'rgba(10, 206, 130, 0.08)'
-                      : 'transparent',
-                  '&:hover': {
-                    backgroundColor: isSelected
-                      ? 'rgba(10, 206, 130, 0.28)'
-                      : 'rgba(10, 206, 130, 0.08)',
-                  },
-                  opacity: isCollectionNoSchema ? 0.5 : 1,
-                  color: isCollectionNoSchema ? 'text.disabled' : 'inherit',
-                  pointerEvents: isCollectionNoSchema ? 'none' : 'auto',
-                }}
-                onClick={() => handleNodeClick(node)}
-                onContextMenu={e => handleContextMenu(e, node)}
-              >
-                {/* Expand/Collapse Icon */}
-                {node.hasChildren ? (
-                  <IconButton
-                    size="small"
-                    onClick={e => {
-                      e.stopPropagation();
-                      handleToggleExpand(node.id);
-                    }}
-                    sx={{ mr: 0 }}
-                  >
-                    {node.isExpanded ? (
-                      <ExpandIcon sx={{ transform: `rotate(90deg)` }} />
-                    ) : (
-                      <ExpandIcon />
-                    )}
-                  </IconButton>
-                ) : (
-                  // Placeholder for alignment if needed
-                  <Box sx={{ width: 0, mr: 0.5 }} /> // Adjust width to match IconButton
-                )}
-
-                {/* Node Type Icon */}
-                <Box sx={{ mr: 1, display: 'flex', alignItems: 'center' }}>
-                  {node.type === 'db' ? <DatabaseIcon /> : <CollectionIcon />}
-                </Box>
-
-                {/* Label */}
-                {node.type === 'collection' && node.data ? (
-                  <CollectionNode data={node.data as CollectionObject} />
-                ) : (
-                  <Tooltip title={node.name} placement="top">
-                    <Typography noWrap sx={{ width: 'calc(100% - 30px)' }}>
-                      {/* Reuse dbName style or create a generic one */}
-                      {node.name}
-                      {node.type === 'db' && (
-                        <span
-                          style={{ marginLeft: 8, color: '#888', fontSize: 12 }}
-                        >
-                          ({collections.length})
-                        </span>
+          {flattenedNodes.length === 0 ? (
+            <NoResultsBox>{t('search.noResults')}</NoResultsBox>
+          ) : (
+            rowVirtualizer.getVirtualItems().map(virtualItem => {
+              const node = flattenedNodes[virtualItem.index];
+              if (!node) return null;
+
+              const isSelected = node.id === selectedItemId;
+              const isContextMenuTarget = contextMenu?.nodeId === node.id;
+              const isCollectionNoSchema =
+                node.type === 'collection' &&
+                (!node.data || !(node.data as CollectionObject).schema);
+
+              return (
+                <TreeNodeBox
+                  key={node.id}
+                  isSelected={isSelected}
+                  isContextMenuTarget={isContextMenuTarget}
+                  isCollectionNoSchema={isCollectionNoSchema}
+                  depth={node.depth}
+                  sx={{
+                    transform: `translateY(${virtualItem.start}px)`,
+                  }}
+                  onClick={() => handleNodeClick(node)}
+                  onContextMenu={e => handleContextMenu(e, node)}
+                >
+                  {node.hasChildren && node.type !== 'db' ? (
+                    <IconButton
+                      size="small"
+                      onClick={e => {
+                        e.stopPropagation();
+                        handleToggleExpand(node.id);
+                      }}
+                      sx={{ mr: 0 }}
+                    >
+                      {node.isExpanded ? (
+                        <ExpandIcon sx={{ transform: `rotate(90deg)` }} />
+                      ) : (
+                        <ExpandIcon />
                       )}
-                    </Typography>
-                  </Tooltip>
-                )}
-              </Box>
-            );
-          })}
-        </Box>
-      </Box>
-      {/* Context Menu Popover (keep existing) */}
+                    </IconButton>
+                  ) : (
+                    <Box sx={{ width: 0 }} />
+                  )}
+
+                  {node.type === 'collection' && node.data && (
+                    <CollectionNode
+                      data={node.data as CollectionObject}
+                      highlight={debouncedSearchQuery}
+                      isSelected={isSelected}
+                      isContextMenuTarget={isContextMenuTarget}
+                      hasChildren={node.hasChildren}
+                      isExpanded={node.isExpanded}
+                      depth={node.depth}
+                      onToggleExpand={e => {
+                        e.stopPropagation();
+                        handleToggleExpand(node.id);
+                      }}
+                      onClick={() => handleNodeClick(node)}
+                      onContextMenu={e => handleContextMenu(e, node)}
+                    />
+                  )}
+                </TreeNodeBox>
+              );
+            })
+          )}
+        </TreeContent>
+      </TreeContainer>
+
+      <Fade in={showScrollTop}>
+        <IconButton
+          onClick={handleScrollToTop}
+          size="small"
+          sx={{
+            position: 'absolute',
+            bottom: 16,
+            right: 16,
+            backgroundColor: 'rgba(255, 255, 255, 0.9)',
+            '&:hover': {
+              backgroundColor: 'rgba(255, 255, 255, 1)',
+            },
+            boxShadow: '0 2px 8px rgba(0,0,0,0.15)',
+            zIndex: 1000,
+            width: 32,
+            height: 32,
+            border: '1px solid rgba(0,0,0,0.1)',
+          }}
+        >
+          <ExpandIcon sx={{ transform: 'rotate(-90deg)', fontSize: 20 }} />
+        </IconButton>
+      </Fade>
+
       <Popover
         open={Boolean(contextMenu)}
         onClose={handleCloseContextMenu}
@@ -474,7 +508,6 @@ const DatabaseTree: React.FC<DatabaseTreeProps> = props => {
           },
         }}
       >
-        {/* Pass the correct contextMenu object */}
         <TreeContextMenu
           onClick={handleCloseContextMenu}
           contextMenu={contextMenu!}

client/src/pages/databases/tree/types.ts

@@ -1,7 +1,7 @@
 import { Params } from 'react-router-dom';
 import type { CollectionObject, DatabaseObject } from '@server/types';
 
-export type TreeNodeType = 'db' | 'collection' | 'partition' | 'segment';
+export type TreeNodeType = 'db' | 'collection' | 'partition' | 'segment' | 'search';
 export type TreeNodeObject = CollectionObject | DatabaseObject | null;
 
 export interface DatabaseTreeItem {

client/src/styles/theme.ts

@@ -28,6 +28,10 @@ declare module '@mui/material/styles' {
       800: string;
       900: string;
     };
+    highlight: {
+      light: string;
+      dark: string;
+    };
   }
 }
 
@@ -120,6 +124,10 @@ const colors = {
     disabled: '#9CA3AF',
   },
   divider: '#E5E7EB',
+  highlight: {
+    light: '#FFE082',
+    dark: '#003A8C',
+  },
 };
 
 const spacing = (factor: number) => `${8 * factor}px`;
@@ -262,6 +270,10 @@ const getCommonThemes = (mode: PaletteMode) => ({
       disabled: mode === 'light' ? colors.text.disabled : colors.neutral[600],
     },
     divider: mode === 'light' ? colors.divider : colors.neutral[700],
+    highlight: {
+      light: colors.highlight.light,
+      dark: colors.highlight.dark,
+    },
   },
   spacing,
 });