Skip to content

Commit 75c3f12

Browse files
committed
fix(plugin-snowflake): poll long queries with backoff up to 45 minutes and add unit tests for auth parsing and type mapping
1 parent c30a5b1 commit 75c3f12

12 files changed

Lines changed: 239 additions & 63 deletions

File tree

Plugins/SnowflakeDriverPlugin/SnowflakeAuth.swift

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -216,11 +216,13 @@ enum SnowflakeConnectionsTOML {
216216
}
217217

218218
private static func stripComment(_ line: String) -> String {
219-
var inString = false
219+
var inDoubleQuotes = false
220+
var inSingleQuotes = false
220221
var result = ""
221222
for char in line {
222-
if char == "\"" { inString.toggle() }
223-
if char == "#" && !inString { break }
223+
if char == "\"" && !inSingleQuotes { inDoubleQuotes.toggle() }
224+
if char == "'" && !inDoubleQuotes { inSingleQuotes.toggle() }
225+
if char == "#" && !inDoubleQuotes && !inSingleQuotes { break }
224226
result.append(char)
225227
}
226228
return result

Plugins/SnowflakeDriverPlugin/SnowflakeConnection.swift

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -412,18 +412,28 @@ final class SnowflakeConnection: @unchecked Sendable {
412412
]
413413
}
414414

415+
private static let queryPollTimeout: TimeInterval = 2_700
416+
private static let queryPollMaxInterval: UInt64 = 5_000_000_000
417+
415418
private func pollIfInProgress(_ initial: [String: Any], token: String) async throws -> [String: Any] {
416419
var response = initial
417-
var attempts = 0
418-
while Self.isInProgress(response), attempts < 600 {
419-
attempts += 1
420+
let deadline = Date().addingTimeInterval(Self.queryPollTimeout)
421+
var interval: UInt64 = 500_000_000
422+
while Self.isInProgress(response) {
420423
guard let data = response["data"] as? [String: Any],
421424
let resultPath = data["getResultUrl"] as? String else {
422425
break
423426
}
424-
try await Task.sleep(nanoseconds: 500_000_000)
427+
guard Date() < deadline else {
428+
throw SnowflakeError.timeout("Query did not finish within 45 minutes")
429+
}
430+
try await Task.sleep(nanoseconds: interval)
431+
interval = min(interval * 2, Self.queryPollMaxInterval)
425432
response = try await getJSON(path: resultPath, token: token)
426433
}
434+
if Self.isInProgress(response) {
435+
throw SnowflakeError.invalidResponse("Query is still running but Snowflake returned no result URL")
436+
}
427437
return response
428438
}
429439

Plugins/SnowflakeDriverPlugin/SnowflakePluginDriver.swift

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,11 +118,13 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
118118
func switchDatabase(to database: String) async throws {
119119
guard let conn = connection else { throw SnowflakeError.notConnected }
120120
try await conn.switchDatabase(to: database)
121+
lock.withLock { resolvedSchemaCache.removeAll() }
121122
}
122123

123124
func switchSchema(to schema: String) async throws {
124125
guard let conn = connection else { throw SnowflakeError.notConnected }
125126
try await conn.switchSchema(to: schema)
127+
lock.withLock { resolvedSchemaCache.removeAll() }
126128
}
127129

128130
// MARK: - Database Management

Plugins/SnowflakeDriverPlugin/SnowflakeTypeMapper.swift

Lines changed: 1 addition & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,10 @@
22
// SnowflakeTypeMapper.swift
33
// SnowflakeDriverPlugin
44
//
5-
// Maps Snowflake's internal row metadata types to display type names and
6-
// decodes JSON cell values into PluginCellValue.
5+
// Maps Snowflake's internal row metadata types to display type names.
76
//
87

98
import Foundation
10-
import TableProPluginKit
119

1210
struct SnowflakeColumnMeta: Sendable {
1311
let name: String
@@ -62,29 +60,4 @@ enum SnowflakeTypeMapper {
6260
return column.internalType.uppercased()
6361
}
6462
}
65-
66-
/// Snowflake's v1 query protocol returns every cell as a JSON string or null.
67-
static func cellValue(from json: Any?) -> PluginCellValue {
68-
switch json {
69-
case nil, is NSNull:
70-
return .null
71-
case let string as String:
72-
return .text(string)
73-
case let number as NSNumber:
74-
return .text(number.stringValue)
75-
case let bool as Bool:
76-
return .text(bool ? "true" : "false")
77-
default:
78-
return .text(String(describing: json ?? ""))
79-
}
80-
}
81-
82-
static func columnInfo(from column: SnowflakeColumnMeta, primaryKeys: Set<String>) -> PluginColumnInfo {
83-
PluginColumnInfo(
84-
name: column.name,
85-
dataType: displayType(for: column),
86-
isNullable: column.nullable,
87-
isPrimaryKey: primaryKeys.contains(column.name.uppercased())
88-
)
89-
}
9063
}

TablePro.xcodeproj/project.pbxproj

Lines changed: 13 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -592,8 +592,6 @@
592592
};
593593
5A32BC012F9D5F1300BAEB5F /* mcp-server */ = {
594594
isa = PBXFileSystemSynchronizedRootGroup;
595-
exceptions = (
596-
);
597595
path = "mcp-server";
598596
sourceTree = "<group>";
599597
};
@@ -769,8 +767,6 @@
769767
};
770768
5ABCC5A82F43856700EAF3FC /* TableProTests */ = {
771769
isa = PBXFileSystemSynchronizedRootGroup;
772-
exceptions = (
773-
);
774770
path = TableProTests;
775771
sourceTree = "<group>";
776772
};
@@ -784,8 +780,6 @@
784780
};
785781
5AF00A122FB9000000000001 /* TableProUITests */ = {
786782
isa = PBXFileSystemSynchronizedRootGroup;
787-
exceptions = (
788-
);
789783
path = TableProUITests;
790784
sourceTree = "<group>";
791785
};
@@ -1057,9 +1051,9 @@
10571051
5A1091BE2EF17EDC0055EA7C = {
10581052
isa = PBXGroup;
10591053
children = (
1060-
5ADDB00500000000000000B0 /* DynamoDBDriverPlugin */,
1061-
5ABQR00500000000000000B0 /* BigQueryDriverPlugin */,
1062-
5AEA8B412F6808CA0040461A /* EtcdDriverPlugin */,
1054+
5ADDB00500000000000000B0 /* Plugins/DynamoDBDriverPlugin */,
1055+
5ABQR00500000000000000B0 /* Plugins/BigQueryDriverPlugin */,
1056+
5AEA8B412F6808CA0040461A /* Plugins/EtcdDriverPlugin */,
10631057
5AE4F4812F6BC0640097AC5B /* Plugins/CloudflareD1DriverPlugin */,
10641058
5A3BE6FE2F97DB0100611C1F /* Plugins/LibSQLDriverPlugin */,
10651059
5A1091C92EF17EDC0055EA7C /* TablePro */,
@@ -1130,7 +1124,7 @@
11301124
name = Products;
11311125
sourceTree = "<group>";
11321126
};
1133-
5ABQR00500000000000000B0 /* BigQueryDriverPlugin */ = {
1127+
5ABQR00500000000000000B0 /* Plugins/BigQueryDriverPlugin */ = {
11341128
isa = PBXGroup;
11351129
children = (
11361130
5ABQR00200000000000000A1 /* BigQueryAuth.swift */,
@@ -1145,7 +1139,7 @@
11451139
path = Plugins/BigQueryDriverPlugin;
11461140
sourceTree = "<group>";
11471141
};
1148-
5ADDB00500000000000000B0 /* DynamoDBDriverPlugin */ = {
1142+
5ADDB00500000000000000B0 /* Plugins/DynamoDBDriverPlugin */ = {
11491143
isa = PBXGroup;
11501144
children = (
11511145
5ADDB00200000000000000A1 /* DynamoDBConnection.swift */,
@@ -1159,7 +1153,7 @@
11591153
path = Plugins/DynamoDBDriverPlugin;
11601154
sourceTree = "<group>";
11611155
};
1162-
5AEA8B412F6808CA0040461A /* EtcdDriverPlugin */ = {
1156+
5AEA8B412F6808CA0040461A /* Plugins/EtcdDriverPlugin */ = {
11631157
isa = PBXGroup;
11641158
children = (
11651159
5AEA8B3B2F6808CA0040461A /* EtcdCommandParser.swift */,
@@ -1175,18 +1169,10 @@
11751169
5AEA8B482F6808E90040461A /* Frameworks */ = {
11761170
isa = PBXGroup;
11771171
children = (
1178-
665759805EBF5C4BF064ADAB /* OS X */,
11791172
);
11801173
name = Frameworks;
11811174
sourceTree = "<group>";
11821175
};
1183-
665759805EBF5C4BF064ADAB /* OS X */ = {
1184-
isa = PBXGroup;
1185-
children = (
1186-
);
1187-
name = "OS X";
1188-
sourceTree = "<group>";
1189-
};
11901176
8FE5E1F9D0550A0E0AACD3EB /* SnowflakeDriverPlugin */ = {
11911177
isa = PBXGroup;
11921178
children = (
@@ -1940,11 +1926,11 @@
19401926
mainGroup = 5A1091BE2EF17EDC0055EA7C;
19411927
minimizedProjectReferenceProxies = 1;
19421928
packageReferences = (
1943-
5A0000012F4F000000000101 /* XCLocalSwiftPackageReference "CodeEditSourceEditor" */,
1929+
5A0000012F4F000000000101 /* XCLocalSwiftPackageReference "LocalPackages/CodeEditSourceEditor" */,
19441930
5ACE00012F4F000000000008 /* XCRemoteSwiftPackageReference "Sparkle" */,
1945-
5A0000012F4F000000000100 /* XCLocalSwiftPackageReference "CodeEditLanguages" */,
1931+
5A0000012F4F000000000100 /* XCLocalSwiftPackageReference "LocalPackages/CodeEditLanguages" */,
19461932
5ACE00012F4F00000000000E /* XCRemoteSwiftPackageReference "oracle-nio" */,
1947-
5A0000012F4F000000000102 /* XCLocalSwiftPackageReference "TableProCore" */,
1933+
5A0000012F4F000000000102 /* XCLocalSwiftPackageReference "Packages/TableProCore" */,
19481934
5A32BBF92F9D5EAB00BAEB5F /* XCRemoteSwiftPackageReference "swift-certificates" */,
19491935
);
19501936
preferredProjectObjectVersion = 77;
@@ -4704,15 +4690,15 @@
47044690
/* End XCConfigurationList section */
47054691

47064692
/* Begin XCLocalSwiftPackageReference section */
4707-
5A0000012F4F000000000100 /* XCLocalSwiftPackageReference "CodeEditLanguages" */ = {
4693+
5A0000012F4F000000000100 /* XCLocalSwiftPackageReference "LocalPackages/CodeEditLanguages" */ = {
47084694
isa = XCLocalSwiftPackageReference;
47094695
relativePath = LocalPackages/CodeEditLanguages;
47104696
};
4711-
5A0000012F4F000000000101 /* XCLocalSwiftPackageReference "CodeEditSourceEditor" */ = {
4697+
5A0000012F4F000000000101 /* XCLocalSwiftPackageReference "LocalPackages/CodeEditSourceEditor" */ = {
47124698
isa = XCLocalSwiftPackageReference;
47134699
relativePath = LocalPackages/CodeEditSourceEditor;
47144700
};
4715-
5A0000012F4F000000000102 /* XCLocalSwiftPackageReference "TableProCore" */ = {
4701+
5A0000012F4F000000000102 /* XCLocalSwiftPackageReference "Packages/TableProCore" */ = {
47164702
isa = XCLocalSwiftPackageReference;
47174703
relativePath = Packages/TableProCore;
47184704
};
@@ -4779,7 +4765,7 @@
47794765
};
47804766
5AD1D8C12FB5000000000002 /* TableProMSSQLCore */ = {
47814767
isa = XCSwiftPackageProductDependency;
4782-
package = 5A0000012F4F000000000102 /* XCLocalSwiftPackageReference "TableProCore" */;
4768+
package = 5A0000012F4F000000000102 /* XCLocalSwiftPackageReference "Packages/TableProCore" */;
47834769
productName = TableProMSSQLCore;
47844770
};
47854771
/* End XCSwiftPackageProductDependency section */
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
../../Plugins/SnowflakeDriverPlugin/SnowflakeAuth.swift
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
../../Plugins/SnowflakeDriverPlugin/SnowflakeError.swift
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
../../Plugins/SnowflakeDriverPlugin/SnowflakeTypeMapper.swift
Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
//
2+
// SnowflakeAuthTests.swift
3+
// TableProTests
4+
//
5+
// Tests for SnowflakeAccount, SnowflakeConnectionsTOML, and the SPKI
6+
// wrapping used for key-pair JWT fingerprints (compiled via symlink from
7+
// SnowflakeDriverPlugin).
8+
//
9+
10+
import Foundation
11+
import Testing
12+
13+
@Suite("Snowflake Account Parsing")
14+
struct SnowflakeAccountTests {
15+
@Test("Plain locator gets the Snowflake domain appended")
16+
func testHostFromLocator() {
17+
#expect(SnowflakeAccount.host(forAccount: "xy12345.us-east-1") == "xy12345.us-east-1.snowflakecomputing.com")
18+
#expect(SnowflakeAccount.host(forAccount: "myorg-myaccount") == "myorg-myaccount.snowflakecomputing.com")
19+
}
20+
21+
@Test("Full hostnames pass through unchanged, case-insensitively")
22+
func testHostPassthrough() {
23+
#expect(
24+
SnowflakeAccount.host(forAccount: "abc.snowflakecomputing.com") == "abc.snowflakecomputing.com"
25+
)
26+
#expect(
27+
SnowflakeAccount.host(forAccount: "Abc.SnowflakeComputing.Com") == "Abc.SnowflakeComputing.Com"
28+
)
29+
}
30+
31+
@Test("URL forms resolve to their host")
32+
func testHostFromURL() {
33+
#expect(
34+
SnowflakeAccount.host(forAccount: "https://abc.snowflakecomputing.com/console") ==
35+
"abc.snowflakecomputing.com"
36+
)
37+
}
38+
39+
@Test("Whitespace is trimmed before resolution")
40+
func testHostTrimsWhitespace() {
41+
#expect(SnowflakeAccount.host(forAccount: " abc \n") == "abc.snowflakecomputing.com")
42+
}
43+
44+
@Test("Issuer account name drops domain and region, then uppercases")
45+
func testIssuerAccountName() {
46+
#expect(SnowflakeAccount.issuerAccountName(forAccount: "xy12345.us-east-1") == "XY12345")
47+
#expect(
48+
SnowflakeAccount.issuerAccountName(forAccount: "xy12345.us-east-1.snowflakecomputing.com") == "XY12345"
49+
)
50+
#expect(SnowflakeAccount.issuerAccountName(forAccount: "myorg-myaccount") == "MYORG-MYACCOUNT")
51+
}
52+
}
53+
54+
@Suite("Snowflake Connections TOML")
55+
struct SnowflakeConnectionsTOMLTests {
56+
@Test("Parses sections with key-value pairs")
57+
func testBasicSection() {
58+
let toml = """
59+
[default]
60+
account = "xy12345"
61+
user = jane
62+
"""
63+
let parsed = SnowflakeConnectionsTOML.parse(toml)
64+
#expect(parsed["default"]?["account"] == "xy12345")
65+
#expect(parsed["default"]?["user"] == "jane")
66+
}
67+
68+
@Test("Strips the connections. prefix from config.toml sections")
69+
func testConnectionsPrefix() {
70+
let parsed = SnowflakeConnectionsTOML.parse("[connections.dev]\naccount = 'abc'\n")
71+
#expect(parsed["dev"]?["account"] == "abc")
72+
}
73+
74+
@Test("Quoted section names are unquoted")
75+
func testQuotedSectionName() {
76+
let parsed = SnowflakeConnectionsTOML.parse("[connections.\"my conn\"]\nrole = \"ADMIN\"\n")
77+
#expect(parsed["my conn"]?["role"] == "ADMIN")
78+
}
79+
80+
@Test("Comments are stripped outside strings and kept inside both quote styles")
81+
func testCommentHandling() {
82+
let toml = """
83+
[default]
84+
account = "abc" # trailing comment
85+
password = "p#ss"
86+
token = 'a#b'
87+
"""
88+
let parsed = SnowflakeConnectionsTOML.parse(toml)
89+
#expect(parsed["default"]?["account"] == "abc")
90+
#expect(parsed["default"]?["password"] == "p#ss")
91+
#expect(parsed["default"]?["token"] == "a#b")
92+
}
93+
94+
@Test("Key-value pairs before any section are ignored")
95+
func testKeysOutsideSectionIgnored() {
96+
let parsed = SnowflakeConnectionsTOML.parse("account = \"abc\"\n[dev]\nuser = \"u\"\n")
97+
#expect(parsed.count == 1)
98+
#expect(parsed["dev"]?["user"] == "u")
99+
}
100+
}
101+
102+
@Suite("Snowflake SPKI Wrapping")
103+
struct SnowflakeSPKIWrappingTests {
104+
private static let rsaAlgorithmID: [UInt8] = [
105+
0x30, 0x0D, 0x06, 0x09, 0x2A, 0x86, 0x48, 0x86,
106+
0xF7, 0x0D, 0x01, 0x01, 0x01, 0x05, 0x00
107+
]
108+
109+
@Test("Short keys use single-byte DER lengths")
110+
func testShortFormLength() {
111+
let pkcs1 = Data((0..<10).map { UInt8($0) })
112+
let spki = [UInt8](SnowflakeKeyPairAuth.wrapPKCS1IntoSPKI(pkcs1))
113+
114+
#expect(spki[0] == 0x30)
115+
#expect(Int(spki[1]) == spki.count - 2)
116+
#expect(Array(spki[2..<17]) == Self.rsaAlgorithmID)
117+
#expect(spki[17] == 0x03)
118+
#expect(Int(spki[18]) == pkcs1.count + 1)
119+
#expect(spki[19] == 0x00)
120+
#expect(Array(spki.suffix(pkcs1.count)) == [UInt8](pkcs1))
121+
}
122+
123+
@Test("Keys past 127 bytes use long-form DER lengths")
124+
func testLongFormLength() {
125+
let pkcs1 = Data(repeating: 0xAB, count: 270)
126+
let spki = [UInt8](SnowflakeKeyPairAuth.wrapPKCS1IntoSPKI(pkcs1))
127+
128+
#expect(spki[0] == 0x30)
129+
#expect(spki[1] == 0x82)
130+
let bodyLength = (Int(spki[2]) << 8) | Int(spki[3])
131+
#expect(bodyLength == spki.count - 4)
132+
#expect(Array(spki[4..<19]) == Self.rsaAlgorithmID)
133+
#expect(spki[19] == 0x03)
134+
#expect(spki[20] == 0x82)
135+
let bitStringLength = (Int(spki[21]) << 8) | Int(spki[22])
136+
#expect(bitStringLength == pkcs1.count + 1)
137+
#expect(spki[23] == 0x00)
138+
#expect(Array(spki.suffix(pkcs1.count)) == [UInt8](pkcs1))
139+
}
140+
}

0 commit comments

Comments
 (0)