@@ -41,6 +41,11 @@ import {
4141import type { SqlLabRootState } from 'src/SqlLab/types' ;
4242import useQueryEditor from 'src/SqlLab/hooks/useQueryEditor' ;
4343import { addTable , removeTables } from 'src/SqlLab/actions/sqlLab' ;
44+ import {
45+ getItem ,
46+ setItem ,
47+ LocalStorageKeys ,
48+ } from 'src/utils/localStorageHelpers' ;
4449import PanelToolbar from 'src/components/PanelToolbar' ;
4550import { ViewLocations } from 'src/SqlLab/contributions' ;
4651import TreeNodeRenderer from './TreeNodeRenderer' ;
@@ -126,6 +131,36 @@ const StyledTreeContainer = styled.div`
126131
127132const ROW_HEIGHT = 28 ;
128133
134+ const getPinnedSchemasStorageKey = (
135+ dbId : number | undefined ,
136+ catalog : string | null | undefined ,
137+ ) : string => `${ dbId ?? '' } ${ catalog ?? '' } ;
138+
139+ const getPinnedSchemasFromStorage = (
140+ dbId : number | undefined ,
141+ catalog : string | null | undefined ,
142+ ) : Set < string > => {
143+ if ( ! dbId ) return new Set ( ) ;
144+ const stored = getItem ( LocalStorageKeys . SqllabPinnedSchemas , { } ) ;
145+ const key = getPinnedSchemasStorageKey ( dbId , catalog ) ;
146+ const schemas = stored [ key ] ;
147+ return Array . isArray ( schemas ) ? new Set < string > ( schemas ) : new Set ( ) ;
148+ } ;
149+
150+ const savePinnedSchemasToStorage = (
151+ dbId : number | undefined ,
152+ catalog : string | null | undefined ,
153+ schemas : Set < string > ,
154+ ) => {
155+ if ( ! dbId ) return ;
156+ const stored = getItem ( LocalStorageKeys . SqllabPinnedSchemas , { } ) ;
157+ const key = getPinnedSchemasStorageKey ( dbId , catalog ) ;
158+ setItem ( LocalStorageKeys . SqllabPinnedSchemas , {
159+ ...stored ,
160+ [ key ] : [ ...schemas ] ,
161+ } ) ;
162+ } ;
163+
129164const TableExploreTree : React . FC < Props > = ( { queryEditorId } ) => {
130165 const dispatch = useDispatch ( ) ;
131166 const theme = useTheme ( ) ;
@@ -161,6 +196,7 @@ const TableExploreTree: React.FC<Props> = ({ queryEditorId }) => {
161196 selectStarMap,
162197 handleToggle,
163198 handleRefreshTables,
199+ refreshTableSchema,
164200 errorPayload,
165201 } = useTreeData ( {
166202 dbId,
@@ -199,6 +235,83 @@ const TableExploreTree: React.FC<Props> = ({ queryEditorId }) => {
199235 } ,
200236 [ dispatch , tables , editorId , dbId ] ,
201237 ) ;
238+ const [ pinnedSchemas , setPinnedSchemas ] = useState < Set < string > > ( ( ) =>
239+ getPinnedSchemasFromStorage ( dbId , catalog ) ,
240+ ) ;
241+
242+ const previousDbIdRef = useRef < number | undefined > ( dbId ) ;
243+ const previousCatalogRef = useRef < string | null | undefined > ( catalog ) ;
244+
245+ // Single effect handles both loading and persisting pinned schemas.
246+ // Using refs to detect source changes avoids the race condition where the
247+ // persist branch would run with stale pinnedSchemas right after a dbId/catalog
248+ // change, corrupting the new source's stored pins.
249+ useEffect ( ( ) => {
250+ const dbChanged = previousDbIdRef . current !== dbId ;
251+ const catalogChanged = previousCatalogRef . current !== catalog ;
252+
253+ if ( dbChanged || catalogChanged ) {
254+ previousDbIdRef . current = dbId ;
255+ previousCatalogRef . current = catalog ;
256+ setPinnedSchemas ( getPinnedSchemasFromStorage ( dbId , catalog ) ) ;
257+ return ;
258+ }
259+
260+ savePinnedSchemasToStorage ( dbId , catalog , pinnedSchemas ) ;
261+ } , [ dbId , catalog , pinnedSchemas ] ) ;
262+
263+ const handlePinSchema = useCallback ( ( schemaName : string ) => {
264+ setPinnedSchemas ( prev => new Set ( [ ...prev , schemaName ] ) ) ;
265+ } , [ ] ) ;
266+
267+ const handleUnpinSchema = useCallback ( ( schemaName : string ) => {
268+ setPinnedSchemas ( prev => {
269+ const next = new Set ( prev ) ;
270+ next . delete ( schemaName ) ;
271+ return next ;
272+ } ) ;
273+ } , [ ] ) ;
274+
275+ const sortedTreeData = useMemo ( ( ) => {
276+ if ( pinnedSchemas . size === 0 ) return treeData ;
277+ const pinned = treeData . filter ( node => pinnedSchemas . has ( node . name ) ) ;
278+ const rest = treeData . filter ( node => ! pinnedSchemas . has ( node . name ) ) ;
279+ return [ ...pinned , ...rest ] ;
280+ } , [ treeData , pinnedSchemas ] ) ;
281+
282+ const [ sortedTables , setSortedTables ] = useState < Record < string , boolean > > ( { } ) ;
283+
284+ useEffect ( ( ) => {
285+ setSortedTables ( { } ) ;
286+ } , [ dbId , catalog ] ) ;
287+
288+ const toggleSortColumns = useCallback ( ( tableId : string ) => {
289+ setSortedTables ( prev => ( { ...prev , [ tableId ] : ! prev [ tableId ] } ) ) ;
290+ } , [ ] ) ;
291+
292+ const displayTreeData = useMemo ( ( ) => {
293+ const activeSorted = Object . keys ( sortedTables ) . filter (
294+ id => sortedTables [ id ] ,
295+ ) ;
296+ if ( activeSorted . length === 0 ) return sortedTreeData ;
297+
298+ const sortedSet = new Set ( activeSorted ) ;
299+ return sortedTreeData . map ( schemaNode => ( {
300+ ...schemaNode ,
301+ children : schemaNode . children ?. map ( tableNode => {
302+ if ( tableNode . type !== 'table' || ! sortedSet . has ( tableNode . id ) ) {
303+ return tableNode ;
304+ }
305+ const { children } = tableNode ;
306+ if ( ! children || children . length <= 1 ) return tableNode ;
307+ return {
308+ ...tableNode ,
309+ children : [ ...children ] . sort ( ( a , b ) => a . name . localeCompare ( b . name ) ) ,
310+ } ;
311+ } ) ,
312+ } ) ) ;
313+ } , [ sortedTreeData , sortedTables ] ) ;
314+
202315 const [ searchTerm , setSearchTerm ] = useState ( '' ) ;
203316 const handleSearchChange = useCallback (
204317 ( { target } : ChangeEvent < HTMLInputElement > ) => setSearchTerm ( target . value ) ,
@@ -270,8 +383,8 @@ const TableExploreTree: React.FC<Props> = ({ queryEditorId }) => {
270383 return false ;
271384 } ;
272385
273- return treeData . some ( node => checkNode ( node ) ) ;
274- } , [ searchTerm , treeData ] ) ;
386+ return displayTreeData . some ( node => checkNode ( node ) ) ;
387+ } , [ searchTerm , displayTreeData ] ) ;
275388
276389 // Node renderer for react-arborist
277390 const renderNode = useCallback (
@@ -283,19 +396,31 @@ const TableExploreTree: React.FC<Props> = ({ queryEditorId }) => {
283396 searchTerm = { searchTerm }
284397 catalog = { catalog }
285398 pinnedTableKeys = { pinnedTableKeys }
399+ pinnedSchemas = { pinnedSchemas }
286400 selectStarMap = { selectStarMap }
287401 handleRefreshTables = { handleRefreshTables }
288402 handlePinTable = { handlePinTable }
289403 handleUnpinTable = { handleUnpinTable }
404+ handlePinSchema = { handlePinSchema }
405+ handleUnpinSchema = { handleUnpinSchema }
406+ refreshTableSchema = { refreshTableSchema }
407+ sortedTables = { sortedTables }
408+ toggleSortColumns = { toggleSortColumns }
290409 />
291410 ) ,
292411 [
293412 catalog ,
294413 pinnedTableKeys ,
414+ pinnedSchemas ,
295415 selectStarMap ,
296416 handleRefreshTables ,
297417 handlePinTable ,
298418 handleUnpinTable ,
419+ handlePinSchema ,
420+ handleUnpinSchema ,
421+ refreshTableSchema ,
422+ sortedTables ,
423+ toggleSortColumns ,
299424 loadingNodes ,
300425 manuallyOpenedNodes ,
301426 searchTerm ,
@@ -369,7 +494,7 @@ const TableExploreTree: React.FC<Props> = ({ queryEditorId }) => {
369494 return (
370495 < Tree < TreeNodeData >
371496 ref = { treeRef }
372- data = { treeData }
497+ data = { displayTreeData }
373498 width = "100%"
374499 height = { height || 500 }
375500 rowHeight = { ROW_HEIGHT }
0 commit comments