@@ -59,9 +59,29 @@ final class QueryViewModel {
5959 var columns: [String] = []
6060 var columnOrder: [String] = []
6161 var hiddenColumns: Set<String> = []
62+ /// The full persisted set of hidden columns for the current stream, including
63+ /// ones absent from the current result. Hiding a column drops it from the
64+ /// SELECT list, so it's missing from the very next result and `hiddenColumns`
65+ /// (only the columns present) forgets it — the next save would then erase the
66+ /// preference. This keeps the intent, mirroring LiveTailViewModel.
67+ private(set) var savedHiddenColumns: Set<String> = []
68+ /// The persisted column order including columns absent from the current
69+ /// result (the order counterpart of `savedHiddenColumns`), so saving doesn't
70+ /// drop a hidden column's position and it returns to its place when shown.
71+ private var savedColumnOrder: [String] = []
6272 var autoHiddenColumns: Set<String> = []
6373 /// Column config from a saved query, applied on next query execution.
6474 var pendingColumnConfig: ColumnConfiguration?
75+
76+ /// The connection and stream the editor text and results belong to. The view
77+ /// model outlives QueryView (switching tabs keeps the query), so the view
78+ /// compares this on appear to catch a stream/server change made while another
79+ /// tab was showing.
80+ struct BoundContext: Equatable {
81+ let connectionID: UUID?
82+ let stream: String
83+ }
84+ var boundContext: BoundContext?
6585 private var currentStream: String?
6686 var isLoading = false
6787 var errorMessage: String?
@@ -192,9 +212,16 @@ final class QueryViewModel {
192212 return columnOrder.filter { !allHidden.contains($0) }
193213 }
194214
215+ /// Whether any column is hidden — including saved-hidden columns that aren't
216+ /// in the current result, which "Show All" can bring back.
217+ var hasHiddenColumns: Bool {
218+ !hiddenColumns.isEmpty || !autoHiddenColumns.isEmpty || !savedHiddenColumns.isEmpty
219+ }
220+
195221 func toggleColumnVisibility(_ column: String) {
196222 if hiddenColumns.contains(column) {
197223 hiddenColumns.remove(column)
224+ savedHiddenColumns.remove(column)
198225 } else if autoHiddenColumns.contains(column) {
199226 autoHiddenColumns.remove(column)
200227 } else {
@@ -203,6 +230,7 @@ final class QueryViewModel {
203230 let visibleCount = columnOrder.count - allHidden.count
204231 if visibleCount > 1 {
205232 hiddenColumns.insert(column)
233+ savedHiddenColumns.insert(column)
206234 }
207235 }
208236 saveColumnConfig()
@@ -211,6 +239,7 @@ final class QueryViewModel {
211239
212240 func showAllColumns() {
213241 hiddenColumns.removeAll()
242+ savedHiddenColumns.removeAll()
214243 autoHiddenColumns.removeAll()
215244 saveColumnConfig()
216245 updateSQLColumns()
@@ -240,32 +269,116 @@ final class QueryViewModel {
240269
241270 func resetColumnConfig() {
242271 columnOrder = columns
272+ savedColumnOrder = []
243273 hiddenColumns.removeAll()
274+ savedHiddenColumns.removeAll()
244275 autoHiddenColumns.removeAll()
245276 saveColumnConfig()
246277 updateSQLColumns()
247278 }
248279
249− /// Rewrites the SELECT clause of `sqlQuery` to reflect the current `visibleColumns`.
280+ /// Rewrites the SELECT clause of `sqlQuery` to reflect the user's column
281+ /// choices (order and hidden columns; auto-hidden columns stay display-only).
250282 func updateSQLColumns() {
251− guard let range = SQLTokenizer.selectColumnListRange(in: sqlQuery) else { return }
252−
253− let visible = visibleColumns
283+ guard let range = SQLTokenizer.selectColumnListRange(in: sqlQuery),
284+ let currentList = Self.parseColumnList(sqlQuery[range]) else { return }
285+
286+ // Project the user's choices only. `autoHiddenColumns` is a display
287+ // heuristic for the current rows (the field is empty in all of them);
288+ // baking it into the SQL would stop the field from ever being fetched
289+ // again, even for a time range where it has values.
290+ let visible = columnOrder.filter { !hiddenColumns.contains($0) }
254291 guard !visible.isEmpty else { return }
255292
293+ // Columns an explicit list names that aren't in the results: Parseable
294+ // omits all-null fields from query JSON, so a sparse listed field can be
295+ // missing from `columns`. Keep it (unless the user hid it) rather than
296+ // silently dropping it from the SELECT on the next rewrite.
297+ var listedButAbsent: [String] = []
298+ var excludesHiddenColumn = false
299+ if case .identifiers(let names) = currentList {
300+ let present = Set(columns)
301+ listedButAbsent = names.filter { !present.contains($0) && !savedHiddenColumns.contains($0) }
302+ // A column the user hid is absent from these results precisely
303+ // because this list was narrowed to exclude it.
304+ excludesHiddenColumn = !savedHiddenColumns.isSubset(of: present)
305+ }
306+
256307 let newColumnList: String
257308 // Use `*` whenever every column is visible, regardless of display order —
258309 // reordering columns (a set-preserving permutation) must not pin the query
259− // to an explicit snapshot that hides newly-appearing fields on re-query.
260− if visible.count == columns.count && Set(visible) == Set(columns) {
310+ // to an explicit snapshot that hides newly-appearing fields on re-query —
311+ // unless the list is explicit because it leaves out a hidden column, which
312+ // `*` would fetch again.
313+ if !excludesHiddenColumn && visible.count == columns.count && Set(visible) == Set(columns) {
261314 newColumnList = "*"
262315 } else {
263− newColumnList = visible.map { Self.escapeSQLIdentifier($0) }.joined(separator: ", ")
316+ newColumnList = (visible + listedButAbsent).map { Self.escapeSQLIdentifier($0) }.joined(separator: ", ")
264317 }
265318
266319 sqlQuery.replaceSubrange(range, with: newColumnList)
267320 }
268321
322+ enum ColumnListShape: Equatable {
323+ case star
324+ /// Plain column references, as the column names they resolve to.
325+ case identifiers([String])
326+ }
327+
328+ /// Parses a SELECT list the column manager can regenerate without changing
329+ /// what the query computes: `*`, or a comma-separated list of plain column
330+ /// references. Anything else — aggregates, expressions, aliases
331+ /// (`count(*) AS n`), qualified names — returns nil: rewriting it with `*` or
332+ /// bare result keys would break the query, so the column manager then only
333+ /// hides/reorders client-side.
334+ nonisolated static func parseColumnList(_ list: Substring) -> ColumnListShape? {
335+ let tokens = SQLTokenizer.tokenize(String(list)).filter { !$0.kind.isTrivia }
336+ if tokens.count == 1, tokens[0].kind == .star { return .star }
337+ guard tokens.count % 2 == 1 else { return nil }
338+ var names: [String] = []
339+ for (index, token) in tokens.enumerated() {
340+ if index % 2 == 1 {
341+ guard token.kind == .comma else { return nil }
342+ continue
343+ }
344+ switch token.kind {
345+ // DataFusion folds unquoted identifiers to lowercase.
346+ case .identifier(let name): names.append(name.lowercased())
347+ case .quotedIdentifier(let name): names.append(name)
348+ default: return nil
349+ }
350+ }
351+ return .identifiers(names)
352+ }
353+
354+ /// Inserts `newColumns` into `order`, placing each one that appears in
355+ /// `reference` (a saved order) right after its nearest saved predecessor
356+ /// that's already placed — or first if none is — and appending the rest.
357+ /// Used to keep the positions of columns that are absent for a while.
358+ nonisolated static func mergedColumnOrder(
359+ _ order: [String], inserting newColumns: [String], reference: [String]
360+ ) -> [String] {
361+ var result = order
362+ let referenceIndex = Dictionary(reference.enumerated().map { ($1, $0) }, uniquingKeysWith: { first, _ in first })
363+ let pending = newColumns.filter { referenceIndex[$0] == nil }
364+ let positioned = newColumns
365+ .filter { referenceIndex[$0] != nil }
366+ .sorted { referenceIndex[$0]! < referenceIndex[$1]! }
367+ for column in positioned where !result.contains(column) {
368+ let savedIndex = referenceIndex[column]!
369+ let predecessor = reference[..<savedIndex].reversed().first { result.contains($0) }
370+ if let predecessor, let position = result.firstIndex(of: predecessor) {
371+ result.insert(column, at: position + 1)
372+ } else {
373+ result.insert(column, at: 0)
374+ }
375+ }
376+ for column in pending where !result.contains(column) {
377+ result.append(column)
378+ }
379+ return result
380+ }
381+
269382 // MARK: - Column Configuration Persistence
270383
271384 struct ColumnConfiguration: Codable {
@@ -279,7 +392,16 @@ final class QueryViewModel {
279392
280393 private func saveColumnConfig() {
281394 guard let stream = currentStream else { return }
282− let config = ColumnConfiguration(order: columnOrder, hidden: hiddenColumns)
395+ // Persist the full intended hidden set and order, not just the currently
396+ // present columns (see `savedHiddenColumns` / `savedColumnOrder`).
397+ let present = Set(columnOrder)
398+ let order = Self.mergedColumnOrder(
399+ columnOrder,
400+ inserting: savedColumnOrder.filter { !present.contains($0) },
401+ reference: savedColumnOrder
402+ )
403+ savedColumnOrder = order
404+ let config = ColumnConfiguration(order: order, hidden: savedHiddenColumns.union(hiddenColumns))
283405 if let data = try? JSONEncoder().encode(config) {
284406 UserDefaults.standard.set(data, forKey: Self.columnConfigKey(for: stream))
285407 }
@@ -308,7 +430,9 @@ final class QueryViewModel {
308430
309431 guard let config else {
310432 columnOrder = extractedColumns
433+ savedColumnOrder = []
311434 hiddenColumns = []
435+ savedHiddenColumns = []
312436 return
313437 }
314438
@@ -321,6 +445,8 @@ final class QueryViewModel {
321445 merged.append(col)
322446 }
323447 columnOrder = merged
448+ savedColumnOrder = config.order
449+ savedHiddenColumns = config.hidden
324450 hiddenColumns = config.hidden.intersection(extractedSet)
325451
326452 // The `visibleCount > 1` guard in toggleColumnVisibility only covers
@@ -519,6 +645,8 @@ final class QueryViewModel {
519645 // report stale column names after the failure.
520646 columnOrder = []
521647 hiddenColumns = []
648+ savedHiddenColumns = []
649+ savedColumnOrder = []
522650 autoHiddenColumns = []
523651 }
524652 }
@@ -685,6 +813,8 @@ final class QueryViewModel {
685813 columns = []
686814 columnOrder = []
687815 hiddenColumns = []
816+ savedHiddenColumns = []
817+ savedColumnOrder = []
688818 autoHiddenColumns = []
689819 currentStream = nil
690820 resultCount = 0
@@ -698,6 +828,7 @@ final class QueryViewModel {
698828 // a later reload guard never sees a non-nil stream paired with empty data.
699829 latestSchemaStream = nil
700830 latestQueryStream = nil
831+ boundContext = nil
701832 }
702833
703834 /// Modifies `sqlQuery` to add a column-value filter condition.
@@ -707,75 +838,139 @@ final class QueryViewModel {
707838 let condition: String
708839 if let value, value != .null {
709840 let literal = value.sqlLiteral
710− condition = exclude ? "\(escapedCol) <> \(literal)" : "\(escapedCol) = \(literal)"
841+ // Exclude with `IS DISTINCT FROM`, not `<>`: under SQL's three-valued
842+ // logic `NULL <> 'x'` is NULL, so `<>` also dropped every row where the
843+ // column is missing/NULL — common for sparse fields — which "exclude
844+ // this value" doesn't ask for (the live-tail exclude keeps them too).
845+ condition = exclude
846+ ? "\(escapedCol) IS DISTINCT FROM \(literal)"
847+ : "\(escapedCol) = \(literal)"
711848 } else {
712849 condition = exclude ? "\(escapedCol) IS NOT NULL" : "\(escapedCol) IS NULL"
713850 }
714851
715− var sql = sqlQuery.trimmingCharacters(in: .whitespacesAndNewlines)
852+ let sql = sqlQuery.trimmingCharacters(in: .whitespacesAndNewlines)
716853 if sql.isEmpty { return }
717854
718855 // Locate the top-level WHERE and the first trailing clause using the
719856 // tokenizer rather than a raw regex, so keywords inside quoted
720857 // identifiers, string literals, comments, or subqueries are ignored.
721− let tailKeywords: Set<String> = ["ORDER", "GROUP", "HAVING", "LIMIT", "OFFSET"]
858+ let tokens = SQLTokenizer.tokenize(sql)
722859 var depth = 0
723− var whereEnd: String.Index?
724− var tailStart: String.Index?
860+ // WHERE and the trailing clauses only follow the top-level FROM. Keywords
861+ // before it belong to the SELECT list — e.g. a field named `offset` or
862+ // `limit` (which the tokenizer reports as keywords) — and must not be
863+ // mistaken for clause starts.
864+ var fromSeen = false
865+ var whereIndex: Int?
866+ var tailIndex: Int?
725867 // Track whether the existing WHERE body has a top-level OR — only then does
726868 // AND-ing a new condition need parentheses to preserve precedence.
727869 var bodyHasTopLevelOr = false
728− for token in SQLTokenizer.tokenize(sql) {
870+ // The last token that is actual SQL — not whitespace, a comment, or a
871+ // statement-terminating `;`. New text is spliced in right after such
872+ // tokens (never at the raw end of the string), so a trailing `;` or a
873+ // trailing `-- comment` stays at the end instead of swallowing the
874+ // appended condition: `... LIMIT 10;` must not become `... LIMIT 10; WHERE`,
875+ // and `... -- note` must not become `... -- note WHERE <commented out>`.
876+ func isSignificant(_ token: SQLTokenizer.Token) -> Bool {
877+ !token.kind.isTrivia && token.kind != .other(";")
878+ }
879+ func nextSignificantKind(after index: Int) -> SQLTokenizer.TokenKind? {
880+ tokens[(index + 1)...].first(where: { isSignificant($0) })?.kind
881+ }
882+ // Whether the keyword at `index` actually opens a trailing clause. ORDER and
883+ // GROUP must be followed by BY (so `WITHIN GROUP (…)` or a column named
884+ // `order` doesn't count), and LIMIT/OFFSET by their row count (so a column
885+ // named `offset` in `WHERE offset > 5` doesn't cut the WHERE body short).
886+ func startsTrailingClause(_ keyword: String, at index: Int) -> Bool {
887+ switch keyword {
888+ case "ORDER", "GROUP":
889+ return nextSignificantKind(after: index) == .keyword("BY")
890+ case "LIMIT", "OFFSET":
891+ if case .number = nextSignificantKind(after: index) { return true }
892+ return false
893+ case "HAVING":
894+ return true
895+ default:
896+ return false
897+ }
898+ }
899+ var lastSignificant: Int?
900+ for (idx, token) in tokens.enumerated() {
901+ if isSignificant(token) { lastSignificant = idx }
729902 switch token.kind {
730903 case .leftParen:
731904 depth += 1
732905 case .rightParen:
733906 depth = max(0, depth - 1)
734− case .keyword(let kw) where depth == 0:
735− if kw == "WHERE", whereEnd == nil {
736− whereEnd = token.range.upperBound
737− } else if tailKeywords.contains(kw), tailStart == nil {
738− tailStart = token.range.lowerBound
739− } else if kw == "OR", whereEnd != nil, tailStart == nil {
907+ case .keyword("FROM") where depth == 0:
908+ fromSeen = true
909+ case .keyword(let kw) where depth == 0 && fromSeen:
910+ if kw == "WHERE", whereIndex == nil {
911+ whereIndex = idx
912+ // A trailing clause seen *before* the WHERE (e.g. the GROUP BY of
913+ // the first SELECT in `... GROUP BY a UNION ALL SELECT ... WHERE`)
914+ // doesn't end this WHERE's body. Keeping it would make the body
915+ // range run backwards (a trapping `Range`), so look for the
916+ // body's end afresh after the WHERE.
917+ tailIndex = nil
918+ } else if tailIndex == nil, startsTrailingClause(kw, at: idx) {
919+ tailIndex = idx
920+ } else if kw == "OR", whereIndex != nil, tailIndex == nil {
740921 bodyHasTopLevelOr = true
741922 }
742923 default:
743924 break
744925 }
745926 }
746−
747− if let whereEnd {
748− // Parenthesize the existing WHERE body before AND-ing the new
749− // condition only when it has a top-level OR, otherwise the OR
750− // silently changes meaning: `WHERE a = 1 OR b = 2` would become
751− // `WHERE a = 1 OR b = 2 AND c = 3`, which SQL binds as
752− // `a = 1 OR (b = 2 AND c = 3)` rather than the intended
753− // `(a = 1 OR b = 2) AND c = 3`. A body without a top-level OR is
754− // left un-parenthesized so simple filters stay readable.
755− let bodyEnd = tailStart ?? sql.endIndex
756− let head = sql[..<whereEnd]
757− let body = sql[whereEnd..<bodyEnd].trimmingCharacters(in: .whitespacesAndNewlines)
758− let tail = sql[bodyEnd...].trimmingCharacters(in: .whitespacesAndNewlines)
759− // An empty WHERE body (e.g. a half-written `... WHERE` with no
760− // condition yet, or `WHERE` immediately followed by ORDER/GROUP/
761− // LIMIT) would otherwise produce `WHERE AND <condition>` — a
762− // dangling AND with no left operand. Drop the AND in that case.
763− if body.isEmpty {
764− sql = "\(head) \(condition)"
927+ // Nothing but comments/semicolons — there's no statement to filter.
928+ guard let lastSignificant else { return }
929+
930+ // Each branch rebuilds the query from slices of the original text, keeping
931+ // the user's line breaks: joining pieces with single spaces would pull the
932+ // text after a `-- comment` onto the comment's line and comment it out.
933+ // Text spliced directly in front of another token (`(a = 1)ORDER BY`) needs
934+ // a space so the two don't fuse (`... IS NULLORDER BY`).
935+ func separated(_ rest: Substring) -> String {
936+ guard let next = rest.first, !next.isWhitespace, next != ";" else { return String(rest) }
937+ return " " + rest
938+ }
939+ if let whereIndex {
940+ // The WHERE body runs up to the first trailing clause, or to the end
941+ // of the statement. Only its significant tokens bound the splice, so a
942+ // comment between the body and ORDER BY stays after the new condition.
943+ let bodyEndIndex = tailIndex ?? (lastSignificant + 1)
944+ let bodyIndices = (whereIndex + 1)..<bodyEndIndex
945+ if let first = bodyIndices.first(where: { isSignificant(tokens[$0]) }),
946+ let last = bodyIndices.last(where: { isSignificant(tokens[$0]) }) {
947+ // Parenthesize the existing WHERE body before AND-ing the new
948+ // condition only when it has a top-level OR, otherwise the OR
949+ // silently changes meaning: `WHERE a = 1 OR b = 2` would become
950+ // `WHERE a = 1 OR b = 2 AND c = 3`, which SQL binds as
951+ // `a = 1 OR (b = 2 AND c = 3)` rather than the intended
952+ // `(a = 1 OR b = 2) AND c = 3`. A body without a top-level OR is
953+ // left un-parenthesized so simple filters stay readable.
954+ let bodyStart = tokens[first].range.lowerBound
955+ let bodyEnd = tokens[last].range.upperBound
956+ let body = sql[bodyStart..<bodyEnd]
957+ let wrappedBody = bodyHasTopLevelOr ? "(\(body))" : String(body)
958+ sqlQuery = "\(sql[..<bodyStart])\(wrappedBody) AND \(condition)\(separated(sql[bodyEnd...]))"
765959 } else {
766− let wrappedBody = bodyHasTopLevelOr ? "(\(body))" : body
767− sql = "\(head) \(wrappedBody) AND \(condition)"
960+ // An empty WHERE body (e.g. a half-written `... WHERE` with no
961+ // condition yet, or `WHERE` immediately followed by ORDER/GROUP/
962+ // LIMIT) would otherwise produce `WHERE AND <condition>` — a
963+ // dangling AND with no left operand. Drop the AND in that case.
964+ let whereEnd = tokens[whereIndex].range.upperBound
965+ sqlQuery = "\(sql[..<whereEnd]) \(condition)\(separated(sql[whereEnd...]))"
768966 }
769− if !tail.isEmpty {
770− sql += " \(tail)"
771− }
772− } else if let tailStart {
773− sql.insert(contentsOf: "WHERE \(condition) ", at: tailStart)
967+ } else if let tailIndex {
968+ let tailStart = tokens[tailIndex].range.lowerBound
969+ sqlQuery = "\(sql[..<tailStart])WHERE \(condition) \(sql[tailStart...])"
774970 } else {
775− sql += " WHERE \(condition)"
971+ let statementEnd = tokens[lastSignificant].range.upperBound
972+ sqlQuery = "\(sql[..<statementEnd]) WHERE \(condition)\(separated(sql[statementEnd...]))"
776973 }
777−
778− sqlQuery = sql
779974 }
780975
781976 /// Sets the default query for the given stream, returning `true` if the
@@ -819,6 +1014,10 @@ final class QueryViewModel {
8191014 // MARK: - Query History
8201015
8211016 private func addToHistory(_ entry: QueryHistoryEntry) {
1017+ // Start from what's persisted, not this instance's copy: a model created
1018+ // earlier (e.g. one whose query finished after a reconnect replaced it)
1019+ // would otherwise overwrite newer entries saved by the current one.
1020+ queryHistory = Self.loadHistory()
8221021 // Deduplicate consecutive identical queries
8231022 if queryHistory.first?.sql == entry.sql { return }
8241023 queryHistory.insert(entry, at: 0)