From 31402ca00c4bcaa72568733ab030bcafd82c1e87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ng=C3=B4=20Qu=E1=BB=91c=20=C4=90=E1=BA=A1t?= Date: Wed, 8 Apr 2026 10:43:58 +0700 Subject: [PATCH 1/2] perf: concurrent schema fetch, memoize code highlighting, cap tables before fetch --- .../Core/Autocomplete/SQLSchemaProvider.swift | 3 +- TablePro/ViewModels/AIChatViewModel.swift | 32 +++++++++++-------- .../Views/AIChat/AIChatCodeBlockView.swift | 26 +++++++++++++-- 3 files changed, 44 insertions(+), 17 deletions(-) diff --git a/TablePro/Core/Autocomplete/SQLSchemaProvider.swift b/TablePro/Core/Autocomplete/SQLSchemaProvider.swift index eae7ded5..e5968672 100644 --- a/TablePro/Core/Autocomplete/SQLSchemaProvider.swift +++ b/TablePro/Core/Autocomplete/SQLSchemaProvider.swift @@ -164,7 +164,8 @@ actor SQLSchemaProvider { guard !tables.isEmpty, let connection = connectionInfo else { return nil } var columnsByTable: [String: [ColumnInfo]] = [:] - for table in tables { + let tablesToFetch = Array(tables.prefix(settings.maxSchemaTables)) + for table in tablesToFetch { let columns = await getColumns(for: table.name) if !columns.isEmpty { columnsByTable[table.name.lowercased()] = columns diff --git a/TablePro/ViewModels/AIChatViewModel.swift b/TablePro/ViewModels/AIChatViewModel.swift index 89129ac0..baf45f8e 100644 --- a/TablePro/ViewModels/AIChatViewModel.swift +++ b/TablePro/ViewModels/AIChatViewModel.swift @@ -445,22 +445,26 @@ final class AIChatViewModel { var columns: [String: [ColumnInfo]] = [:] var foreignKeys: [String: [ForeignKeyInfo]] = [:] - for table in tablesToFetch { - if let schemaProvider { - let cached = await schemaProvider.getColumns(for: table.name) - if !cached.isEmpty { - columns[table.name] = cached + await withTaskGroup(of: (String, [ColumnInfo]).self) { group in + for table in tablesToFetch { + group.addTask { [schemaProvider] in + if let schemaProvider { + let cached = await schemaProvider.getColumns(for: table.name) + if !cached.isEmpty { + return (table.name, cached) + } + } + do { + let cols = try await driver.fetchColumns(table: table.name) + return (table.name, cols) + } catch { + return (table.name, []) + } } } - - if columns[table.name] == nil { - do { - let cols = try await driver.fetchColumns(table: table.name) - columns[table.name] = cols - } catch { - Self.logger.warning( - "Failed to fetch columns for table '\(table.name)': \(error.localizedDescription)" - ) + for await (tableName, cols) in group { + if !cols.isEmpty { + columns[tableName] = cols } } } diff --git a/TablePro/Views/AIChat/AIChatCodeBlockView.swift b/TablePro/Views/AIChat/AIChatCodeBlockView.swift index 50aa2071..dc82937d 100644 --- a/TablePro/Views/AIChat/AIChatCodeBlockView.swift +++ b/TablePro/Views/AIChat/AIChatCodeBlockView.swift @@ -14,6 +14,8 @@ struct AIChatCodeBlockView: View { let language: String? @State private var isCopied: Bool = false + @State private var cachedHighlight: AttributedString? + @State private var cachedCodeLength: Int = 0 @FocusedValue(\.commandActions) private var actions var body: some View { @@ -77,11 +79,11 @@ struct AIChatCodeBlockView: View { private var codeContent: some View { ScrollView(.horizontal, showsIndicators: false) { if isSQL { - Text(highlightedSQL(code)) + Text(currentHighlight { highlightedSQL(code) }) .textSelection(.enabled) .padding(10) } else if isMongoDB { - Text(highlightedJavaScript(code)) + Text(currentHighlight { highlightedJavaScript(code) }) .textSelection(.enabled) .padding(10) } else if isRedis { @@ -96,6 +98,26 @@ struct AIChatCodeBlockView: View { .padding(10) } } + .onChange(of: code) { + let newLen = (code as NSString).length + // Re-highlight when code grows by 50+ chars to avoid per-token work during streaming + if newLen - cachedCodeLength >= 50 { + cachedHighlight = isSQL ? highlightedSQL(code) : highlightedJavaScript(code) + cachedCodeLength = newLen + } + } + } + + private func currentHighlight(_ compute: () -> AttributedString) -> AttributedString { + if let cached = cachedHighlight, (code as NSString).length - cachedCodeLength < 50 { + return cached + } + let result = compute() + DispatchQueue.main.async { + cachedHighlight = result + cachedCodeLength = (code as NSString).length + } + return result } private var isSQL: Bool { From c18716eca7599073f0aa464b99256a8dfcc82330 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ng=C3=B4=20Qu=E1=BB=91c=20=C4=90=E1=BA=A1t?= Date: Wed, 8 Apr 2026 10:49:22 +0700 Subject: [PATCH 2/2] fix: simplify highlight cache to avoid state writes during body, add debug log for fetch errors --- TablePro/ViewModels/AIChatViewModel.swift | 1 + .../Views/AIChat/AIChatCodeBlockView.swift | 28 ++++++------------- 2 files changed, 10 insertions(+), 19 deletions(-) diff --git a/TablePro/ViewModels/AIChatViewModel.swift b/TablePro/ViewModels/AIChatViewModel.swift index baf45f8e..6331023f 100644 --- a/TablePro/ViewModels/AIChatViewModel.swift +++ b/TablePro/ViewModels/AIChatViewModel.swift @@ -458,6 +458,7 @@ final class AIChatViewModel { let cols = try await driver.fetchColumns(table: table.name) return (table.name, cols) } catch { + Self.logger.debug("Schema context: failed to fetch columns for '\(table.name)'") return (table.name, []) } } diff --git a/TablePro/Views/AIChat/AIChatCodeBlockView.swift b/TablePro/Views/AIChat/AIChatCodeBlockView.swift index dc82937d..1de90861 100644 --- a/TablePro/Views/AIChat/AIChatCodeBlockView.swift +++ b/TablePro/Views/AIChat/AIChatCodeBlockView.swift @@ -78,12 +78,8 @@ struct AIChatCodeBlockView: View { private var codeContent: some View { ScrollView(.horizontal, showsIndicators: false) { - if isSQL { - Text(currentHighlight { highlightedSQL(code) }) - .textSelection(.enabled) - .padding(10) - } else if isMongoDB { - Text(currentHighlight { highlightedJavaScript(code) }) + if isSQL || isMongoDB { + Text(cachedHighlight ?? AttributedString(code)) .textSelection(.enabled) .padding(10) } else if isRedis { @@ -98,26 +94,20 @@ struct AIChatCodeBlockView: View { .padding(10) } } + .onAppear { + recomputeHighlight() + } .onChange(of: code) { let newLen = (code as NSString).length - // Re-highlight when code grows by 50+ chars to avoid per-token work during streaming if newLen - cachedCodeLength >= 50 { - cachedHighlight = isSQL ? highlightedSQL(code) : highlightedJavaScript(code) - cachedCodeLength = newLen + recomputeHighlight() } } } - private func currentHighlight(_ compute: () -> AttributedString) -> AttributedString { - if let cached = cachedHighlight, (code as NSString).length - cachedCodeLength < 50 { - return cached - } - let result = compute() - DispatchQueue.main.async { - cachedHighlight = result - cachedCodeLength = (code as NSString).length - } - return result + private func recomputeHighlight() { + cachedHighlight = isSQL ? highlightedSQL(code) : highlightedJavaScript(code) + cachedCodeLength = (code as NSString).length } private var isSQL: Bool {