Skip to content

Commit ad82403

Browse files
datlechindatlechin
authored
fix(editor): make statement splitting dollar-quote aware for PostgreSQL (#1559) (#1560)
* fix(editor): make statement splitting dollar-quote aware for PostgreSQL (#1559) * refactor(editor): share the dollar-quote token-start guard across both SQL splitters --------- Signed-off-by: Ngô Quốc Đạt <datlechin@gmail.com>
1 parent 699cca4 commit ad82403

14 files changed

Lines changed: 207 additions & 85 deletions

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
3434
- The Delete shortcut in the data grid now follows a custom binding.
3535
- Find Next (Cmd+G) and Find Previous (Cmd+Shift+G) now work in the editor.
3636
- Pagination buttons no longer fire their page shortcut twice.
37+
- Running a PostgreSQL script with a `DO $$ ... $$` block or a dollar-quoted function body no longer fails with an unterminated dollar-quoted string error. (#1559)
3738
- AWS IAM connections no longer ask for a password on connect or reconnect. IAM supplies the credentials, so the prompt was never needed. The same now holds for any auth mode that replaces the password, such as a Postgres password file.
3839
- Oracle connection failures show the listener's actual reason (such as an unknown service name) instead of a generic "server closed the connection" message. (#483)
3940

‎TablePro/Core/AI/Chat/Tools/ConfirmDestructiveOperationChatTool.swift‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,14 +36,15 @@ struct ConfirmDestructiveOperationChatTool: ChatTool {
3636
isError: true
3737
)
3838
}
39-
guard !QueryClassifier.isMultiStatement(query) else {
39+
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
40+
41+
guard !QueryClassifier.isMultiStatement(query, databaseType: meta.databaseType) else {
4042
return ChatToolResult(
4143
content: "Multi-statement queries are not supported. Send one statement at a time.",
4244
isError: true
4345
)
4446
}
4547

46-
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
4748
let tier = QueryClassifier.classifyTier(query, databaseType: meta.databaseType)
4849
guard tier == .destructive else {
4950
return ChatToolResult(

‎TablePro/Core/AI/Chat/Tools/ExecuteQueryChatTool.swift‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,10 @@ struct ExecuteQueryChatTool: ChatTool {
4545
guard (query as NSString).length <= 102_400 else {
4646
return ChatToolResult(content: "Query exceeds 100KB limit", isError: true)
4747
}
48-
guard !QueryClassifier.isMultiStatement(query) else {
48+
49+
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
50+
51+
guard !QueryClassifier.isMultiStatement(query, databaseType: meta.databaseType) else {
4952
return ChatToolResult(
5053
content: "Multi-statement queries are not supported. Send one statement at a time.",
5154
isError: true
@@ -66,8 +69,6 @@ struct ExecuteQueryChatTool: ChatTool {
6669
clamp: 1...300
6770
) ?? mcpSettings.queryTimeoutSeconds
6871

69-
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
70-
7172
let tier = QueryClassifier.classifyTier(query, databaseType: meta.databaseType)
7273
if tier == .destructive {
7374
return ChatToolResult(

‎TablePro/Core/Coordinators/QueryExecutionCoordinator.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ final class QueryExecutionCoordinator {
2323
let fullQuery = tab.content.query
2424
guard !fullQuery.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
2525

26-
let statements = SQLStatementScanner.allStatements(in: fullQuery)
26+
let statements = SQLStatementScanner.allStatements(in: fullQuery, dialect: parent.sqlDialect)
2727
guard !statements.isEmpty else { return }
2828

2929
if AppSettingsManager.shared.editor.queryParametersEnabled {

‎TablePro/Core/MCP/Protocol/Tools/ConfirmDestructiveOperationTool.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,14 +57,14 @@ public struct ConfirmDestructiveOperationTool: MCPToolImplementation {
5757
)
5858
}
5959

60-
guard !QueryClassifier.isMultiStatement(query) else {
60+
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
61+
62+
guard !QueryClassifier.isMultiStatement(query, databaseType: meta.databaseType) else {
6163
throw MCPProtocolError.invalidParams(
6264
detail: "Multi-statement queries are not supported. Send one statement at a time."
6365
)
6466
}
6567

66-
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
67-
6868
let tier = QueryClassifier.classifyTier(query, databaseType: meta.databaseType)
6969
guard tier == .destructive else {
7070
throw MCPProtocolError.invalidParams(

‎TablePro/Core/MCP/Protocol/Tools/ExecuteQueryTool.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,9 @@ public struct ExecuteQueryTool: MCPToolImplementation {
7777
throw MCPProtocolError.invalidParams(detail: "Query exceeds 100KB limit")
7878
}
7979

80-
guard !QueryClassifier.isMultiStatement(query) else {
80+
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
81+
82+
guard !QueryClassifier.isMultiStatement(query, databaseType: meta.databaseType) else {
8183
throw MCPProtocolError.invalidParams(
8284
detail: "Multi-statement queries are not supported. Send one statement at a time."
8385
)
@@ -86,8 +88,6 @@ public struct ExecuteQueryTool: MCPToolImplementation {
8688
try await throwIfCancelled(context)
8789
await context.progress.emit(progress: 0.0, total: 1.0, message: "Connecting")
8890

89-
let meta = try await ToolConnectionMetadata.resolve(connectionId: connectionId)
90-
9191
if let database {
9292
_ = try await services.connectionBridge.switchDatabase(
9393
connectionId: connectionId,

‎TablePro/Core/Services/Execution/DefaultExecutionGate.swift‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ internal actor DefaultExecutionGate: ExecutionGate {
3030
let tier = request.sql.map { QueryClassifier.classifyTier($0, databaseType: request.databaseType) }
3131
let isDangerous = request.sql.map { QueryClassifier.isDangerousQuery($0, databaseType: request.databaseType) } ?? false
3232
let isDestructive = request.kind.declaresDestructive || tier == .destructive || isDangerous
33-
let isMultiStatement = request.sql.map { QueryClassifier.isMultiStatement($0) } ?? false
33+
let isMultiStatement = request.sql.map {
34+
QueryClassifier.isMultiStatement($0, databaseType: request.databaseType)
35+
} ?? false
3436
let effectiveWrite = await resolveEffectiveWrite(request, tier: tier)
3537

3638
if let denial = capabilityDenial(

‎TablePro/Core/Utilities/SQL/QueryClassifier.swift‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
//
55

66
import Foundation
7+
import TableProPluginKit
78

89
enum QueryTier {
910
case safe
@@ -129,8 +130,11 @@ enum QueryClassifier {
129130
return .safe
130131
}
131132

132-
static func isMultiStatement(_ sql: String) -> Bool {
133-
SQLStatementScanner.allStatements(in: sql).count > 1
133+
static func isMultiStatement(_ sql: String, databaseType: DatabaseType) -> Bool {
134+
SQLStatementScanner.allStatements(
135+
in: sql,
136+
dialect: SqlDialect.from(databaseTypeId: databaseType.rawValue)
137+
).count > 1
134138
}
135139

136140
static func isExplainStatement(_ sql: String) -> Bool {

‎TablePro/Core/Utilities/SQL/SQLFileParser.swift‎

Lines changed: 2 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -38,20 +38,6 @@ final class SQLFileParser: Sendable {
3838
private static let kCapitalE: unichar = 0x45
3939
private static let kSmallE: unichar = 0x65
4040

41-
private static func isIdentifierStart(_ ch: unichar) -> Bool {
42-
(ch >= 0x41 && ch <= 0x5A) || (ch >= 0x61 && ch <= 0x7A) || ch == 0x5F
43-
}
44-
45-
private static func isIdentifierPart(_ ch: unichar) -> Bool {
46-
isIdentifierStart(ch) || (ch >= 0x30 && ch <= 0x39)
47-
}
48-
49-
private enum DollarQuoteScan {
50-
case opener(length: Int, tag: String)
51-
case notOpener
52-
case needsMoreData
53-
}
54-
5541
nonisolated private static func needsLookahead(
5642
_ char: unichar,
5743
state: ParserState,
@@ -163,44 +149,6 @@ final class SQLFileParser: Sendable {
163149
}
164150
}
165151

166-
private static func scanDollarQuoteOpener(
167-
at pos: Int, in buffer: NSString, bufLen: Int
168-
) -> DollarQuoteScan {
169-
var p = pos + 1
170-
while p < bufLen {
171-
let ch = buffer.character(at: p)
172-
if ch == kDollar {
173-
let tagLen = p - pos - 1
174-
if tagLen == 0 {
175-
return .opener(length: 2, tag: "")
176-
}
177-
let firstChar = buffer.character(at: pos + 1)
178-
if !isIdentifierStart(firstChar) {
179-
return .notOpener
180-
}
181-
let tag = buffer.substring(with: NSRange(location: pos + 1, length: tagLen))
182-
return .opener(length: tagLen + 2, tag: tag)
183-
}
184-
if !isIdentifierPart(ch) {
185-
return .notOpener
186-
}
187-
p += 1
188-
}
189-
return .needsMoreData
190-
}
191-
192-
private static func matchesDollarClose(
193-
at pos: Int, tag: String, in buffer: NSString, bufLen: Int
194-
) -> Bool {
195-
let closeLen = (tag as NSString).length + 2
196-
guard pos + closeLen <= bufLen else { return false }
197-
if buffer.character(at: pos) != kDollar { return false }
198-
if buffer.character(at: pos + closeLen - 1) != kDollar { return false }
199-
if tag.isEmpty { return true }
200-
let tagRange = NSRange(location: pos + 1, length: (tag as NSString).length)
201-
return buffer.substring(with: tagRange) == tag
202-
}
203-
204152
private struct StepResult {
205153
var advanced: Bool
206154
var deferred: Bool
@@ -255,7 +203,7 @@ final class SQLFileParser: Sendable {
255203
}
256204

257205
if ctx.dialect.supportsDollarQuotes && char == kDollar {
258-
switch scanDollarQuoteOpener(at: i, in: nsBuffer, bufLen: bufLen) {
206+
switch SqlDollarQuote.scanOpener(at: i, in: nsBuffer, bufLen: bufLen) {
259207
case .opener(let length, let tag):
260208
(ctx.hasStatementContent, ctx.statementStartLine) = markContent(
261209
ctx.hasStatementContent, ctx.statementStartLine, ctx.currentLine)
@@ -452,7 +400,7 @@ final class SQLFileParser: Sendable {
452400
i = pos
453401
return StepResult(advanced: true, deferred: true)
454402
}
455-
if matchesDollarClose(at: pos, tag: ctx.dollarTag, in: nsBuffer, bufLen: bufLen) {
403+
if SqlDollarQuote.matchesClose(at: pos, tag: ctx.dollarTag, in: nsBuffer, bufLen: bufLen) {
456404
pos += closeLen
457405
ctx.state = .normal
458406
ctx.dollarTag = ""

‎TablePro/Core/Utilities/SQL/SQLStatementScanner.swift‎

Lines changed: 34 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,18 @@
44
//
55

66
import Foundation
7+
import TableProPluginKit
78

89
enum SQLStatementScanner {
910
struct LocatedStatement {
1011
let sql: String
1112
let offset: Int
1213
}
1314

14-
/// Returns statements with trailing semicolons stripped — for driver execution.
15-
static func allStatements(in sql: String) -> [String] {
15+
/// Returns statements with trailing semicolons stripped, for driver execution.
16+
static func allStatements(in sql: String, dialect: SqlDialect = .generic) -> [String] {
1617
var results: [String] = []
17-
scan(sql: sql, cursorPosition: nil) { rawSQL, _ in
18+
scan(sql: sql, cursorPosition: nil, dialect: dialect) { rawSQL, _ in
1819
var trimmed = rawSQL.trimmingCharacters(in: .whitespacesAndNewlines)
1920
if trimmed.hasSuffix(";") {
2021
trimmed = String(trimmed.dropLast())
@@ -28,7 +29,7 @@ enum SQLStatementScanner {
2829
return results
2930
}
3031

31-
/// Returns statements preserving trailing semicolons — for display/history/favorites.
32+
/// Returns statements preserving trailing semicolons, for display/history/favorites.
3233
static func allStatementsPreservingSemicolons(in sql: String) -> [String] {
3334
var results: [String] = []
3435
scan(sql: sql, cursorPosition: nil) { rawSQL, _ in
@@ -44,8 +45,8 @@ enum SQLStatementScanner {
4445
return results
4546
}
4647

47-
static func statementAtCursor(in sql: String, cursorPosition: Int) -> String {
48-
var result = locatedStatementAtCursor(in: sql, cursorPosition: cursorPosition)
48+
static func statementAtCursor(in sql: String, cursorPosition: Int, dialect: SqlDialect = .generic) -> String {
49+
var result = locatedStatementAtCursor(in: sql, cursorPosition: cursorPosition, dialect: dialect)
4950
.sql
5051
.trimmingCharacters(in: .whitespacesAndNewlines)
5152
if result.hasSuffix(";") {
@@ -55,9 +56,9 @@ enum SQLStatementScanner {
5556
return result
5657
}
5758

58-
static func locatedStatementAtCursor(in sql: String, cursorPosition: Int) -> LocatedStatement {
59+
static func locatedStatementAtCursor(in sql: String, cursorPosition: Int, dialect: SqlDialect = .generic) -> LocatedStatement {
5960
var result = LocatedStatement(sql: "", offset: 0)
60-
scan(sql: sql, cursorPosition: cursorPosition) { rawSQL, offset in
61+
scan(sql: sql, cursorPosition: cursorPosition, dialect: dialect) { rawSQL, offset in
6162
result = LocatedStatement(sql: rawSQL, offset: offset)
6263
return false
6364
}
@@ -75,10 +76,12 @@ enum SQLStatementScanner {
7576
private static let star = UInt16(UnicodeScalar("*").value)
7677
private static let newline = UInt16(UnicodeScalar("\n").value)
7778
private static let backslash = UInt16(UnicodeScalar("\\").value)
79+
private static let dollar = UInt16(UnicodeScalar("$").value)
7880

7981
private static func scan(
8082
sql: String,
8183
cursorPosition: Int?,
84+
dialect: SqlDialect = .generic,
8285
onStatement: (_ rawSQL: String, _ offset: Int) -> Bool
8386
) {
8487
let nsQuery = sql as NSString
@@ -97,6 +100,9 @@ enum SQLStatementScanner {
97100
var stringCharVal: UInt16 = 0
98101
var inLineComment = false
99102
var inBlockComment = false
103+
var inDollarQuote = false
104+
var dollarTag = ""
105+
let dollarQuotesEnabled = dialect.supportsDollarQuotes
100106
var i = 0
101107

102108
while i < length {
@@ -118,6 +124,18 @@ enum SQLStatementScanner {
118124
continue
119125
}
120126

127+
if inDollarQuote {
128+
if ch == dollar,
129+
SqlDollarQuote.matchesClose(at: i, tag: dollarTag, in: nsQuery, bufLen: length) {
130+
inDollarQuote = false
131+
i += (dollarTag as NSString).length + 2
132+
dollarTag = ""
133+
continue
134+
}
135+
i += 1
136+
continue
137+
}
138+
121139
if !inString && ch == dash && i + 1 < length && nsQuery.character(at: i + 1) == dash {
122140
inLineComment = true
123141
i += 2
@@ -148,6 +166,14 @@ enum SQLStatementScanner {
148166
}
149167
}
150168

169+
if dollarQuotesEnabled, !inString, ch == dollar,
170+
case .opener(let openerLength, let tag) = SqlDollarQuote.scanOpener(at: i, in: nsQuery, bufLen: length) {
171+
inDollarQuote = true
172+
dollarTag = tag
173+
i += openerLength
174+
continue
175+
}
176+
151177
if ch == semicolonChar && !inString {
152178
let stmtEnd = i + 1
153179

0 commit comments

Comments
 (0)