@@ -41,6 +41,11 @@ final class DatabaseManager {
4141 /// Health monitors for active connections (MySQL/PostgreSQL only)
4242 private var healthMonitors : [ UUID : ConnectionHealthMonitor ] = [ : ]
4343
44+ /// Tracks connections with user queries currently in-flight.
45+ /// The health monitor skips pings while a query is running to avoid
46+ /// racing on non-thread-safe driver connections.
47+ private var queriesInFlight : [ UUID : Int ] = [ : ]
48+
4449 /// Current session (computed from currentSessionId)
4550 var currentSession : ConnectionSession ? {
4651 guard let sessionId = currentSessionId else { return nil }
@@ -311,10 +316,18 @@ final class DatabaseManager {
311316
312317 /// Execute a query on the current session
313318 func execute( query: String ) async throws -> QueryResult {
314- guard let driver = activeDriver else {
319+ guard let sessionId = currentSessionId , let driver = activeDriver else {
315320 throw DatabaseError . notConnected
316321 }
317322
323+ queriesInFlight [ sessionId, default: 0 ] += 1
324+ defer {
325+ if let count = queriesInFlight [ sessionId] , count > 1 {
326+ queriesInFlight [ sessionId] = count - 1
327+ } else {
328+ queriesInFlight. removeValue ( forKey: sessionId)
329+ }
330+ }
318331 return try await driver. execute ( query: query)
319332 }
320333
@@ -346,17 +359,22 @@ final class DatabaseManager {
346359 sshPasswordOverride: sshPassword
347360 )
348361
349- defer {
350- // Close tunnel after test
362+ let result : Bool
363+ do {
364+ let driver = try DatabaseDriverFactory . createDriver ( for: testConnection)
365+ result = try await driver. testConnection ( )
366+ } catch {
351367 if connection. sshConfig. enabled {
352- Task {
353- try ? await SSHTunnelManager . shared. closeTunnel ( connectionId: connection. id)
354- }
368+ try ? await SSHTunnelManager . shared. closeTunnel ( connectionId: connection. id)
355369 }
370+ throw error
356371 }
357372
358- let driver = try DatabaseDriverFactory . createDriver ( for: testConnection)
359- return try await driver. testConnection ( )
373+ if connection. sshConfig. enabled {
374+ try ? await SSHTunnelManager . shared. closeTunnel ( connectionId: connection. id)
375+ }
376+
377+ return result
360378 }
361379
362380 // MARK: - SSH Tunnel Helper
@@ -447,6 +465,9 @@ final class DatabaseManager {
447465 connectionId: connectionId,
448466 pingHandler: { [ weak self] in
449467 guard let self else { return false }
468+ // Skip ping while a user query is in-flight to avoid racing
469+ // on the same non-thread-safe driver connection.
470+ guard await self . queriesInFlight [ connectionId] == nil else { return true }
450471 guard let mainDriver = await self . activeSessions [ connectionId] ? . driver else {
451472 return false
452473 }
@@ -638,8 +659,14 @@ final class DatabaseManager {
638659
639660 Self . logger. warning ( " SSH tunnel died for connection: \( session. connection. name) " )
640661
641- // Mark connection as reconnecting
662+ // Stop health monitor before retrying to prevent stale pings during reconnect
663+ await stopHealthMonitor ( for: connectionId)
664+
665+ // Disconnect the stale driver and invalidate it so connectToSession
666+ // creates a fresh connection instead of short-circuiting on driver != nil
667+ session. driver? . disconnect ( )
642668 updateSession ( connectionId) { session in
669+ session. driver = nil
643670 session. status = . connecting
644671 }
645672
0 commit comments