From 5a3baf65d8b2bbb7f3f9b33767255fc6f3543ff0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ng=C3=B4=20Qu=E1=BB=91c=20=C4=90=E1=BA=A1t?=
 <datlechin@gmail.com>
Date: Wed, 3 Jun 2026 17:06:36 +0700
Subject: [PATCH] fix(datagrid): keep query result columns in SELECT order
 (#1565)

---
 CHANGELOG.md                                  |  4 +++
 .../Views/Results/DataGridColumnPool.swift    |  2 +-
 .../Views/Results/DataGridCoordinator.swift   | 10 ++++--
 .../TableViewCoordinatorLayoutTests.swift     | 35 +++++++++++++++----
 4 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 13e469c93..bb3cb0ce3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Query result columns now follow the order in the SELECT. Adding or removing a column no longer leaves new columns stuck at the end of the grid. (#1565)
+
 ## [0.48.0] - 2026-06-02
 
 ### Added
diff --git a/TablePro/Views/Results/DataGridColumnPool.swift b/TablePro/Views/Results/DataGridColumnPool.swift
index 20220313c..d6e81741e 100644
--- a/TablePro/Views/Results/DataGridColumnPool.swift
+++ b/TablePro/Views/Results/DataGridColumnPool.swift
@@ -46,7 +46,7 @@ final class DataGridColumnPool {
             if slot < visibleCount {
                 let columnName = schema.columnNames[slot]
                 let resolvedWidth = willRestoreWidths
-                    ? (savedLayout?.columnWidths[columnName] ?? 100)
+                    ? (savedLayout?.columnWidths[columnName] ?? widthCalculator(columnName, slot))
                     : widthCalculator(columnName, slot)
                 configureColumn(
                     column,
diff --git a/TablePro/Views/Results/DataGridCoordinator.swift b/TablePro/Views/Results/DataGridCoordinator.swift
index 05f8dc9a8..475f8c0f1 100644
--- a/TablePro/Views/Results/DataGridCoordinator.swift
+++ b/TablePro/Views/Results/DataGridCoordinator.swift
@@ -51,8 +51,14 @@ final class TableViewCoordinator: NSObject, NSTableViewDelegate, NSTableViewData
     }
 
     func savedColumnLayout(binding: ColumnLayoutState) -> ColumnLayoutState? {
-        if tabType == .table,
-           let connectionId,
+        guard tabType == .table else {
+            guard !binding.columnWidths.isEmpty else { return nil }
+            var layout = binding
+            layout.columnOrder = nil
+            return layout
+        }
+
+        if let connectionId,
            let tableName,
            !tableName.isEmpty,
            let stored = layoutPersister.load(for: tableName, connectionId: connectionId) {
diff --git a/TableProTests/Views/Results/TableViewCoordinatorLayoutTests.swift b/TableProTests/Views/Results/TableViewCoordinatorLayoutTests.swift
index 2091593ab..3d9dc948f 100644
--- a/TableProTests/Views/Results/TableViewCoordinatorLayoutTests.swift
+++ b/TableProTests/Views/Results/TableViewCoordinatorLayoutTests.swift
@@ -4,8 +4,8 @@
 //
 
 import Foundation
-import TableProPluginKit
 import SwiftUI
+import TableProPluginKit
 import Testing
 
 @testable import TablePro
@@ -97,20 +97,41 @@ struct TableViewCoordinatorLayoutTests {
         #expect(coordinator.savedColumnLayout(binding: ColumnLayoutState()) == nil)
     }
 
-    @Test("Non-table tab uses the binding directly")
-    func nonTableTabUsesBinding() {
+    @Test("Query tab drops a stale saved column order so new columns keep their query position")
+    func queryTabDropsStaleColumnOrder() {
         let coordinator = makeCoordinator(
             tabType: .query,
             connectionId: nil,
             tableName: nil,
             persister: FakeColumnLayoutPersister()
         )
-        let resolved = coordinator.savedColumnLayout(binding: nonEmptyLayout())
-        #expect(resolved?.columnWidths == ["id": 60])
+        var binding = ColumnLayoutState()
+        binding.columnWidths = ["id": 60, "business_model": 120]
+        binding.columnOrder = ["id", "business_model"]
+
+        var expected = ColumnLayoutState()
+        expected.columnWidths = ["id": 60, "business_model": 120]
+
+        #expect(coordinator.savedColumnLayout(binding: binding) == expected)
+    }
+
+    @Test("Query tab keeps remembered widths when there is no saved order")
+    func queryTabKeepsWidths() {
+        let coordinator = makeCoordinator(
+            tabType: .query,
+            connectionId: nil,
+            tableName: nil,
+            persister: FakeColumnLayoutPersister()
+        )
+
+        var expected = ColumnLayoutState()
+        expected.columnWidths = ["id": 60]
+
+        #expect(coordinator.savedColumnLayout(binding: nonEmptyLayout()) == expected)
     }
 
-    @Test("Non-table tab returns nil when binding is empty")
-    func nonTableTabEmptyReturnsNil() {
+    @Test("Query tab returns nil when binding is empty")
+    func queryTabEmptyReturnsNil() {
         let coordinator = makeCoordinator(
             tabType: .query,
             connectionId: nil,