Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -133,8 +133,11 @@ public final class SQLCompletionProvider {
guard let catalog else { return [] }
if let derived = context.tableReferences.first(where: {
$0.isDerived && $0.identifier.caseInsensitiveCompare(dotPrefix) == .orderedSame
}), let columns = derived.derivedColumns, !columns.isEmpty {
return columns.map { SQLCompletionItem.column($0, dataType: nil, tableName: derived.identifier) }
}) {
let columns = await derivedColumns(of: derived)
if !columns.isEmpty {
return columns.map { SQLCompletionItem.column($0, dataType: nil, tableName: derived.identifier) }
}
}
if let tableName = await catalog.resolveAlias(dotPrefix, in: context.tableReferences) {
let schema = context.tableReferences.first {
Expand Down Expand Up @@ -286,22 +289,25 @@ public final class SQLCompletionProvider {
])

case .orderBy:
items += await columnItems(for: context.tableReferences)
let columns = await columnItems(for: context.tableReferences)
// Only a closed quote stays outside the replaced text; an open one is replaced with the token.
items += outputAliasItems(context.outputAliases, besides: columns, insideQuotes: context.quotedName?.close != nil)
items += columns
items += filterKeywords([
"ASC", "DESC", "NULLS FIRST", "NULLS LAST",
"LIMIT", "OFFSET",
"UNION", "INTERSECT", "EXCEPT"
])

case .set:
if let firstTable = context.tableReferences.first {
items = await catalog?.columnCompletionItems(for: firstTable.tableName, schema: firstTable.schema) ?? []
if let target = context.targetTable ?? context.tableReferences.first {
items = await catalog?.columnCompletionItems(for: target.tableName, schema: target.schema) ?? []
}
items += filterKeywords(["WHERE", "RETURNING"])

case .insertColumns:
if let firstTable = context.tableReferences.first {
items = await catalog?.columnCompletionItems(for: firstTable.tableName, schema: firstTable.schema) ?? []
if let target = context.targetTable ?? context.tableReferences.first {
items = await catalog?.columnCompletionItems(for: target.tableName, schema: target.schema) ?? []
}

case .values:
Expand Down Expand Up @@ -534,7 +540,59 @@ public final class SQLCompletionProvider {
if references.isEmpty {
return await catalog?.allColumnsFromCachedTables() ?? []
}
return await catalog?.allColumnsInScope(for: references) ?? []
var resolved: [TableReference] = []
for reference in references {
guard reference.derivedProjection != nil else {
resolved.append(reference)
continue
}
resolved.append(reference.resolvingDerivedColumns(await derivedColumns(of: reference)))
}
return await catalog?.allColumnsInScope(for: resolved) ?? []
}

// ORDER BY sorts by an output alias before a source column of the same name, so the alias leads.
private func outputAliasItems(
_ spellings: [String], besides columns: [SQLCompletionItem], insideQuotes: Bool
) -> [SQLCompletionItem] {
let columnNames = Set(columns.map(\.label))
let foldedColumnNames = Set(columns.map { $0.label.lowercased() })
var seen = Set<String>()
return spellings.compactMap { spelling in
let name = SQLTokenStream.unquoted(spelling)
// A quoted alias keeps its case, so `"Status"` is another name than a `status` column.
let isQuoted = name != spelling
let repeatsColumn = isQuoted ? columnNames.contains(name) : foldedColumnNames.contains(name.lowercased())
guard !name.isEmpty, !repeatsColumn, seen.insert(spelling).inserted else { return nil }
return SQLCompletionItem(
label: name,
kind: .column,
insertText: insideQuotes ? name : identifierSpelling(of: spelling, name: name),
sortPriority: 10
)
}
}

// `AS 'Total'` names the column on MySQL and SQLite, but ORDER BY 'Total' sorts by a constant.
private func identifierSpelling(of spelling: String, name: String) -> String {
guard spelling.hasPrefix("'") else { return spelling }
let isBare = name.unicodeScalars.first.map { CharacterSet.letters.contains($0) || $0 == "_" } == true
&& name.unicodeScalars.allSatisfy { CharacterSet.alphanumerics.contains($0) || $0 == "_" }
guard !isBare || SQLKeywords.isKeyword(name) else { return name }
// The analyzer reads a backslash-escaped `"` as a string, which is where a backtick quotes a name instead.
let quote = grammar.isSuperset(of: [.backslashEscapesInDoubleQuotes, .backtickQuotes]) ? "`" : "\""
return quote + name.replacingOccurrences(of: quote, with: quote + quote) + quote
}

// CTEs were bound to their own scope by the parser, so what is left to expand is a table the catalog knows.
private func derivedColumns(of reference: TableReference) async -> [String] {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

derivedColumns(of:) awaits the catalog sequentially per table, and columnItems(for:) calls it once per reference. The same tables are queried again in allColumnsInScope. N derived tables using * mean N sequential awaits plus repeated fetches on every SELECT-list or WHERE completion. Use a TaskGroup or batch the lookups and reuse the results.

guard let projection = reference.derivedProjection else { return reference.derivedColumns ?? [] }
var expanded: [DerivedProjection.Table: [String]] = [:]
for table in projection.tables where expanded[table] == nil {
expanded[table] = await catalog?.columnCompletionItems(for: table.name, schema: table.schema).map(\.label) ?? []
}
// The catalog answers no columns for a table it does not know, which must not count as a zero-width `*`.
return projection.columns { table in expanded[table].flatMap { $0.isEmpty ? nil : $0 } }
}

// A reference is spelled as the statement wrote it, and stays quoted when it was, or it would fold to another name.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@ public struct TableReference: Hashable, Sendable {
public let derivedColumns: [String]?
// Whether `identifier` was quoted in the statement, since a qualifier inserted bare would fold to another name.
public let isIdentifierQuoted: Bool
// Set while a `*` in the derived table still needs the catalog; the provider resolves it into `derivedColumns`.
let derivedProjection: DerivedProjection?

public init(
tableName: String,
Expand All @@ -66,6 +68,26 @@ public struct TableReference: Hashable, Sendable {
self.schema = schema
self.derivedColumns = derivedColumns
self.isIdentifierQuoted = isIdentifierQuoted
self.derivedProjection = nil
}

init(tableName: String, alias: String?, schema: String?, projection: DerivedProjection, isIdentifierQuoted: Bool) {
self.tableName = tableName
self.alias = alias
self.schema = schema
self.derivedColumns = projection.columns { _ in nil }
self.isIdentifierQuoted = isIdentifierQuoted
self.derivedProjection = projection.tables.isEmpty ? nil : projection
}

func resolvingDerivedColumns(_ columns: [String]) -> TableReference {
TableReference(
tableName: tableName,
alias: alias,
schema: schema,
derivedColumns: columns,
isIdentifierQuoted: isIdentifierQuoted
)
}

public var identifier: String {
Expand Down Expand Up @@ -105,6 +127,8 @@ public struct SQLContext: Sendable {
public let expectsObjectName: Bool // The table-operand slot of FROM, JOIN or INTO
public let comparisonColumn: String? // The column on the other side of a comparison the cursor is a value in
public private(set) var quotedName: SQLQuotedName?
public let outputAliases: [String] // SELECT-list aliases an ORDER BY at the cursor may name, as written
public let targetTable: TableReference? // The table the INSERT, UPDATE or MERGE at the cursor writes

public init(
clauseType: SQLClauseType,
Expand All @@ -120,7 +144,9 @@ public struct SQLContext: Sendable {
isAfterComma: Bool = false,
expectsObjectName: Bool = false,
comparisonColumn: String? = nil,
quotedName: SQLQuotedName? = nil
quotedName: SQLQuotedName? = nil,
outputAliases: [String] = [],
targetTable: TableReference? = nil
) {
self.clauseType = clauseType
self.prefix = prefix
Expand All @@ -136,6 +162,8 @@ public struct SQLContext: Sendable {
self.expectsObjectName = expectsObjectName
self.comparisonColumn = comparisonColumn
self.quotedName = quotedName
self.outputAliases = outputAliases
self.targetTable = targetTable
}

public func replacingTableReferences(_ references: [TableReference]) -> SQLContext {
Expand Down Expand Up @@ -184,6 +212,7 @@ public final class SQLContextAnalyzer {
// The linear DML clauses are left to `scanLinearClause`: a fixed-priority sweep cannot tell the clause nearest the
// cursor from an earlier one. Order matters here, the more specific pattern first.
private static let structuralClauseRegexes: [(regex: NSRegularExpression, clause: SQLClauseType)] = {
let insertTarget = "\(quotedSegment)(?:\\.\(quotedSegment))*(?:\\s+AS\\s+\(quotedSegment))?"
let patterns: [(String, SQLClauseType)] = [
("\\bADD\\s+(?:COLUMN\\s+)?[`\"']?\\w+[`\"']?\\s+\\w+.*?\\b(?:AFTER|BEFORE)(?:\\s+\\w*)?$",
.alterTableColumn),
Expand Down Expand Up @@ -217,19 +246,13 @@ public final class SQLContextAnalyzer {
("\\b(LIMIT|OFFSET)\\s+\\d*$", .limit),
("\\bVALUES\\s*(?:\\([^)]*\\)\\s*+(?:,\\s*+)?)+\\w*$", .values),
("\\bVALUES\\s*\\([^)]*$", .values),
("\\bINSERT\\s+INTO\\s+\\w+\\s*\\([^)]*$", .insertColumns),
("\\bINSERT\\s+INTO\\s+\(insertTarget)\\s*\\([^)]*$", .insertColumns),
]
return patterns.map { pattern, clause in
(compileRegex(pattern, options: .caseInsensitive), clause)
}
}()

private static let cteFirstRegex = compileRegex(
"(?i)\\bWITH\\s+(?:RECURSIVE\\s+)?([\\w]+)\\s+AS\\s*\\("
)

private static let cteCommaRegex = compileRegex("(?i),\\s*([\\w]+)\\s+AS\\s*\\(")

// Read on the code view, where a quoted name's content is masked to word characters, so a segment may be quoted
// and hold any character.
private static let quotedSegment = "[`\"'\\[]?\\w+[`\"'\\]]?"
Expand Down Expand Up @@ -298,10 +321,10 @@ public final class SQLContextAnalyzer {

var tableReferences = extractTableReferences(from: text.forReferences)

let cteNames = extractCTENames(from: text)

let derivedTables = Self.derivedTableParser.parse(located.sql)
mergeDerivedTables(derivedTables, cteNames: cteNames, into: &tableReferences)
let tokens = Self.tokens(of: text, grammar: grammar)
let derivedTables = Self.derivedTableParser.visibleTables(in: tokens, cursor: cursor)
let cteNames = derivedTables.filter(\.isCommonTableExpression).map(\.alias)
mergeDerivedTables(derivedTables, into: &tableReferences)

var seenReferences = Set<TableReference>(tableReferences)

Expand Down Expand Up @@ -340,6 +363,14 @@ public final class SQLContextAnalyzer {

let isCastTarget = endsWithCastOperator(text.code, before: clauseEnd)
let comparisonColumn = comparisonTarget(in: text, before: clauseEnd)
let outputAliases = resolution.clause == .orderBy
? Self.derivedTableParser.outputAliases(in: tokens, cursor: clauseEnd)
: []
let targetTable = resolution.clause == .set || resolution.clause == .insertColumns
? Self.derivedTableParser.targetTable(in: tokens, cursor: clauseEnd).map {
writtenTable($0, among: tableReferences)
}
: nil

return SQLContext(
clauseType: isCastTarget ? .castTarget : resolution.clause,
Expand All @@ -355,13 +386,24 @@ public final class SQLContextAnalyzer {
isAfterComma: isAfterComma,
expectsObjectName: resolution.expectsObjectName,
comparisonColumn: comparisonColumn,
quotedName: quotedName
quotedName: quotedName,
outputAliases: outputAliases,
targetTable: targetTable
)
}

// MARK: - Lexing

private enum CursorState: Equatable {
static func tokens(of text: LexedStatement, grammar: SQLLexicalGrammar) -> SQLTokenStream {
let references = text.forReferences
return SQLTokenStream(
code: references.code,
original: references.original,
bracketQuotes: grammar.contains(.bracketQuotedIdentifiers)
)
}

enum CursorState: Equatable {
case code
case comment
case literal
Expand All @@ -370,7 +412,7 @@ public final class SQLContextAnalyzer {

// `code` blanks every comment and literal body and masks every quoted name's content to `_`, offsets kept, so a
// scan or pattern reads structure from it and a name from `original` at the same range.
private struct LexedStatement {
struct LexedStatement {
let original: NSString
let code: NSString
let cursor: CursorState
Expand All @@ -392,7 +434,7 @@ public final class SQLContextAnalyzer {
}
}

private static func lex(_ statement: NSString, cursor: Int, grammar: SQLLexicalGrammar) -> LexedStatement {
static func lex(_ statement: NSString, cursor: Int, grammar: SQLLexicalGrammar) -> LexedStatement {
let length = statement.length
var units = [UInt16](repeating: 0, count: length)
statement.getCharacters(&units, range: NSRange(location: 0, length: length))
Expand Down Expand Up @@ -525,31 +567,6 @@ public final class SQLContextAnalyzer {
}
}

// MARK: - CTE Support

private func extractCTENames(from text: LexedStatement) -> [String] {
var cteNames: [String] = []
let nsRange = NSRange(location: 0, length: text.code.length)

if let match = Self.cteFirstRegex.firstMatch(in: text.codeString, range: nsRange) {
let nameNSRange = match.range(at: 1)
if nameNSRange.location != NSNotFound {
cteNames.append(text.original.substring(with: nameNSRange))
}
}

Self.cteCommaRegex.enumerateMatches(in: text.codeString, range: nsRange) { match, _, _ in
if let match = match {
let nameNSRange = match.range(at: 1)
if nameNSRange.location != NSNotFound {
cteNames.append(text.original.substring(with: nameNSRange))
}
}
}

return cteNames
}

// MARK: - Subquery Support

private func calculateNestingLevel(in code: NSString, end: Int) -> Int {
Expand Down Expand Up @@ -751,7 +768,7 @@ public final class SQLContextAnalyzer {

private static let tableRefKeywords: Set<String> = [
"LEFT", "RIGHT", "INNER", "OUTER", "FULL", "CROSS", "NATURAL",
"JOIN", "ON", "AND", "OR", "WHERE", "SELECT", "FROM", "AS", "SET"
"JOIN", "ON", "AND", "OR", "WHERE", "SELECT", "FROM", "AS", "SET", "LATERAL"
]

private static let bareIdentifierRegex = compileRegex("^\\w+$")
Expand Down Expand Up @@ -889,25 +906,31 @@ public final class SQLContextAnalyzer {

// Columns attach to a matching reference, such as a CTE aliased in FROM, and the derived table or CTE is still
// registered under its own name so it resolves by either identifier.
private func mergeDerivedTables(
_ derivedTables: [DerivedTable],
cteNames: [String],
into references: inout [TableReference]
) {
let derivedColumnsByAlias = Dictionary(
derivedTables.map { ($0.alias.lowercased(), $0.columns) },
// SQL Server writes `UPDATE o SET ... FROM orders o`, naming the table by its FROM alias.
private func writtenTable(_ target: TableReference, among references: [TableReference]) -> TableReference {
guard target.schema == nil, let aliased = references.first(where: {
!$0.isDerived && $0.alias?.caseInsensitiveCompare(target.tableName) == .orderedSame
}) else { return target }
return aliased
}

private func mergeDerivedTables(_ derivedTables: [DerivedTable], into references: inout [TableReference]) {
let projectionsByAlias = Dictionary(
derivedTables.map { ($0.alias.lowercased(), $0.projection) },
uniquingKeysWith: { first, _ in first }
)

for index in references.indices {
let ref = references[index]
// A schema-qualified name is a real table even when a CTE shares its name.
guard ref.schema == nil else { continue }
for key in [ref.tableName.lowercased(), ref.identifier.lowercased()] {
guard let columns = derivedColumnsByAlias[key] else { continue }
guard let projection = projectionsByAlias[key] else { continue }
references[index] = TableReference(
tableName: ref.tableName,
alias: ref.alias,
schema: ref.schema,
derivedColumns: columns,
schema: nil,
projection: projection,
isIdentifierQuoted: ref.isIdentifierQuoted
)
break
Expand All @@ -916,12 +939,13 @@ public final class SQLContextAnalyzer {

var present = Set(references.map { $0.identifier.lowercased() })
for derived in derivedTables where present.insert(derived.alias.lowercased()).inserted {
references.append(
TableReference(tableName: derived.alias, alias: derived.alias, derivedColumns: derived.columns)
)
}
for cteName in cteNames where present.insert(cteName.lowercased()).inserted {
references.append(TableReference(tableName: cteName, alias: nil))
references.append(TableReference(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Every visible CTE is appended to tableReferences with its columns, even when the query never reads it. WITH a AS (SELECT x FROM t), b AS (SELECT y FROM u) SELECT | FROM a now offers y from b. Only derived tables used in FROM or JOIN should add unqualified columns. Unreferenced CTEs should stay resolvable for b. and FROM completion only.

tableName: derived.alias,
alias: derived.alias,
schema: nil,
projection: derived.projection,
isIdentifierQuoted: false
))
}
}

Expand Down
Loading
Loading