SwiftDBAI: natural language queries for any SQLite database
Drop-in SwiftUI chat view, headless ChatEngine, LLM-agnostic via AnyLanguageModel. Read-only by default with configurable allowlists. Robust SQL parser with 63 tests. Includes demo app with GitHub stars dataset.
This commit is contained in:
Krishna Kumar
629
Tests/SwiftDBAITests/SQLQueryParserTests.swift
Normal file
629
Tests/SwiftDBAITests/SQLQueryParserTests.swift
Normal file
@@ -0,0 +1,629 @@
|
||||
// SQLQueryParserTests.swift
|
||||
// SwiftDBAITests
|
||||
|
||||
import Testing
|
||||
@testable import SwiftDBAI
|
||||
|
||||
@Suite("SQLQueryParser")
|
||||
struct SQLQueryParserTests {
|
||||
|
||||
let readOnlyParser = SQLQueryParser(allowlist: .readOnly)
|
||||
let standardParser = SQLQueryParser(allowlist: .standard)
|
||||
let unrestrictedParser = SQLQueryParser(allowlist: .unrestricted)
|
||||
|
||||
// MARK: - Extraction from code blocks
|
||||
|
||||
@Test("Extracts SQL from markdown sql code block")
|
||||
func extractFromSQLCodeBlock() throws {
|
||||
let text = """
|
||||
Here's the query to find the top users:
|
||||
|
||||
```sql
|
||||
SELECT name, COUNT(*) as count FROM users GROUP BY name ORDER BY count DESC
|
||||
```
|
||||
|
||||
This will give you the results.
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT name, COUNT(*) as count FROM users GROUP BY name ORDER BY count DESC")
|
||||
#expect(result.operation == .select)
|
||||
#expect(result.requiresConfirmation == false)
|
||||
}
|
||||
|
||||
@Test("Extracts SQL from generic code block")
|
||||
func extractFromGenericCodeBlock() throws {
|
||||
let text = """
|
||||
Here you go:
|
||||
|
||||
```
|
||||
SELECT * FROM products WHERE price > 100
|
||||
```
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM products WHERE price > 100")
|
||||
}
|
||||
|
||||
@Test("Extracts SQL from labeled text")
|
||||
func extractFromLabel() throws {
|
||||
let text = """
|
||||
I can help with that.
|
||||
SQL: SELECT id, name FROM categories WHERE active = 1
|
||||
That should work.
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT id, name FROM categories WHERE active = 1")
|
||||
}
|
||||
|
||||
@Test("Extracts direct SQL from plain text")
|
||||
func extractDirectSQL() throws {
|
||||
let text = "SELECT COUNT(*) FROM orders WHERE status = 'shipped'"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT COUNT(*) FROM orders WHERE status = 'shipped'")
|
||||
}
|
||||
|
||||
@Test("Handles SQL with trailing semicolons")
|
||||
func trailingSemicolon() throws {
|
||||
let text = "```sql\nSELECT * FROM users;\n```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Handles multiline SQL in code block")
|
||||
func multilineSQL() throws {
|
||||
let text = """
|
||||
```sql
|
||||
SELECT u.name, COUNT(o.id) as order_count
|
||||
FROM users u
|
||||
JOIN orders o ON u.id = o.user_id
|
||||
GROUP BY u.name
|
||||
ORDER BY order_count DESC
|
||||
LIMIT 10
|
||||
```
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("SELECT u.name"))
|
||||
#expect(result.sql.contains("LIMIT 10"))
|
||||
}
|
||||
|
||||
@Test("Handles WITH (CTE) queries as SELECT")
|
||||
func cteQuery() throws {
|
||||
let text = """
|
||||
```sql
|
||||
WITH top_users AS (
|
||||
SELECT user_id, COUNT(*) as cnt FROM orders GROUP BY user_id
|
||||
)
|
||||
SELECT * FROM top_users WHERE cnt > 5
|
||||
```
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.operation == .select)
|
||||
}
|
||||
|
||||
// MARK: - No SQL found
|
||||
|
||||
@Test("Throws noSQLFound for text without SQL")
|
||||
func noSQLFound() throws {
|
||||
let text = "I'm sorry, I can't help with that request."
|
||||
#expect(throws: SQLParsingError.noSQLFound) {
|
||||
try readOnlyParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Throws noSQLFound for empty input")
|
||||
func emptyInput() throws {
|
||||
#expect(throws: SQLParsingError.noSQLFound) {
|
||||
try readOnlyParser.parse("")
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Operation detection
|
||||
|
||||
@Test("Detects INSERT operation")
|
||||
func detectInsert() throws {
|
||||
let text = "```sql\nINSERT INTO users (name) VALUES ('Alice')\n```"
|
||||
let result = try standardParser.parse(text)
|
||||
#expect(result.operation == .insert)
|
||||
}
|
||||
|
||||
@Test("Detects UPDATE operation")
|
||||
func detectUpdate() throws {
|
||||
let text = "```sql\nUPDATE users SET name = 'Bob' WHERE id = 1\n```"
|
||||
let result = try standardParser.parse(text)
|
||||
#expect(result.operation == .update)
|
||||
}
|
||||
|
||||
@Test("Detects DELETE operation and requires confirmation")
|
||||
func detectDeleteRequiresConfirmation() throws {
|
||||
let text = "```sql\nDELETE FROM users WHERE id = 99\n```"
|
||||
let result = try unrestrictedParser.parse(text)
|
||||
#expect(result.operation == .delete)
|
||||
#expect(result.requiresConfirmation == true)
|
||||
}
|
||||
|
||||
// MARK: - Allowlist enforcement
|
||||
|
||||
@Test("Rejects INSERT on read-only allowlist")
|
||||
func rejectInsertOnReadOnly() throws {
|
||||
let text = "```sql\nINSERT INTO users (name) VALUES ('Mallory')\n```"
|
||||
#expect(throws: SQLParsingError.operationNotAllowed(.insert)) {
|
||||
try readOnlyParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Rejects UPDATE on read-only allowlist")
|
||||
func rejectUpdateOnReadOnly() {
|
||||
let text = "```sql\nUPDATE users SET name = 'Eve' WHERE id = 1\n```"
|
||||
#expect(throws: SQLParsingError.operationNotAllowed(.update)) {
|
||||
try readOnlyParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Rejects DELETE on standard allowlist")
|
||||
func rejectDeleteOnStandard() {
|
||||
let text = "```sql\nDELETE FROM users WHERE id = 1\n```"
|
||||
#expect(throws: SQLParsingError.operationNotAllowed(.delete)) {
|
||||
try standardParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Dangerous operations
|
||||
|
||||
@Test("Rejects DROP TABLE")
|
||||
func rejectDrop() {
|
||||
let text = "```sql\nDROP TABLE users\n```"
|
||||
#expect(throws: SQLParsingError.dangerousOperation("DROP")) {
|
||||
try unrestrictedParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Rejects ALTER TABLE")
|
||||
func rejectAlter() {
|
||||
let text = "```sql\nALTER TABLE users ADD COLUMN age INTEGER\n```"
|
||||
#expect(throws: SQLParsingError.dangerousOperation("ALTER")) {
|
||||
try unrestrictedParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Rejects PRAGMA")
|
||||
func rejectPragma() {
|
||||
let text = "```sql\nPRAGMA table_info(users)\n```"
|
||||
#expect(throws: SQLParsingError.dangerousOperation("PRAGMA")) {
|
||||
try unrestrictedParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Does not match dangerous keywords inside identifiers")
|
||||
func noFalsePositiveOnSubstring() throws {
|
||||
// "DROPDOWN" contains "DROP" as substring but is not the keyword
|
||||
let text = "SELECT dropdown_value FROM settings"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("dropdown_value"))
|
||||
}
|
||||
|
||||
// MARK: - Multiple statements
|
||||
|
||||
@Test("Rejects multiple statements separated by semicolons")
|
||||
func rejectMultipleStatements() {
|
||||
let text = "```sql\nSELECT * FROM users; SELECT * FROM orders\n```"
|
||||
#expect(throws: SQLParsingError.multipleStatements) {
|
||||
try readOnlyParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Allows semicolons inside string literals")
|
||||
func allowSemicolonInString() throws {
|
||||
let text = "SELECT * FROM users WHERE bio = 'hello; world'"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("hello; world"))
|
||||
}
|
||||
|
||||
// MARK: - ParsedSQL equality
|
||||
|
||||
@Test("ParsedSQL equality works")
|
||||
func parsedSQLEquality() {
|
||||
let a = ParsedSQL(sql: "SELECT 1", operation: .select)
|
||||
let b = ParsedSQL(sql: "SELECT 1", operation: .select)
|
||||
#expect(a == b)
|
||||
}
|
||||
|
||||
// MARK: - Error descriptions
|
||||
|
||||
@Test("Error descriptions are meaningful")
|
||||
func errorDescriptions() {
|
||||
#expect(SQLParsingError.noSQLFound.description.contains("No SQL"))
|
||||
#expect(SQLParsingError.operationNotAllowed(.insert).description.contains("INSERT"))
|
||||
#expect(SQLParsingError.dangerousOperation("DROP").description.contains("DROP"))
|
||||
#expect(SQLParsingError.multipleStatements.description.contains("single"))
|
||||
}
|
||||
|
||||
// MARK: - MutationPolicy integration
|
||||
|
||||
@Test("MutationPolicy allows INSERT on permitted table")
|
||||
func mutationPolicyAllowsInsertOnPermittedTable() throws {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.insert, .update],
|
||||
allowedTables: ["orders", "order_items"]
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "```sql\nINSERT INTO orders (product, qty) VALUES ('Widget', 3)\n```"
|
||||
let result = try parser.parse(text)
|
||||
#expect(result.operation == .insert)
|
||||
#expect(result.requiresConfirmation == false)
|
||||
}
|
||||
|
||||
@Test("MutationPolicy rejects INSERT on non-permitted table")
|
||||
func mutationPolicyRejectsInsertOnForbiddenTable() {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.insert, .update],
|
||||
allowedTables: ["orders"]
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "```sql\nINSERT INTO users (name) VALUES ('Alice')\n```"
|
||||
#expect(throws: SQLParsingError.tableNotAllowed(table: "users", operation: .insert)) {
|
||||
try parser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("MutationPolicy rejects UPDATE on non-permitted table")
|
||||
func mutationPolicyRejectsUpdateOnForbiddenTable() {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.insert, .update],
|
||||
allowedTables: ["orders"]
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "```sql\nUPDATE users SET name = 'Bob' WHERE id = 1\n```"
|
||||
#expect(throws: SQLParsingError.tableNotAllowed(table: "users", operation: .update)) {
|
||||
try parser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("MutationPolicy rejects DELETE on non-permitted table")
|
||||
func mutationPolicyRejectsDeleteOnForbiddenTable() {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.insert, .update, .delete],
|
||||
allowedTables: ["temp_data"]
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "```sql\nDELETE FROM users WHERE id = 99\n```"
|
||||
#expect(throws: SQLParsingError.tableNotAllowed(table: "users", operation: .delete)) {
|
||||
try parser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("MutationPolicy allows mutation on any table when allowedTables is nil")
|
||||
func mutationPolicyAllowsAllTablesWhenNil() throws {
|
||||
let policy = MutationPolicy(allowedOperations: [.insert, .update])
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "```sql\nINSERT INTO any_table (col) VALUES ('val')\n```"
|
||||
let result = try parser.parse(text)
|
||||
#expect(result.operation == .insert)
|
||||
}
|
||||
|
||||
@Test("MutationPolicy SELECT is never restricted by table allowlist")
|
||||
func mutationPolicySelectIgnoresTableRestrictions() throws {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.insert],
|
||||
allowedTables: ["orders"]
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
// SELECT from a table NOT in allowedTables — should still work
|
||||
let text = "```sql\nSELECT * FROM users\n```"
|
||||
let result = try parser.parse(text)
|
||||
#expect(result.operation == .select)
|
||||
#expect(result.requiresConfirmation == false)
|
||||
}
|
||||
|
||||
@Test("MutationPolicy DELETE requires confirmation by default")
|
||||
func mutationPolicyDeleteRequiresConfirmation() throws {
|
||||
let policy = MutationPolicy(allowedOperations: [.delete])
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "```sql\nDELETE FROM users WHERE id = 1\n```"
|
||||
let result = try parser.parse(text)
|
||||
#expect(result.operation == .delete)
|
||||
#expect(result.requiresConfirmation == true)
|
||||
}
|
||||
|
||||
@Test("MutationPolicy DELETE skips confirmation when configured")
|
||||
func mutationPolicyDeleteNoConfirmation() throws {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.delete],
|
||||
requiresDestructiveConfirmation: false
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "```sql\nDELETE FROM users WHERE id = 1\n```"
|
||||
let result = try parser.parse(text)
|
||||
#expect(result.operation == .delete)
|
||||
#expect(result.requiresConfirmation == false)
|
||||
}
|
||||
|
||||
@Test("MutationPolicy readOnly preset rejects all mutations")
|
||||
func mutationPolicyReadOnlyRejectsAll() {
|
||||
let parser = SQLQueryParser(mutationPolicy: .readOnly)
|
||||
#expect(throws: SQLParsingError.operationNotAllowed(.insert)) {
|
||||
try parser.parse("INSERT INTO t (a) VALUES (1)")
|
||||
}
|
||||
#expect(throws: SQLParsingError.operationNotAllowed(.update)) {
|
||||
try parser.parse("UPDATE t SET a = 1")
|
||||
}
|
||||
#expect(throws: SQLParsingError.operationNotAllowed(.delete)) {
|
||||
try parser.parse("DELETE FROM t WHERE id = 1")
|
||||
}
|
||||
}
|
||||
|
||||
@Test("MutationPolicy table matching is case-insensitive")
|
||||
func mutationPolicyTableCaseInsensitive() throws {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.insert],
|
||||
allowedTables: ["Orders"]
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
let text = "INSERT INTO orders (product) VALUES ('Widget')"
|
||||
let result = try parser.parse(text)
|
||||
#expect(result.operation == .insert)
|
||||
}
|
||||
|
||||
@Test("MutationPolicy handles quoted table names")
|
||||
func mutationPolicyQuotedTableNames() throws {
|
||||
let policy = MutationPolicy(
|
||||
allowedOperations: [.insert, .update],
|
||||
allowedTables: ["order_items"]
|
||||
)
|
||||
let parser = SQLQueryParser(mutationPolicy: policy)
|
||||
|
||||
// Backtick-quoted
|
||||
let backtick = "INSERT INTO `order_items` (qty) VALUES (5)"
|
||||
let r1 = try parser.parse(backtick)
|
||||
#expect(r1.operation == .insert)
|
||||
|
||||
// Double-quote-quoted
|
||||
let doubleQuote = "UPDATE \"order_items\" SET qty = 10 WHERE id = 1"
|
||||
let r2 = try parser.parse(doubleQuote)
|
||||
#expect(r2.operation == .update)
|
||||
}
|
||||
|
||||
@Test("Error description for tableNotAllowed is meaningful")
|
||||
func tableNotAllowedDescription() {
|
||||
let error = SQLParsingError.tableNotAllowed(table: "secret", operation: .delete)
|
||||
#expect(error.description.contains("secret"))
|
||||
#expect(error.description.contains("DELETE"))
|
||||
}
|
||||
|
||||
@Test("Error description for confirmationRequired is meaningful")
|
||||
func confirmationRequiredDescription() {
|
||||
let error = SQLParsingError.confirmationRequired(sql: "DELETE FROM x", operation: .delete)
|
||||
#expect(error.description.contains("DELETE"))
|
||||
#expect(error.description.contains("confirmation"))
|
||||
}
|
||||
|
||||
// MARK: - Robust extraction edge cases
|
||||
|
||||
@Test("Extracts plain SQL without any wrapping")
|
||||
func plainSQL() throws {
|
||||
let text = "SELECT * FROM users"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Extracts SQL from markdown sql code block")
|
||||
func markdownSQLBlock() throws {
|
||||
let text = "```sql\nSELECT * FROM users\n```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Extracts SQL from generic code block")
|
||||
func genericCodeBlock() throws {
|
||||
let text = "```\nSELECT * FROM users\n```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Strips trailing semicolons")
|
||||
func trailingSemicolonEdge() throws {
|
||||
let text = "SELECT * FROM users;"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Extracts SQL with preamble text")
|
||||
func preambleText() throws {
|
||||
let text = "Here's the query:\nSELECT * FROM users"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Handles trailing backticks only (no opening fence)")
|
||||
func trailingBackticksOnly() throws {
|
||||
let text = "SELECT * FROM users\n```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Extracts SQL from single-line code block")
|
||||
func singleLineCodeBlock() throws {
|
||||
let text = "```sql SELECT * FROM users ```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Handles no newline before closing fence")
|
||||
func noNewlineBeforeClosingFence() throws {
|
||||
let text = "```sql\nSELECT * FROM users```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Extracts SQL inline with text prefix")
|
||||
func inlineWithText() throws {
|
||||
let text = "The SQL query is: SELECT * FROM users"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Handles extra whitespace around SQL")
|
||||
func extraWhitespace() throws {
|
||||
let text = "\n\nSELECT * FROM users\n\n"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Extracts SQL from chatty LLM response with preamble and postamble")
|
||||
func chattyLLMResponse() throws {
|
||||
let text = "Sure! Here's the SQL:\n\n```sql\nSELECT * FROM users\n```\n\nThis will return all users."
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Preserves SQL comments")
|
||||
func sqlWithComments() throws {
|
||||
let text = "SELECT * FROM users -- get all users"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("-- get all users"))
|
||||
}
|
||||
|
||||
@Test("Preserves backtick-quoted identifiers in SQL")
|
||||
func backtickQuotedIdentifiers() throws {
|
||||
let text = "SELECT `column name` FROM users"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("`column name`"))
|
||||
}
|
||||
|
||||
@Test("Strips think tags from Qwen-style models")
|
||||
func thinkTags() throws {
|
||||
let text = "<think>I need to query the users table</think>\nSELECT * FROM users"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
#expect(!result.sql.contains("think"))
|
||||
}
|
||||
|
||||
@Test("Handles 4 or 5 backtick fences")
|
||||
func extraBacktickFences() throws {
|
||||
let text4 = "````sql\nSELECT * FROM users\n````"
|
||||
let result4 = try readOnlyParser.parse(text4)
|
||||
#expect(result4.sql == "SELECT * FROM users")
|
||||
|
||||
let text5 = "`````\nSELECT * FROM users\n`````"
|
||||
let result5 = try readOnlyParser.parse(text5)
|
||||
#expect(result5.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Handles mixed case SQL keywords")
|
||||
func mixedCaseSQL() throws {
|
||||
let text = "select * from USERS"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "select * from USERS")
|
||||
}
|
||||
|
||||
@Test("Handles WITH clause (CTE) queries")
|
||||
func withClause() throws {
|
||||
let text = "WITH cte AS (SELECT id FROM orders) SELECT * FROM cte"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.hasPrefix("WITH"))
|
||||
#expect(result.operation == .select)
|
||||
}
|
||||
|
||||
@Test("Handles WITH clause in code block")
|
||||
func withClauseInCodeBlock() throws {
|
||||
let text = "```sql\nWITH top AS (\n SELECT user_id, COUNT(*) as cnt FROM orders GROUP BY user_id\n)\nSELECT * FROM top WHERE cnt > 5\n```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.hasPrefix("WITH"))
|
||||
#expect(result.operation == .select)
|
||||
}
|
||||
|
||||
@Test("Multi-line SQL with JOINs and subqueries in code block")
|
||||
func multiLineJoinsAndSubqueries() throws {
|
||||
let text = """
|
||||
```sql
|
||||
SELECT u.name, o.total
|
||||
FROM users u
|
||||
INNER JOIN orders o ON u.id = o.user_id
|
||||
WHERE o.total > (SELECT AVG(total) FROM orders)
|
||||
ORDER BY o.total DESC
|
||||
```
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("INNER JOIN"))
|
||||
#expect(result.sql.contains("SELECT AVG(total)"))
|
||||
#expect(result.sql.contains("ORDER BY"))
|
||||
}
|
||||
|
||||
@Test("Handles response with both explanation text and SQL")
|
||||
func explanationAndSQL() throws {
|
||||
let text = """
|
||||
To find all active users, we need to query the users table
|
||||
and filter by the active column. Here's the query:
|
||||
|
||||
SELECT * FROM users WHERE active = 1
|
||||
|
||||
This should give you the results you're looking for.
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users WHERE active = 1")
|
||||
}
|
||||
|
||||
@Test("Throws noSQLFound for empty response")
|
||||
func emptyResponse() throws {
|
||||
#expect(throws: SQLParsingError.noSQLFound) {
|
||||
try readOnlyParser.parse("")
|
||||
}
|
||||
#expect(throws: SQLParsingError.noSQLFound) {
|
||||
try readOnlyParser.parse(" \n\n ")
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Throws noSQLFound for response with no SQL at all")
|
||||
func noSQLAtAll() throws {
|
||||
#expect(throws: SQLParsingError.noSQLFound) {
|
||||
try readOnlyParser.parse("I cannot help with that question. Please try asking about your data.")
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Handles response with multiple SQL statements in code block (rejects them)")
|
||||
func multipleStatementsInCodeBlock() throws {
|
||||
// When multiple statements are in a code block, the parser sees both and rejects
|
||||
let text = "```sql\nSELECT * FROM users; SELECT * FROM orders\n```"
|
||||
#expect(throws: SQLParsingError.multipleStatements) {
|
||||
try readOnlyParser.parse(text)
|
||||
}
|
||||
}
|
||||
|
||||
@Test("Extracts first SQL statement from plain text with multiple statements")
|
||||
func multipleStatementsPlainText() throws {
|
||||
// In plain text, the direct extraction stops at the semicolon and extracts the first statement
|
||||
let text = "SELECT * FROM users; SELECT * FROM orders"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Preserves backtick identifiers inside code blocks")
|
||||
func backtickIdentifiersInCodeBlock() throws {
|
||||
let text = "```sql\nSELECT `first name`, `last name` FROM `user data`\n```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("`first name`"))
|
||||
#expect(result.sql.contains("`last name`"))
|
||||
#expect(result.sql.contains("`user data`"))
|
||||
}
|
||||
|
||||
@Test("Strips think tags with multiline reasoning content")
|
||||
func multilineThinkTags() throws {
|
||||
let text = """
|
||||
<think>
|
||||
The user wants to find all users.
|
||||
I should use SELECT * FROM users.
|
||||
Let me think about which columns to include...
|
||||
</think>
|
||||
SELECT * FROM users
|
||||
"""
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql == "SELECT * FROM users")
|
||||
}
|
||||
|
||||
@Test("Handles mixed backtick styles in response")
|
||||
func mixedBacktickStyles() throws {
|
||||
// Code fences + backtick-quoted identifiers inside
|
||||
let text = "```sql\nSELECT `user name` FROM users WHERE `is active` = 1\n```"
|
||||
let result = try readOnlyParser.parse(text)
|
||||
#expect(result.sql.contains("`user name`"))
|
||||
#expect(result.sql.contains("`is active`"))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user