feat(ai): give the schema prompt to the tables the question is about

apple/Packages/TableProCore/Sources/TableProAutocomplete/SQLContextAnalyzer.swift

@@ -278,7 +278,9 @@ public final class SQLContextAnalyzer {
         "|WINDOW|FETCH|FOR)\\b|[;()]|$)"
     )
 
-    private static let alterTableRegex = compileRegex("(?i)\\bALTER\\s+TABLE\\s+(\(quotedSegment))")
+    private static let alterTableRegex = compileRegex(
+        "(?i)\\bALTER\\s+TABLE\\s+(\(quotedSegment)(?:\\.\(quotedSegment))*)"
+    )
 
     private static let derivedTableParser = DerivedTableParser()
 
@@ -288,6 +290,19 @@ public final class SQLContextAnalyzer {
 
     public init() {}
 
+    // `analyze` reads no tables when its cursor sits in a comment or literal; this reads them wherever the statement
+    // ends, so `SELECT * FROM orders -- note` still names `orders`. A CTE's name comes back like a table's, since a
+    // CTE can share a table's name and still read that table in its own body.
+    public func tableReferences(in statement: String, grammar: SQLLexicalGrammar) -> [TableReference] {
+        let nsStatement = statement as NSString
+        let text = Self.lex(nsStatement, cursor: nsStatement.length, grammar: grammar)
+        var references = extractTableReferences(from: text.forReferences)
+        if let alterRef = extractAlterTableReference(from: text.forReferences), !references.contains(alterRef) {
+            references.append(alterRef)
+        }
+        return references
+    }
+
     // MARK: - Main Analysis
 
     public func analyze(query: String, cursorPosition: Int, grammar: SQLLexicalGrammar) -> SQLContext {
@@ -328,8 +343,7 @@ public final class SQLContextAnalyzer {
 
         var seenReferences = Set<TableReference>(tableReferences)
 
-        if let alterTableName = extractAlterTableName(from: text.forReferences) {
-            let alterRef = TableReference(tableName: alterTableName, alias: nil)
+        if let alterRef = extractAlterTableReference(from: text.forReferences) {
             if seenReferences.insert(alterRef).inserted {
                 tableReferences.append(alterRef)
             }
@@ -1003,12 +1017,19 @@ public final class SQLContextAnalyzer {
         )
     }
 
-    private func extractAlterTableName(from text: LexedStatement) -> String? {
+    private func extractAlterTableReference(from text: LexedStatement) -> TableReference? {
         let nsRange = NSRange(location: 0, length: text.code.length)
         guard let match = Self.alterTableRegex.firstMatch(in: text.codeString, range: nsRange) else { return nil }
-        let tableNSRange = match.range(at: 1)
-        guard tableNSRange.location != NSNotFound else { return nil }
-        return Self.unquoted(text.original.substring(with: tableNSRange))
+        let pathRange = match.range(at: 1)
+        guard pathRange.location != NSNotFound else { return nil }
+        let segments = Self.pathSegments(pathRange, in: text)
+        guard let table = segments.last, !table.name.isEmpty else { return nil }
+        return TableReference(
+            tableName: table.name,
+            alias: nil,
+            schema: segments.count >= 2 ? segments[segments.count - 2].name : nil,
+            isIdentifierQuoted: table.isQuoted
+        )
     }
 
     // MARK: - Clause Detection

apple/Packages/TableProCore/Sources/TableProAutocomplete/SchemaRelevanceRanker.swift

@@ -0,0 +1,200 @@
+import Foundation
+
+public struct SchemaRelevanceDocument: Sendable, Equatable {
+    public let name: String
+    public let details: [String]
+
+    public init(name: String, details: [String] = []) {
+        self.name = name
+        self.details = details
+    }
+}
+
+// Orders tables by how well their names and columns match a question: BM25 over identifier
+// words, with trigram matching for words that miss every exact term (typos, inflections).
+public enum SchemaRelevanceRanker {
+    static let nameWeight = 3
+    static let fullNameBoost = 5.0
+    static let fuzzyThreshold = 0.5
+    static let fuzzyMinLength = 4
+    static let k1 = 1.2
+    static let b = 0.75
+
+    public static func rank(_ documents: [SchemaRelevanceDocument], for query: String) -> [Int] {
+        let queryWords = Set(words(in: query))
+        guard !queryWords.isEmpty, !documents.isEmpty else { return [] }
+
+        let index = Index(documents: documents)
+        let matches = queryWords.flatMap { index.matches(for: $0) }
+        let orderedQueryWords = words(in: query)
+
+        var scored: [(offset: Int, score: Double)] = []
+        for (offset, terms) in index.termCounts.enumerated() {
+            var score = 0.0
+            for match in matches {
+                guard let frequency = terms[match.term] else { continue }
+                score += match.weight * index.bm25(frequency: frequency, idfTerm: match.term, length: index.lengths[offset])
+            }
+            let nameWords = index.nameWords[offset]
+            if nameWords.joined().count >= fuzzyMinLength, containsSequence(orderedQueryWords, nameWords) {
+                score += fullNameBoost
+            }
+            if score > 0 {
+                scored.append((offset, score))
+            }
+        }
+
+        return scored
+            .sorted { $0.score != $1.score ? $0.score > $1.score : $0.offset < $1.offset }
+            .map(\.offset)
+    }
+
+    private static func containsSequence(_ words: [String], _ sequence: [String]) -> Bool {
+        guard !sequence.isEmpty, words.count >= sequence.count else { return false }
+        return (0...(words.count - sequence.count)).contains { start in
+            Array(words[start..<(start + sequence.count)]) == sequence
+        }
+    }
+
+    static func words(in text: String) -> [String] {
+        var result: [String] = []
+        var current = ""
+        var previous: Character?
+
+        func flush() {
+            if !current.isEmpty {
+                result.append(stem(current))
+                current = ""
+            }
+        }
+
+        for character in text {
+            guard character.isLetter || character.isNumber else {
+                flush()
+                previous = nil
+                continue
+            }
+            if let previous, startsNewWord(previous: previous, current: character) {
+                flush()
+            }
+            current += fold(character)
+            previous = character
+        }
+        flush()
+        return result
+    }
+
+    // Folding leaves "đ" alone because Unicode treats it as its own letter, not a "d" with a mark.
+    private static func fold(_ character: Character) -> String {
+        if character == "đ" || character == "Đ" { return "d" }
+        return String(character).folding(options: [.diacriticInsensitive, .caseInsensitive, .widthInsensitive], locale: nil)
+    }
+
+    private static func startsNewWord(previous: Character, current: Character) -> Bool {
+        if previous.isLowercase, current.isUppercase { return true }
+        if previous.isLetter, current.isNumber { return true }
+        if previous.isNumber, current.isLetter { return true }
+        return false
+    }
+
+    static func stem(_ word: String) -> String {
+        guard word.count > 3, word.allSatisfy(\.isASCII) else { return word }
+        if word.hasSuffix("ies") { return String(word.dropLast(3)) + "y" }
+        if word.hasSuffix("sses") || word.hasSuffix("xes") || word.hasSuffix("ches") || word.hasSuffix("shes") {
+            return String(word.dropLast(2))
+        }
+        if word.hasSuffix("s"), !word.hasSuffix("ss"), !word.hasSuffix("us"), !word.hasSuffix("is") {
+            return String(word.dropLast())
+        }
+        return word
+    }
+
+    static func trigramSimilarity(_ lhs: String, _ rhs: String) -> Double {
+        trigramSimilarity(trigrams(lhs), trigrams(rhs))
+    }
+
+    private static func trigramSimilarity(_ left: Set<String>, _ right: Set<String>) -> Double {
+        guard !left.isEmpty, !right.isEmpty else { return 0 }
+        return Double(left.intersection(right).count) / Double(left.union(right).count)
+    }
+
+    private static func trigrams(_ word: String) -> Set<String> {
+        let padded = Array("  \(word) ")
+        guard padded.count >= 3 else { return [] }
+        return Set((0...(padded.count - 3)).map { String(padded[$0..<($0 + 3)]) })
+    }
+
+    private struct Match {
+        let term: String
+        let weight: Double
+    }
+
+    private struct Index {
+        let termCounts: [[String: Int]]
+        let nameWords: [[String]]
+        let lengths: [Int]
+        let documentFrequency: [String: Int]
+        let fuzzyTerms: [(term: String, trigrams: Set<String>)]
+        let averageLength: Double
+
+        init(documents: [SchemaRelevanceDocument]) {
+            var termCounts: [[String: Int]] = []
+            var nameWords: [[String]] = []
+            var lengths: [Int] = []
+            var documentFrequency: [String: Int] = [:]
+            termCounts.reserveCapacity(documents.count)
+
+            for document in documents {
+                let name = SchemaRelevanceRanker.words(in: document.name)
+                var counts: [String: Int] = [:]
+                for word in name {
+                    counts[word, default: 0] += SchemaRelevanceRanker.nameWeight
+                }
+                for detail in document.details {
+                    for word in SchemaRelevanceRanker.words(in: detail) {
+                        counts[word, default: 0] += 1
+                    }
+                }
+                for term in counts.keys {
+                    documentFrequency[term, default: 0] += 1
+                }
+                termCounts.append(counts)
+                nameWords.append(name)
+                lengths.append(counts.values.reduce(0, +))
+            }
+
+            self.termCounts = termCounts
+            self.nameWords = nameWords
+            self.lengths = lengths
+            self.documentFrequency = documentFrequency
+            self.fuzzyTerms = documentFrequency.keys
+                .filter { $0.count >= SchemaRelevanceRanker.fuzzyMinLength }
+                .map { ($0, SchemaRelevanceRanker.trigrams($0)) }
+            let total = lengths.reduce(0, +)
+            self.averageLength = documents.isEmpty ? 0 : max(Double(total) / Double(documents.count), 1)
+        }
+
+        func matches(for word: String) -> [Match] {
+            if documentFrequency[word] != nil {
+                return [Match(term: word, weight: 1)]
+            }
+            guard word.count >= SchemaRelevanceRanker.fuzzyMinLength else { return [] }
+            let wordTrigrams = SchemaRelevanceRanker.trigrams(word)
+            return fuzzyTerms.compactMap { entry in
+                let similarity = SchemaRelevanceRanker.trigramSimilarity(wordTrigrams, entry.trigrams)
+                return similarity >= SchemaRelevanceRanker.fuzzyThreshold ? Match(term: entry.term, weight: similarity) : nil
+            }
+        }
+
+        func bm25(frequency: Int, idfTerm: String, length: Int) -> Double {
+            let documentCount = Double(termCounts.count)
+            let containing = Double(documentFrequency[idfTerm] ?? 0)
+            let idf = log(1 + (documentCount - containing + 0.5) / (containing + 0.5))
+            let tf = Double(frequency)
+            let k1 = SchemaRelevanceRanker.k1
+            let b = SchemaRelevanceRanker.b
+            let norm = k1 * (1 - b + b * Double(length) / averageLength)
+            return idf * (tf * (k1 + 1)) / (tf + norm)
+        }
+    }
+}

apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SQLContextAnalyzerTableReferenceTests.swift

@@ -0,0 +1,49 @@
+import Foundation
+@testable import TableProAutocomplete
+import TableProSQLGrammar
+import Testing
+
+struct SQLContextAnalyzerTableReferenceTests {
+    private let analyzer = SQLContextAnalyzer()
+
+    private func tables(_ statement: String) -> [String] {
+        analyzer.tableReferences(in: statement, grammar: .ansi).map(\.tableName)
+    }
+
+    @Test("A statement that ends in a comment still names its tables")
+    func trailingComment() {
+        #expect(tables("SELECT * FROM orders -- recent orders") == ["orders"])
+        #expect(tables("SELECT * FROM orders /* recent */") == ["orders"])
+        #expect(analyzer.analyze(query: "SELECT * FROM orders -- recent", cursorPosition: 30).tableReferences.isEmpty)
+    }
+
+    @Test("A statement that ends in a string still names its tables")
+    func trailingLiteral() {
+        #expect(tables("SELECT * FROM orders WHERE note = 'open") == ["orders"])
+    }
+
+    @Test("A CTE that shares a table's name still names the table its body reads")
+    func shadowingCommonTableExpression() {
+        let references = analyzer.tableReferences(
+            in: "WITH orders AS (SELECT * FROM orders WHERE active) SELECT * FROM orders",
+            grammar: .ansi
+        )
+        #expect(references.contains { $0.tableName == "orders" && $0.derivedColumns == nil })
+    }
+
+    @Test("Joined, updated and altered tables are named")
+    func everyClause() {
+        #expect(tables("SELECT * FROM orders o JOIN customers c ON c.id = o.customer_id") == ["orders", "customers"])
+        #expect(tables("UPDATE customers SET name = 'x'").contains("customers"))
+        #expect(tables("ALTER TABLE invoices ADD COLUMN due DATE") == ["invoices"])
+    }
+
+    @Test("A schema-qualified ALTER target keeps its schema")
+    func qualifiedAlterTarget() {
+        let references = analyzer.tableReferences(in: "ALTER TABLE archive.orders ADD COLUMN due DATE", grammar: .ansi)
+        #expect(references.map(\.tableName) == ["orders"])
+        #expect(references.first?.schema == "archive")
+        let context = analyzer.analyze(query: "ALTER TABLE archive.orders DROP COLUMN ", cursorPosition: 39)
+        #expect(context.tableReferences.contains { $0.tableName == "orders" && $0.schema == "archive" })
+    }
+}

apple/Packages/TableProCore/Tests/TableProAutocompleteTests/SchemaRelevanceRankerTests.swift

@@ -0,0 +1,71 @@
+import Foundation
+@testable import TableProAutocomplete
+import Testing
+
+struct SchemaRelevanceRankerTests {
+    private let documents = [
+        SchemaRelevanceDocument(name: "audit_log", details: ["id", "created_at", "message"]),
+        SchemaRelevanceDocument(name: "customers", details: ["id", "email", "full_name"]),
+        SchemaRelevanceDocument(name: "order_items", details: ["order_id", "product_id", "quantity"]),
+        SchemaRelevanceDocument(name: "orders", details: ["id", "customer_id", "total_amount", "created_at"]),
+        SchemaRelevanceDocument(name: "products", details: ["id", "title", "price"]),
+        SchemaRelevanceDocument(name: "ProductCategory", details: ["id", "label"])
+    ]
+
+    private func names(_ query: String) -> [String] {
+        SchemaRelevanceRanker.rank(documents, for: query).map { documents[$0].name }
+    }
+
+    @Test("Splits snake case, camel case and digits into folded, stemmed words")
+    func splitsIdentifiers() {
+        #expect(SchemaRelevanceRanker.words(in: "order_items") == ["order", "item"])
+        #expect(SchemaRelevanceRanker.words(in: "ProductCategories") == ["product", "category"])
+        #expect(SchemaRelevanceRanker.words(in: "userId2fa") == ["user", "id", "2", "fa"])
+        #expect(SchemaRelevanceRanker.words(in: "Đơn hàng") == ["don", "hang"])
+    }
+
+    @Test("A table named in the question ranks first")
+    func namedTableRanksFirst() {
+        #expect(names("total amount of orders per customer").first == "orders")
+    }
+
+    @Test("A word that only a column holds still finds its table")
+    func columnWordFindsTable() {
+        #expect(names("which rows have an email").first == "customers")
+    }
+
+    @Test("A misspelled table name matches by trigrams")
+    func typoMatches() {
+        #expect(names("list all custmers").first == "customers")
+    }
+
+    @Test("A camel case table matches its plural words")
+    func camelCaseMatches() {
+        #expect(names("product categories").contains("ProductCategory"))
+    }
+
+    @Test("A question that matches nothing ranks nothing")
+    func noMatchIsEmpty() {
+        #expect(names("hello there").isEmpty)
+        #expect(SchemaRelevanceRanker.rank(documents, for: "").isEmpty)
+        #expect(SchemaRelevanceRanker.rank([], for: "orders").isEmpty)
+    }
+
+    @Test("A table name inside a longer word gets no name boost")
+    func nameInsideLongerWordIsNotBoosted() {
+        let documents = [
+            SchemaRelevanceDocument(name: "data"),
+            SchemaRelevanceDocument(name: "events", details: ["metadata"])
+        ]
+        #expect(SchemaRelevanceRanker.rank(documents, for: "show the metadata").first == 1)
+    }
+
+    @Test("Equal scores keep the listed order")
+    func tiesKeepListedOrder() {
+        let twins = [
+            SchemaRelevanceDocument(name: "b_events"),
+            SchemaRelevanceDocument(name: "a_events")
+        ]
+        #expect(SchemaRelevanceRanker.rank(twins, for: "events") == [0, 1])
+    }
+}