fix(connections): fail queued driver work on disconnect and gate Redis database selection

CHANGELOG.md

@@ -29,6 +29,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Dropped tables left staged after a save that failed part way, and drops staged in another database unstaged on refresh.
 - Table tab showing another database's rows after a database switch on PostgreSQL, Redshift and CockroachDB.
 - Wrong approximate row count for a PostgreSQL or PGlite table outside the current schema.
+- Queued queries running on a reopened connection, or stalling it, after a disconnect.
+- Redis keys read from the wrong database when switching databases while a load was running.
+- "This tab is on" error on a table tab for the previous database when it reloads during a database switch.
+- Table on another database opening empty with no error while the connection switches database.
 - Reopen Closed Tab removing the closed tab without reopening it when its connection's window already had tabs.
 - Unchecked and soloed filter rows dropped from a table's saved filters after switching tabs.
 - Table opened in another database from a link, MCP or AppleScript bound to the current schema.

TablePro/Core/Concurrency/SessionDriverGate.swift

@@ -13,22 +13,26 @@ import Foundation
 ///
 /// The body runs inline in the caller's own task rather than in a detached one, so
 /// cancellation still reaches the work.
+///
+/// A turn is owned by the ticket that took it, not by the connection. A drain ends the turn of a
+/// holder that may still be stuck in a driver call, and when that holder finally returns, its
+/// release must not free or hand off a turn a later session has taken since.
 @MainActor
 final class SessionDriverGate {
     private struct Waiter {
         let ticket: UUID
         let continuation: CheckedContinuation<Void, Error>
     }
 
-    private var holders: Set<UUID> = []
+    private var owners: [UUID: UUID] = [:]
     private var waiters: [UUID: [Waiter]] = [:]
 
     func withExclusiveAccess<T>(
         _ connectionId: UUID,
         _ body: () async throws -> T
     ) async throws -> T {
-        try await acquire(connectionId)
-        defer { release(connectionId) }
+        let ticket = try await acquire(connectionId)
+        defer { release(connectionId, ticket: ticket) }
         return try await body()
     }
 
@@ -42,19 +46,19 @@ final class SessionDriverGate {
 
     /// Releases a connection that is going away, failing everyone still queued for it.
     func drain(connectionId: UUID) {
-        holders.remove(connectionId)
+        owners.removeValue(forKey: connectionId)
         let pending = waiters.removeValue(forKey: connectionId) ?? []
         for waiter in pending {
             waiter.continuation.resume(throwing: CancellationError())
         }
     }
 
-    private func acquire(_ connectionId: UUID) async throws {
-        guard holders.contains(connectionId) else {
-            holders.insert(connectionId)
-            return
-        }
+    private func acquire(_ connectionId: UUID) async throws -> UUID {
         let ticket = UUID()
+        guard owners[connectionId] != nil else {
+            owners[connectionId] = ticket
+            return ticket
+        }
         try await withTaskCancellationHandler(
             operation: { try await enqueue(ticket: ticket, connectionId: connectionId) },
             onCancel: { [weak self] in
@@ -63,6 +67,12 @@ final class SessionDriverGate {
                 }
             }
         )
+        /// A hand-off resumes this caller before it runs, so a drain can land in between, and the
+        /// turn it was handed ended with that drain.
+        guard owners[connectionId] == ticket else {
+            throw CancellationError()
+        }
+        return ticket
     }
 
     private func enqueue(ticket: UUID, connectionId: UUID) async throws {
@@ -90,14 +100,16 @@ final class SessionDriverGate {
         waiter.continuation.resume(throwing: CancellationError())
     }
 
-    private func release(_ connectionId: UUID) {
+    private func release(_ connectionId: UUID, ticket: UUID) {
+        guard owners[connectionId] == ticket else { return }
         guard var pending = waiters[connectionId], !pending.isEmpty else {
-            holders.remove(connectionId)
+            owners.removeValue(forKey: connectionId)
             waiters.removeValue(forKey: connectionId)
             return
         }
         let next = pending.removeFirst()
         waiters[connectionId] = pending.isEmpty ? nil : pending
+        owners[connectionId] = next.ticket
         next.continuation.resume()
     }
 }

TablePro/Core/Coordinators/PaginationCoordinator.swift

@@ -303,8 +303,7 @@ final class PaginationCoordinator {
         /// would discard its own rows. It registers as unclaimed work instead, which is what keeps
         /// the titlebar reporting it, and releases that on every exit including cancellation.
         let workToken = parent.tabExecution.beginUnclaimedWork(for: tabId)
-
-        let route = DatabaseManager.shared.executionRoute(for: scope)
+        let isTableTab = parent.tabManager.tabs[idx].tabType == .table
 
         let startedAt = ContinuousClock.Instant.now
         let fetchAllTask = Task { [weak self, parent] in
@@ -314,11 +313,7 @@ final class PaginationCoordinator {
             do {
                 let start = CFAbsoluteTimeGetCurrent()
                 progressLog.info("[fetchAll] executing full query: \(baseQuery.prefix(100), privacy: .public)")
-                let result = try await DatabaseManager.shared.withScopedDriver(
-                    scope: scope,
-                    route: route,
-                    cancellation: .cancellableRead
-                ) { driver in
+                let result = try await parent.withExecutionDriver(scope: scope, isTableTab: isTableTab) { driver in
                     try await driver.executeUserQuery(
                         query: baseQuery,
                         rowCap: nil,

TablePro/Core/Database/DatabaseManager+Health.swift

@@ -138,6 +138,11 @@ extension DatabaseManager {
         let attemptedDriver = session.driver
         await SchemaService.shared.prepareForReload(connectionId: connectionId)
         await DatabaseTreeMetadataService.shared.handleReconnect(connectionId: connectionId)
+        /// A connection that stopped answering has most likely taken its pooled connections with it,
+        /// and a rebuilt tunnel moves every one of them to a new port, so pooled work waits for the
+        /// replacement rather than dialing what is being torn down.
+        MetadataConnectionPool.shared.beginTransportReplacement(connectionId: connectionId)
+        defer { MetadataConnectionPool.shared.endTransportReplacement(connectionId: connectionId) }
 
         do {
             guard let result = try await trackOperation(sessionId: connectionId, operation: {
@@ -355,6 +360,19 @@ extension DatabaseManager {
 
         await SchemaService.shared.prepareForReload(connectionId: sessionId)
         await DatabaseTreeMetadataService.shared.handleReconnect(connectionId: sessionId)
+        /// A pooled connection stands on the effective connection and its own database, so the pool is
+        /// only held back when this rebuilds a tunnel on a new port or recovers a session that had
+        /// stopped answering. A database switch over a live, direct connection changes neither, and
+        /// closing the pool there withdrew an open that another database's table load was waiting on.
+        let replacesPooledTransport = session.connection.activeTunnelKind != nil || session.liveness != .live
+        if replacesPooledTransport {
+            MetadataConnectionPool.shared.beginTransportReplacement(connectionId: sessionId)
+        }
+        defer {
+            if replacesPooledTransport {
+                MetadataConnectionPool.shared.endTransportReplacement(connectionId: sessionId)
+            }
+        }
 
         await stopHealthMonitor(for: sessionId)
 

TablePro/Core/Database/DatabaseManager+ScopedDriver.swift

@@ -15,6 +15,13 @@ enum ScopedDriverRoute: Equatable {
     case unavailable(String)
 }
 
+/// What a table read's turn on the session driver came to.
+private enum TableReadTurn<T: Sendable>: Sendable {
+    case ran(T)
+    /// The route decided once the turn came, which is never the session driver.
+    case moved(ScopedDriverRoute)
+}
+
 extension DatabaseManager {
     /// A metadata read needs no transaction, no temp tables and no cancellation handle,
     /// so it takes a pooled connection and leaves the shared driver where it is.
@@ -76,29 +83,7 @@ extension DatabaseManager {
         cancellation: DriverCancellationPolicy,
         _ body: @Sendable @escaping (DatabaseDriver) async throws -> T
     ) async throws -> T {
-        let leased: @Sendable (DatabaseDriver) async throws -> T
-        if cancellation.isTracked {
-            let connectionId = scope.connectionId
-            let token = UUID()
-            let entry = RunningDriver(driver: nil, policy: cancellation)
-            leased = { driver in
-                await MainActor.run {
-                    DatabaseManager.shared.runningDrivers[connectionId, default: [:]][token] =
-                        entry.adopting(driver)
-                }
-                do {
-                    let value = try await body(driver)
-                    await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) }
-                    return value
-                } catch {
-                    await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) }
-                    throw error
-                }
-            }
-        } else {
-            leased = body
-        }
-
+        let leased = trackedLease(for: scope.connectionId, cancellation: cancellation, body)
         switch route {
         case .unavailable(let message):
             throw DatabaseError.queryFailed(message)
@@ -111,6 +96,67 @@ extension DatabaseManager {
         }
     }
 
+    /// A table tab's read is a SELECT the app built from the tab's own table, so it depends on
+    /// nothing the session holds: no transaction, no temp table, no variable. That makes it the one
+    /// kind of work that can follow a route change it waited through. A database switch on an engine
+    /// that reconnects to perform one moves the browsed database while holding the gate, so a read
+    /// queued behind it for the database being left is re-routed once its turn comes, instead of
+    /// being refused by `pin`. User SQL never comes here: the session it was written against is gone,
+    /// and a refusal is the right answer for it.
+    ///
+    /// The gate is left before the read is dispatched again, so a pooled read never holds it. The
+    /// loop ends: a turn only ever moves the read off the session driver, and neither the pool nor an
+    /// unavailable route queues on the gate again.
+    func withTableReadDriver<T: Sendable>(
+        scope: DatabaseScope,
+        cancellation: DriverCancellationPolicy,
+        _ body: @Sendable @escaping (DatabaseDriver) async throws -> T
+    ) async throws -> T {
+        let leased = trackedLease(for: scope.connectionId, cancellation: cancellation, body)
+        var route = executionRoute(for: scope)
+        while true {
+            switch route {
+            case .unavailable(let message):
+                throw DatabaseError.queryFailed(message)
+            case .pooled:
+                return try await MetadataConnectionPool.shared.withDriver(scope: scope, leased)
+            case .sessionDriver:
+                switch try await withTableReadSessionDriver(scope: scope, leased) {
+                case .ran(let value):
+                    return value
+                case .moved(let decided):
+                    route = decided
+                }
+            }
+        }
+    }
+
+    /// Registers the driver a tracked lease runs on for the length of its body, whichever route it
+    /// took, so Stop reaches the handle the work is actually on.
+    private func trackedLease<T: Sendable>(
+        for connectionId: UUID,
+        cancellation: DriverCancellationPolicy,
+        _ body: @Sendable @escaping (DatabaseDriver) async throws -> T
+    ) -> @Sendable (DatabaseDriver) async throws -> T {
+        guard cancellation.isTracked else { return body }
+        let token = UUID()
+        let entry = RunningDriver(driver: nil, policy: cancellation)
+        return { driver in
+            await MainActor.run {
+                DatabaseManager.shared.runningDrivers[connectionId, default: [:]][token] =
+                    entry.adopting(driver)
+            }
+            do {
+                let value = try await body(driver)
+                await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) }
+                return value
+            } catch {
+                await MainActor.run { DatabaseManager.shared.releaseRunningDriver(token, for: connectionId) }
+                throw error
+            }
+        }
+    }
+
     internal func releaseRunningDriver(_ token: UUID, for connectionId: UUID) {
         runningDrivers[connectionId]?.removeValue(forKey: token)
         if runningDrivers[connectionId]?.isEmpty == true {
@@ -183,26 +229,50 @@ extension DatabaseManager {
         scope: DatabaseScope,
         _ body: @Sendable @escaping (DatabaseDriver) async throws -> T
     ) async throws -> T {
+        try await withSessionDriverTurn(connectionId: scope.connectionId) { driver in
+            try await pin(driver, to: scope)
+            return try await body(driver)
+        }
+    }
+
+    /// The route is asked again through `executionRoute` itself, so it reads the same inputs the
+    /// caller's first answer did, the browsed database among them. The driver's own connection is no
+    /// substitute: it carries the database name the connection resolved to, not the one browsed.
+    private func withTableReadSessionDriver<T: Sendable>(
+        scope: DatabaseScope,
+        _ body: @Sendable @escaping (DatabaseDriver) async throws -> T
+    ) async throws -> TableReadTurn<T> {
+        try await withSessionDriverTurn(connectionId: scope.connectionId) { driver in
+            let route = executionRoute(for: scope)
+            guard route == .sessionDriver else { return .moved(route) }
+            try await pin(driver, to: scope)
+            return .ran(try await body(driver))
+        }
+    }
+
+    private func withSessionDriverTurn<R>(
+        connectionId: UUID,
+        _ turn: (DatabaseDriver) async throws -> R
+    ) async throws -> R {
         /// Outside the gate on purpose. A verification that has to reconnect runs the whole
         /// reconnect, which restores the schema and the database on the new driver, and doing that
         /// while holding the gate would deadlock the very thing waiting to be pinned.
-        await verifyBeforeUse(scope.connectionId)
+        await verifyBeforeUse(connectionId)
         /// A check that failed and could not recover left the driver installed and disconnected,
         /// so the presence of a driver below is not enough. Refusing here is the point of checking
         /// at all: without it the user's own work runs on a handle the app already knows is dead.
-        guard isUsable(scope.connectionId) else {
+        guard isUsable(connectionId) else {
             throw DatabaseError.notConnected
         }
-        return try await sessionDriverGate.withExclusiveAccess(scope.connectionId) {
-            try await trackOperation(sessionId: scope.connectionId) {
+        return try await sessionDriverGate.withExclusiveAccess(connectionId) {
+            try await trackOperation(sessionId: connectionId) {
                 try Task.checkCancellation()
                 /// Asked again once the lease has its turn, because a database switch that held the
                 /// gate can have left the driver the same way while this lease waited.
-                guard isUsable(scope.connectionId), let driver = driver(for: scope.connectionId) else {
+                guard isUsable(connectionId), let driver = driver(for: connectionId) else {
                     throw DatabaseError.notConnected
                 }
-                try await pin(driver, to: scope)
-                return try await body(driver)
+                return try await turn(driver)
             }
         }
     }

TablePro/Core/Database/DatabaseManager+Sessions.swift

@@ -333,13 +333,22 @@ extension DatabaseManager {
 
         if pm?.capabilities.requiresReconnectForDatabaseSwitch == true {
             try await reconnectOntoDatabase(database, for: connectionId)
-        } else if let adapter = driver as? PluginDriverAdapter {
+        } else if driver is PluginDriverAdapter {
             let grouping = pm?.schema.databaseGroupingStrategy ?? .byDatabase
-            try await sessionDriverGate.withExclusiveAccess(connectionId) {
+            let sessionStartedAt = session(for: connectionId)?.connectedAt
+            let adapter = try await sessionDriverGate.withExclusiveAccess(connectionId) {
+                try Task.checkCancellation()
+                guard session(for: connectionId)?.connectedAt == sessionStartedAt else {
+                    throw CancellationError()
+                }
+                guard let adapter = self.driver(for: connectionId) as? PluginDriverAdapter else {
+                    throw DatabaseError.notConnected
+                }
                 try await adapter.switchDatabase(to: database)
                 if grouping == .bySchema {
                     await resetSchema(on: adapter, to: pm?.schema.defaultSchemaName)
                 }
+                return adapter
             }
             updateSession(connectionId) { session in
                 session.browseDatabase = database
@@ -376,10 +385,10 @@ extension DatabaseManager {
     /// one the reconnect was about to disconnect. Counting it as an operation keeps the monitor's
     /// ping and a waiting lease's verification off the driver while it is being replaced.
     ///
-    /// A switch can now wait for its turn, and a disconnect does not fail what is waiting, so the
-    /// session it was asked on is checked once the turn comes. A connection closed and opened again
-    /// in between is a new session, and moving it would switch, or disconnect, a session nobody
-    /// asked this of.
+    /// A switch can now wait for its turn. A disconnect fails every caller still waiting for one,
+    /// and the session the switch was asked on is checked again once the turn comes: a connection
+    /// closed and opened again is a new session, and moving it would switch, or disconnect, a
+    /// session nobody asked this of.
     private func reconnectOntoDatabase(_ database: String, for connectionId: UUID) async throws {
         let sessionStartedAt = session(for: connectionId)?.connectedAt
         try await sessionDriverGate.withExclusiveAccess(connectionId) {
@@ -756,8 +765,11 @@ extension DatabaseManager {
         userRequestedDisconnects.contains(connectionId)
     }
 
+    /// Drains the driver gate in the same step the entry goes, so nothing still queued for this
+    /// session wakes to find a reopened one under the same id and runs there.
     internal func removeSessionEntry(for connectionId: UUID) {
         activeSessions.removeValue(forKey: connectionId)
+        sessionDriverGate.drain(connectionId: connectionId)
         connectionStatusVersions.removeValue(forKey: connectionId)
         forgetVerification(for: connectionId)
         AppEvents.shared.connectionStatusChanged.send(