feat: persist column widths and order per table across sessions (#413)

datlechin · web-flow · commit fcf0115e408c · 2026-03-22T02:44:42.000+07:00
* feat: persist column widths and order per table across sessions

* fix: address code review issues for column layout persistence
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Show/hide row numbers column in data grid (Settings > Data Grid)
+- Persist column widths and order per table across tab switches, view toggles, and app restarts
 
 ## [0.22.0] - 2026-03-21
 
diff --git a/TablePro/Core/Storage/ColumnLayoutStorage.swift b/TablePro/Core/Storage/ColumnLayoutStorage.swift
@@ -0,0 +1,59 @@
+//
+//  ColumnLayoutStorage.swift
+//  TablePro
+//
+
+import Foundation
+
+@MainActor
+internal final class ColumnLayoutStorage {
+    static let shared = ColumnLayoutStorage()
+
+    private init() {}
+
+    // MARK: - Types
+
+    private struct PersistedColumnLayout: Codable {
+        var columnWidths: [String: CGFloat]
+        var columnOrder: [String]?
+    }
+
+    // MARK: - Public API
+
+    func save(_ layout: ColumnLayoutState, for tableName: String, connectionId: UUID) {
+        guard !layout.columnWidths.isEmpty else { return }
+
+        let persisted = PersistedColumnLayout(
+            columnWidths: layout.columnWidths,
+            columnOrder: layout.columnOrder
+        )
+        let key = Self.userDefaultsKey(tableName: tableName, connectionId: connectionId)
+        if let data = try? JSONEncoder().encode(persisted) {
+            UserDefaults.standard.set(data, forKey: key)
+        }
+    }
+
+    func load(for tableName: String, connectionId: UUID) -> ColumnLayoutState? {
+        let key = Self.userDefaultsKey(tableName: tableName, connectionId: connectionId)
+        guard let data = UserDefaults.standard.data(forKey: key),
+              let persisted = try? JSONDecoder().decode(PersistedColumnLayout.self, from: data)
+        else {
+            return nil
+        }
+        var state = ColumnLayoutState()
+        state.columnWidths = persisted.columnWidths
+        state.columnOrder = persisted.columnOrder
+        return state
+    }
+
+    func clear(for tableName: String, connectionId: UUID) {
+        let key = Self.userDefaultsKey(tableName: tableName, connectionId: connectionId)
+        UserDefaults.standard.removeObject(forKey: key)
+    }
+
+    // MARK: - Private
+
+    private static func userDefaultsKey(tableName: String, connectionId: UUID) -> String {
+        "com.TablePro.columns.layout.\(connectionId.uuidString).\(tableName)"
+    }
+}
diff --git a/TablePro/Views/Main/Child/MainEditorContentView.swift b/TablePro/Views/Main/Child/MainEditorContentView.swift
@@ -329,6 +329,7 @@ struct MainEditorContentView: View {
             databaseType: connection.type,
             tableName: tab.tableName,
             primaryKeyColumn: changeManager.primaryKeyColumn,
+            tabType: tab.tabType,
             showRowNumbers: AppSettingsManager.shared.dataGrid.showRowNumbers,
             hiddenColumns: columnVisibilityManager.hiddenColumns,
             onHideColumn: { [coordinator] columnName in
@@ -473,6 +474,7 @@ struct MainEditorContentView: View {
                 }
                 DispatchQueue.main.async {
                     coordinator.isUpdatingColumnLayout = false
+                    coordinator.saveColumnLayoutForTable()
                 }
             }
         )
diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+ColumnLayout.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+ColumnLayout.swift
@@ -0,0 +1,27 @@
+//
+//  MainContentCoordinator+ColumnLayout.swift
+//  TablePro
+//
+
+import Foundation
+
+extension MainContentCoordinator {
+    func saveColumnLayoutForTable() {
+        guard let index = tabManager.selectedTabIndex else { return }
+        let tab = tabManager.tabs[index]
+        guard tab.tabType == .table, let tableName = tab.tableName, !tableName.isEmpty else { return }
+
+        ColumnLayoutStorage.shared.save(tab.columnLayout, for: tableName, connectionId: connectionId)
+        columnVisibilityManager.saveLastHiddenColumns(for: tableName, connectionId: connectionId)
+    }
+
+    func restoreColumnLayoutForTable(_ tableName: String) {
+        guard let index = tabManager.selectedTabIndex else { return }
+
+        if let savedLayout = ColumnLayoutStorage.shared.load(for: tableName, connectionId: connectionId) {
+            tabManager.tabs[index].columnLayout.columnWidths = savedLayout.columnWidths
+            tabManager.tabs[index].columnLayout.columnOrder = savedLayout.columnOrder
+        }
+        restoreLastHiddenColumnsForTable(tableName)
+    }
+}
diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+Navigation.swift
@@ -97,6 +97,7 @@ extension MainContentCoordinator {
             }
             // In-place navigation needs selectRedisDatabaseAndQuery to ensure the correct
             // database is SELECTed and session state is updated before querying.
+            restoreColumnLayoutForTable(tableName)
             if navigationModel == .inPlace, let dbIndex = Int(currentDatabase) {
                 selectRedisDatabaseAndQuery(dbIndex)
             } else {
@@ -119,6 +120,7 @@ extension MainContentCoordinator {
                     toolbarState.isTableTab = true
                 AppState.shared.isTableTab = true
                 }
+                restoreColumnLayoutForTable(tableName)
                 if let dbIndex = Int(currentDatabase) {
                     selectRedisDatabaseAndQuery(dbIndex)
                 }
@@ -190,6 +192,7 @@ extension MainContentCoordinator {
                     AppState.shared.isTableTab = true
                 }
                 preview.window.makeKeyAndOrderFront(nil)
+                previewCoordinator.restoreColumnLayoutForTable(tableName)
                 previewCoordinator.runQuery()
                 return
             }
@@ -216,6 +219,7 @@ extension MainContentCoordinator {
                 toolbarState.isTableTab = true
                 AppState.shared.isTableTab = true
             }
+            restoreColumnLayoutForTable(tableName)
             runQuery()
             return
         }
diff --git a/TablePro/Views/Main/Extensions/MainContentCoordinator+TabSwitch.swift b/TablePro/Views/Main/Extensions/MainContentCoordinator+TabSwitch.swift
@@ -26,6 +26,8 @@ extension MainContentCoordinator {
                 tabManager.tabs[oldIndex].pendingChanges = changeManager.saveState()
             }
             tabManager.tabs[oldIndex].filterState = filterStateManager.saveToTabState()
+            saveColumnVisibilityToTab()
+            saveColumnLayoutForTable()
         }
 
         if tabManager.tabs.count > 2 {
diff --git a/TablePro/Views/Main/MainContentView.swift b/TablePro/Views/Main/MainContentView.swift
@@ -501,6 +501,9 @@ struct MainContentView: View {
                             )
                             tabManager.tabs[tabIndex].query = filteredQuery
                         }
+                        if let tableName = selectedTab.tableName {
+                            coordinator.restoreColumnLayoutForTable(tableName)
+                        }
                         coordinator.executeTableTabQueryDirectly()
                     }
                 } else {
@@ -577,6 +580,9 @@ struct MainContentView: View {
                     {
                         Task { await coordinator.switchDatabase(to: selectedTab.databaseName) }
                     } else {
+                        if let tableName = selectedTab.tableName {
+                            coordinator.restoreColumnLayoutForTable(tableName)
+                        }
                         coordinator.executeTableTabQueryDirectly()
                     }
                 } else {
diff --git a/TablePro/Views/Results/DataGridView.swift b/TablePro/Views/Results/DataGridView.swift
@@ -61,6 +61,7 @@ struct DataGridView: NSViewRepresentable {
     var databaseType: DatabaseType?
     var tableName: String?
     var primaryKeyColumn: String?
+    var tabType: TabType?
     var showRowNumbers: Bool = true
     var hiddenColumns: Set<String> = []
     var onHideColumn: ((String) -> Void)?
@@ -271,6 +272,7 @@ struct DataGridView: NSViewRepresentable {
         coordinator.databaseType = databaseType
         coordinator.tableName = tableName
         coordinator.primaryKeyColumn = primaryKeyColumn
+        coordinator.tabType = tabType
 
         coordinator.rebuildVisualStateCache()
 
@@ -336,15 +338,11 @@ struct DataGridView: NSViewRepresentable {
                     column.headerCell.setAccessibilityLabel(
                         String(localized: "Column: \(columnName)")
                     )
-                    if let savedWidth = columnLayout.columnWidths[columnName] {
-                        column.width = savedWidth
-                    } else {
-                        column.width = coordinator.cellFactory.calculateOptimalColumnWidth(
-                            for: columnName,
-                            columnIndex: index,
-                            rowProvider: rowProvider
-                        )
-                    }
+                    column.width = coordinator.cellFactory.calculateOptimalColumnWidth(
+                        for: columnName,
+                        columnIndex: index,
+                        rowProvider: rowProvider
+                    )
                     column.minWidth = 30
                     column.resizingMask = .userResizingMask
                     column.isEditable = isEditable
@@ -358,20 +356,18 @@ struct DataGridView: NSViewRepresentable {
                           colIndex < rowProvider.columns.count else { continue }
                     let columnName = rowProvider.columns[colIndex]
                     column.title = columnName
-                    if let savedWidth = columnLayout.columnWidths[columnName] {
-                        column.width = savedWidth
-                    } else {
-                        column.width = coordinator.cellFactory.calculateOptimalColumnWidth(
-                            for: columnName,
-                            columnIndex: colIndex,
-                            rowProvider: rowProvider
-                        )
-                    }
+                    column.width = coordinator.cellFactory.calculateOptimalColumnWidth(
+                        for: columnName,
+                        columnIndex: colIndex,
+                        rowProvider: rowProvider
+                    )
                     column.isEditable = isEditable
                 }
             }
-            // Restore user-resized column widths after rebuild (only if user explicitly resized)
-            if coordinator.hasUserResizedColumns, !columnLayout.columnWidths.isEmpty {
+            let hasSavedLayout = !columnLayout.columnWidths.isEmpty
+
+            // Restore saved column widths after rebuild (from user resize or persisted layout)
+            if hasSavedLayout {
                 for column in tableView.tableColumns where column.identifier.rawValue != "__rowNumber__" {
                     guard let colIndex = Self.columnIndex(from: column.identifier),
                           colIndex < rowProvider.columns.count else { continue }
@@ -380,16 +376,19 @@ struct DataGridView: NSViewRepresentable {
                         column.width = savedWidth
                     }
                 }
+                coordinator.hasUserResizedColumns = true
             }
 
-            // Restore saved column order after rebuild (only if user explicitly reordered)
-            if coordinator.hasUserResizedColumns, let savedOrder = columnLayout.columnOrder {
+            // Restore saved column order after rebuild
+            if let savedOrder = columnLayout.columnOrder {
                 DataGridView.applyColumnOrder(savedOrder, to: tableView, columns: rowProvider.columns)
+                coordinator.hasUserResizedColumns = true
             }
 
             // Persist calculated widths so subsequent tab switches reuse them
             // instead of calling the expensive calculateOptimalColumnWidth.
-            if !coordinator.hasUserResizedColumns {
+            // Skip when saved layout exists to avoid overwriting persisted values.
+            if !coordinator.hasUserResizedColumns, !hasSavedLayout {
                 var newWidths: [String: CGFloat] = [:]
                 for column in tableView.tableColumns where column.identifier.rawValue != "__rowNumber__" {
                     guard let colIndex = Self.columnIndex(from: column.identifier),
@@ -624,6 +623,7 @@ struct DataGridView: NSViewRepresentable {
 
     static func dismantleNSView(_ nsView: NSScrollView, coordinator: TableViewCoordinator) {
         coordinator.overlayEditor?.dismiss(commit: false)
+        coordinator.persistColumnLayoutToStorage()
         if let observer = coordinator.settingsObserver {
             NotificationCenter.default.removeObserver(observer)
             coordinator.settingsObserver = nil
@@ -681,6 +681,7 @@ final class TableViewCoordinator: NSObject, NSTableViewDelegate, NSTableViewData
     var databaseType: DatabaseType?
     var tableName: String?
     var primaryKeyColumn: String?
+    var tabType: TabType?
 
     /// Check if undo is available
     func canUndo() -> Bool {
@@ -692,6 +693,32 @@ final class TableViewCoordinator: NSObject, NSTableViewDelegate, NSTableViewData
         changeManager.canRedo
     }
 
+    /// Capture current column widths and order from the live NSTableView
+    /// and persist directly to ColumnLayoutStorage. Called from dismantleNSView
+    /// to guarantee layout is saved even when the view is torn down without
+    /// a SwiftUI render cycle (e.g., closing a tab).
+    func persistColumnLayoutToStorage() {
+        guard tabType == .table else { return }
+        guard let tableView, let connectionId, let tableName, !tableName.isEmpty else { return }
+        guard !rowProvider.columns.isEmpty else { return }
+
+        var widths: [String: CGFloat] = [:]
+        var order: [String] = []
+        for column in tableView.tableColumns where column.identifier.rawValue != "__rowNumber__" {
+            guard let colIndex = DataGridView.columnIndex(from: column.identifier),
+                  colIndex < rowProvider.columns.count else { continue }
+            let name = rowProvider.columns[colIndex]
+            widths[name] = column.width
+            order.append(name)
+        }
+
+        guard !widths.isEmpty else { return }
+        var layout = ColumnLayoutState()
+        layout.columnWidths = widths
+        layout.columnOrder = order
+        ColumnLayoutStorage.shared.save(layout, for: tableName, connectionId: connectionId)
+    }
+
     weak var tableView: NSTableView?
     let cellFactory = DataGridCellFactory()
     var overlayEditor: CellOverlayEditor?
@@ -717,6 +744,8 @@ final class TableViewCoordinator: NSObject, NSTableViewDelegate, NSTableViewData
     var hasUserResizedColumns: Bool = false
     /// Guards against two-frame bounce when async column layout write-back triggers updateNSView
     var isWritingColumnLayout: Bool = false
+    /// Debounced work item for persisting column layout after resize/reorder
+    var layoutPersistWorkItem: DispatchWorkItem?
 
     private let cellIdentifier = NSUserInterfaceItemIdentifier("DataCell")
     static let rowViewIdentifier = NSUserInterfaceItemIdentifier("TableRowView")
diff --git a/TablePro/Views/Results/Extensions/DataGridView+Selection.swift b/TablePro/Views/Results/Extensions/DataGridView+Selection.swift
@@ -11,11 +11,22 @@ extension TableViewCoordinator {
         // Only track user-initiated resizes, not programmatic ones during column rebuilds
         guard !isRebuildingColumns else { return }
         hasUserResizedColumns = true
+        scheduleLayoutPersist()
     }
 
     func tableViewColumnDidMove(_ notification: Notification) {
         guard !isRebuildingColumns else { return }
         hasUserResizedColumns = true
+        scheduleLayoutPersist()
+    }
+
+    private func scheduleLayoutPersist() {
+        layoutPersistWorkItem?.cancel()
+        let workItem = DispatchWorkItem { [weak self] in
+            self?.persistColumnLayoutToStorage()
+        }
+        layoutPersistWorkItem = workItem
+        DispatchQueue.main.asyncAfter(deadline: .now() + 0.5, execute: workItem)
     }
 
     func tableViewSelectionDidChange(_ notification: Notification) {

Original file line number	Diff line number	Diff line change
`@@ -329,6 +329,7 @@ struct MainEditorContentView: View {`
`329`	`329`	`databaseType: connection.type,`
`330`	`330`	`tableName: tab.tableName,`
`331`	`331`	`primaryKeyColumn: changeManager.primaryKeyColumn,`
	`332`	`+ tabType: tab.tabType,`
`332`	`333`	`showRowNumbers: AppSettingsManager.shared.dataGrid.showRowNumbers,`
`333`	`334`	`hiddenColumns: columnVisibilityManager.hiddenColumns,`
`334`	`335`	`onHideColumn: { [coordinator] columnName in`
`@@ -473,6 +474,7 @@ struct MainEditorContentView: View {`
`473`	`474`	`}`
`474`	`475`	`DispatchQueue.main.async {`
`475`	`476`	`coordinator.isUpdatingColumnLayout = false`
	`477`	`+ coordinator.saveColumnLayoutForTable()`
`476`	`478`	`}`
`477`	`479`	`}`
`478`	`480`	`)`
Original file line number	Diff line number	Diff line change
`@@ -97,6 +97,7 @@ extension MainContentCoordinator {`
`97`	`97`	`}`
`98`	`98`	`// In-place navigation needs selectRedisDatabaseAndQuery to ensure the correct`
`99`	`99`	`// database is SELECTed and session state is updated before querying.`
	`100`	`+ restoreColumnLayoutForTable(tableName)`
`100`	`101`	`if navigationModel == .inPlace, let dbIndex = Int(currentDatabase) {`
`101`	`102`	`selectRedisDatabaseAndQuery(dbIndex)`
`102`	`103`	`} else {`
`@@ -119,6 +120,7 @@ extension MainContentCoordinator {`
`119`	`120`	`toolbarState.isTableTab = true`
`120`	`121`	`AppState.shared.isTableTab = true`
`121`	`122`	`}`
	`123`	`+ restoreColumnLayoutForTable(tableName)`
`122`	`124`	`if let dbIndex = Int(currentDatabase) {`
`123`	`125`	`selectRedisDatabaseAndQuery(dbIndex)`
`124`	`126`	`}`
`@@ -190,6 +192,7 @@ extension MainContentCoordinator {`
`190`	`192`	`AppState.shared.isTableTab = true`
`191`	`193`	`}`
`192`	`194`	`preview.window.makeKeyAndOrderFront(nil)`
	`195`	`+ previewCoordinator.restoreColumnLayoutForTable(tableName)`
`193`	`196`	`previewCoordinator.runQuery()`
`194`	`197`	`return`
`195`	`198`	`}`
`@@ -216,6 +219,7 @@ extension MainContentCoordinator {`
`216`	`219`	`toolbarState.isTableTab = true`
`217`	`220`	`AppState.shared.isTableTab = true`
`218`	`221`	`}`
	`222`	`+ restoreColumnLayoutForTable(tableName)`
`219`	`223`	`runQuery()`
`220`	`224`	`return`
`221`	`225`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,8 @@ extension MainContentCoordinator {`
`26`	`26`	`tabManager.tabs[oldIndex].pendingChanges = changeManager.saveState()`
`27`	`27`	`}`
`28`	`28`	`tabManager.tabs[oldIndex].filterState = filterStateManager.saveToTabState()`
	`29`	`+ saveColumnVisibilityToTab()`
	`30`	`+ saveColumnLayoutForTable()`
`29`	`31`	`}`
`30`	`32`
`31`	`33`	`if tabManager.tabs.count > 2 {`
Original file line number	Diff line number	Diff line change
`@@ -501,6 +501,9 @@ struct MainContentView: View {`
`501`	`501`	`)`
`502`	`502`	`tabManager.tabs[tabIndex].query = filteredQuery`
`503`	`503`	`}`
	`504`	`+ if let tableName = selectedTab.tableName {`
	`505`	`+ coordinator.restoreColumnLayoutForTable(tableName)`
	`506`	`+ }`
`504`	`507`	`coordinator.executeTableTabQueryDirectly()`
`505`	`508`	`}`
`506`	`509`	`} else {`
`@@ -577,6 +580,9 @@ struct MainContentView: View {`
`577`	`580`	`{`
`578`	`581`	`Task { await coordinator.switchDatabase(to: selectedTab.databaseName) }`
`579`	`582`	`} else {`
	`583`	`+ if let tableName = selectedTab.tableName {`
	`584`	`+ coordinator.restoreColumnLayoutForTable(tableName)`
	`585`	`+ }`
`580`	`586`	`coordinator.executeTableTabQueryDirectly()`
`581`	`587`	`}`
`582`	`588`	`} else {`