fix(plugin-snowflake): key the shared session on the saved connection and scope Stop to its own driver (#2470)

Claude-Session: https://claude.ai/code/session_01Qk1xfY3vnneRifC22eV2r7

Author: datlechin

CHANGELOG.md

@@ -61,6 +61,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Snowflake `BINARY` cells showing the hex of their hex, which the hex editor then wrote back.
 - A date cell the app cannot read opening a picker set to today, which overwrote the value on OK. (#2454)
 - Snowflake reporting no rows changed for an `UPDATE` or `MERGE`, and a `SELECT`'s own result for a column named `number of rows`.
+- Two saved Snowflake connections to one account sharing a session, so switching database in one window moved the other.
+- Stop on a Snowflake query cancelling every other query on the same connection, including a sidebar refresh or a save.
 
 ### Security
 

Plugins/SnowflakeDriverPlugin/SnowflakeConnection.swift

@@ -40,20 +40,31 @@ final class SnowflakeConnection: @unchecked Sendable {
 
     let host: String
     let params: ResolvedParameters
+    private let connectionIdentifier: String
 
     private let session: URLSession
     private let lock = NSLock()
     private let heartbeat = SnowflakeHeartbeat()
     private var sessionToken: String?
     private var renewalToken: String?
-    private var activeRequestIDs: Set<String> = []
+    private var activeRequestIDs: [String: Set<String>] = [:]
     private var sequenceId = 0
     private var connectTask: Task<Void, Error>?
 
     var sessionFingerprint: String {
-        [host, params.user.uppercased(), params.authMethod, params.role.uppercased()].joined(separator: "|")
+        SnowflakeSessionKey.fingerprint(
+            connectionId: connectionIdentifier,
+            host: host,
+            user: params.user,
+            authMethod: params.authMethod,
+            role: params.role
+        )
     }
 
+    /// Statements the connection issues for itself, such as the connect-time `USE` calls, belong to
+    /// no driver and are never the target of a Stop.
+    static let sessionOwner = "session"
+
     private var _currentDatabase: String?
     private var _currentSchema: String?
     private var _currentWarehouse: String?
@@ -71,6 +82,7 @@ final class SnowflakeConnection: @unchecked Sendable {
     init(config: DriverConnectionConfig) {
         self.params = Self.resolveParameters(from: config)
         self.host = SnowflakeAccount.host(forAccount: params.account)
+        self.connectionIdentifier = config.additionalFields["connectionId"] ?? ""
 
         let configuration = URLSessionConfiguration.ephemeral
         configuration.timeoutIntervalForRequest = 120
@@ -419,9 +431,13 @@ final class SnowflakeConnection: @unchecked Sendable {
 
     // MARK: - Query Execution
 
-    func query(_ sql: String, parameters: [PluginCellValue] = []) async throws -> SnowflakeQueryResult {
+    func query(
+        _ sql: String,
+        parameters: [PluginCellValue] = [],
+        owner: String = SnowflakeConnection.sessionOwner
+    ) async throws -> SnowflakeQueryResult {
         try await withReauthentication {
-            try await performQuery(sql, parameters: parameters)
+            try await performQuery(sql, parameters: parameters, owner: owner)
         }
     }
 
@@ -439,8 +455,12 @@ final class SnowflakeConnection: @unchecked Sendable {
         }
     }
 
-    func cancelAllQueries() {
-        let (requestIDs, token) = lock.withLock { (activeRequestIDs, sessionToken) }
+    /// One Snowflake session is shared by the driver the user sees and by every pooled metadata
+    /// driver behind it, so an abort has to name whose work it is stopping. Aborting the whole
+    /// session made Stop on a query cancel a sidebar refresh running beside it, and cancel a save's
+    /// remaining statements.
+    func cancelQueries(owner: String) {
+        let (requestIDs, token) = lock.withLock { (activeRequestIDs[owner] ?? [], sessionToken) }
         guard !requestIDs.isEmpty, let token else { return }
         Task { [weak self] in
             for requestID in requestIDs {
@@ -454,8 +474,12 @@ final class SnowflakeConnection: @unchecked Sendable {
         }
     }
 
-    private func performQuery(_ sql: String, parameters: [PluginCellValue] = []) async throws -> SnowflakeQueryResult {
-        let (data, token) = try await submitQuery(sql, parameters: parameters)
+    private func performQuery(
+        _ sql: String,
+        parameters: [PluginCellValue] = [],
+        owner: String
+    ) async throws -> SnowflakeQueryResult {
+        let (data, token) = try await submitQuery(sql, parameters: parameters, owner: owner)
         if let resultIds = data["resultIds"] as? String, !resultIds.isEmpty {
             return try await collectMultiStatementResults(ids: resultIds, token: token)
         }
@@ -465,7 +489,8 @@ final class SnowflakeConnection: @unchecked Sendable {
 
     private func submitQuery(
         _ sql: String,
-        parameters: [PluginCellValue]
+        parameters: [PluginCellValue],
+        owner: String
     ) async throws -> (data: [String: Any], token: String) {
         guard let token = lock.withLock({ sessionToken }) else {
             throw SnowflakeError.notConnected
@@ -474,11 +499,14 @@ final class SnowflakeConnection: @unchecked Sendable {
         let requestID = UUID().uuidString.lowercased()
         let sequence = lock.withLock { () -> Int in
             sequenceId += 1
-            activeRequestIDs.insert(requestID)
+            activeRequestIDs[owner, default: []].insert(requestID)
             return sequenceId
         }
         defer {
-            lock.withLock { _ = activeRequestIDs.remove(requestID) }
+            lock.withLock {
+                activeRequestIDs[owner]?.remove(requestID)
+                if activeRequestIDs[owner]?.isEmpty == true { activeRequestIDs[owner] = nil }
+            }
         }
 
         var body: [String: Any] = [
@@ -611,9 +639,12 @@ final class SnowflakeConnection: @unchecked Sendable {
         let batches: AsyncThrowingStream<[[PluginCellValueBox]], Error>
     }
 
-    func queryStreamed(_ sql: String) async throws -> StreamedResult {
+    func queryStreamed(
+        _ sql: String,
+        owner: String = SnowflakeConnection.sessionOwner
+    ) async throws -> StreamedResult {
         let (data, _) = try await withReauthentication {
-            try await submitQuery(sql, parameters: [])
+            try await submitQuery(sql, parameters: [], owner: owner)
         }
         applyFinalSessionInfo(data)
 

Plugins/SnowflakeDriverPlugin/SnowflakePluginDriver.swift

@@ -17,6 +17,10 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
     private var resolvedSchemaCache: [String: String] = [:]
     private var columnTypeCache: [String: [String: String]] = [:]
 
+    /// Several drivers share one Snowflake session, so a Stop has to name its own work. This
+    /// identifies the statements this driver issued and nobody else's.
+    private let queryOwner = UUID().uuidString
+
     private static let logger = Logger(subsystem: "com.TablePro", category: "SnowflakePluginDriver")
 
     private var connection: SnowflakeConnection? {
@@ -41,7 +45,7 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
     }
 
     func cancelQuery() throws {
-        connection?.cancelAllQueries()
+        connection?.cancelQueries(owner: queryOwner)
     }
 
     var supportsSchemas: Bool { true }
@@ -55,7 +59,7 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
         guard !parameters.isEmpty else { return try await execute(query: query) }
         guard let conn = connection else { throw SnowflakeError.notConnected }
         let startTime = Date()
-        let result = try await conn.query(query, parameters: parameters)
+        let result = try await conn.query(query, parameters: parameters, owner: queryOwner)
         return PluginQueryResult(
             columns: result.columns.map(\.name),
             columnTypeNames: result.columns.map(SnowflakeTypeMapper.displayType),
@@ -79,7 +83,7 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
         }
         lock.withLock { _connection = conn }
 
-        if let result = try? await conn.query("SELECT CURRENT_VERSION()"),
+        if let result = try? await conn.query("SELECT CURRENT_VERSION()", owner: queryOwner),
            let first = result.rows.first?.first, case .text(let version) = first {
             lock.withLock { _serverVersion = "Snowflake \(version)" }
         } else {
@@ -114,7 +118,7 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
     func execute(query: String) async throws -> PluginQueryResult {
         guard let conn = connection else { throw SnowflakeError.notConnected }
         let startTime = Date()
-        let result = try await conn.query(query)
+        let result = try await conn.query(query, owner: queryOwner)
         let executionTime = Date().timeIntervalSince(startTime)
 
         if result.columns.isEmpty {
@@ -201,9 +205,9 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
         guard let conn = connection else { throw SnowflakeError.notConnected }
         switch id {
         case "warehouse":
-            _ = try await conn.query("USE WAREHOUSE \(quoteIdentifier(value))")
+            _ = try await conn.query("USE WAREHOUSE \(quoteIdentifier(value))", owner: queryOwner)
         case "role":
-            _ = try await conn.query("USE ROLE \(quoteIdentifier(value))")
+            _ = try await conn.query("USE ROLE \(quoteIdentifier(value))", owner: queryOwner)
             lock.withLock {
                 resolvedSchemaCache.removeAll()
                 columnTypeCache.removeAll()
@@ -583,7 +587,7 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
                 do {
                     guard let conn = self.connection else { throw SnowflakeError.notConnected }
                     let trimmed = query.replacingOccurrences(of: ";\\s*\\z", with: "", options: .regularExpression)
-                    let streamed = try await conn.queryStreamed(trimmed)
+                    let streamed = try await conn.queryStreamed(trimmed, owner: queryOwner)
                     continuation.yield(.header(PluginStreamHeader(
                         columns: streamed.columns.map(\.name),
                         columnTypeNames: streamed.columns.map(SnowflakeTypeMapper.displayType),
@@ -682,7 +686,7 @@ final class SnowflakePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
 
     private func rawQuery(_ sql: String) async throws -> SnowflakeQueryResult {
         guard let conn = connection else { throw SnowflakeError.notConnected }
-        return try await conn.query(sql)
+        return try await conn.query(sql, owner: queryOwner)
     }
 
     private func namedValues(in result: SnowflakeQueryResult, column: String) -> [String] {

Plugins/SnowflakeDriverPlugin/SnowflakeSessionKey.swift

@@ -0,0 +1,28 @@
+//
+//  SnowflakeSessionKey.swift
+//  SnowflakeDriverPlugin
+//
+//  Decides which drivers share one authenticated Snowflake session.
+//
+
+import Foundation
+
+enum SnowflakeSessionKey {
+    /// Two saved connections can reach the same account, user and role and still be different
+    /// connections. Keying on the account alone let them share a session, so `USE DATABASE` in one
+    /// window moved the other window's current database.
+    ///
+    /// The saved connection's own identifier separates them. The database cannot: the metadata pool
+    /// rewrites it on its copy of the connection before building a driver, so including it here
+    /// would give that driver a different key, a second login, and another MFA prompt, which is the
+    /// whole reason the session is shared in the first place.
+    static func fingerprint(
+        connectionId: String,
+        host: String,
+        user: String,
+        authMethod: String,
+        role: String
+    ) -> String {
+        [connectionId, host, user.uppercased(), authMethod, role.uppercased()].joined(separator: "|")
+    }
+}

TablePro/Core/Database/DatabaseDriver.swift

@@ -581,6 +581,7 @@ enum DatabaseDriverFactory {
             additionalFields["enableCleartextPlugin"] = "true"
         }
         additionalFields["queryTimeoutSeconds"] = String(AppSettingsManager.shared.general.queryTimeoutSeconds)
+        additionalFields["connectionId"] = connection.id.uuidString
         let config = DriverConnectionConfig(
             host: connection.host,
             port: connection.port,

TableProTests/Plugins/SnowflakeSessionKeyTests.swift

@@ -0,0 +1,62 @@
+//
+//  SnowflakeSessionKeyTests.swift
+//  TableProTests
+//
+//  Tests for SnowflakeSessionKey (compiled via symlink from SnowflakeDriverPlugin).
+//
+
+import Foundation
+import Testing
+
+@Suite("Snowflake Session Key")
+struct SnowflakeSessionKeyTests {
+    private func key(
+        connectionId: String = "A",
+        host: String = "acct.snowflakecomputing.com",
+        user: String = "dana",
+        authMethod: String = "password",
+        role: String = "analyst"
+    ) -> String {
+        SnowflakeSessionKey.fingerprint(
+            connectionId: connectionId,
+            host: host,
+            user: user,
+            authMethod: authMethod,
+            role: role
+        )
+    }
+
+    /// The defect: two saved connections to one account, user and role shared a session, so
+    /// `USE DATABASE` in one moved the other's current database and a save wrote to the wrong one.
+    @Test("Two saved connections to the same account do not share a session")
+    func testDistinctConnectionsDoNotShare() {
+        #expect(key(connectionId: "A") != key(connectionId: "B"))
+    }
+
+    @Test("The same saved connection always resolves to one session")
+    func testSameConnectionShares() {
+        #expect(key(connectionId: "A") == key(connectionId: "A"))
+    }
+
+    /// The metadata pool rewrites the database on its copy before building a driver, so a key that
+    /// varied with the database would give that driver its own login and another MFA prompt.
+    @Test("The database is not part of the key")
+    func testDatabaseIsNotInTheKey() {
+        #expect(!key().contains("ANALYTICS"))
+        #expect(key(connectionId: "A") == key(connectionId: "A"))
+    }
+
+    @Test("Account identity still separates sessions")
+    func testAccountIdentitySeparates() {
+        #expect(key(host: "one.snowflakecomputing.com") != key(host: "two.snowflakecomputing.com"))
+        #expect(key(user: "dana") != key(user: "sam"))
+        #expect(key(authMethod: "password") != key(authMethod: "keyPair"))
+        #expect(key(role: "analyst") != key(role: "admin"))
+    }
+
+    @Test("User and role compare case-insensitively")
+    func testCaseFolding() {
+        #expect(key(user: "dana") == key(user: "DANA"))
+        #expect(key(role: "analyst") == key(role: "ANALYST"))
+    }
+}

project.yml

@@ -486,6 +486,7 @@ targets:
       - Plugins/SnowflakeDriverPlugin/SnowflakeMFATokenStore.swift
       - Plugins/SnowflakeDriverPlugin/SnowflakeObjectQueries.swift
       - Plugins/SnowflakeDriverPlugin/SnowflakeSchemaQueries.swift
+      - Plugins/SnowflakeDriverPlugin/SnowflakeSessionKey.swift
       - Plugins/SnowflakeDriverPlugin/SnowflakeSQL.swift
       - Plugins/SnowflakeDriverPlugin/SnowflakeStatementGenerator.swift
       - Plugins/SnowflakeDriverPlugin/SnowflakeStatementType.swift