fix(plugins): cancel every in-flight request in the D1, libSQL and Trino HTTP clients

CHANGELOG.md

@@ -108,6 +108,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Vim Replace mode writing an invisible character for Backspace and keypad Enter. (#2717)
 - Option and Control chords editing text in Vim Normal and Visual mode, and Vim's Ctrl commands never running. (#2717)
 - Line and paragraph separators (U+2028, U+2029) shown as line breaks the database does not see. (#2717)
+- Stop on Cloudflare D1, libSQL and Trino cancelling a sidebar read instead of the running query.
 - Numeric-looking filter values sent unquoted to text columns when a table first opens or after a foreign key jump.
 
 ## [0.73.0] - 2026-09-09

Packages/TableProCore/Sources/TableProTrinoCore/TrinoStatementClient.swift

@@ -11,8 +11,7 @@ public final class TrinoStatementClient: @unchecked Sendable {
     private let config: TrinoClientConfig
     private let session: TrinoSessionState
     private let lock = NSLock()
-    private var _cancelled = false
-    private var _currentNextUri: String?
+    private var running: [ObjectIdentifier: TrinoRunningStatement] = [:]
 
     private static let maxTransientRetries = 5
     private static let logger = Logger(subsystem: "com.TablePro", category: "TrinoStatementClient")
@@ -62,14 +61,14 @@ public final class TrinoStatementClient: @unchecked Sendable {
         }
     }
 
+    /// Stops every statement running on this client, not only the one that started last: the app
+    /// runs sidebar and autocomplete reads on the same client while a query runs. Each statement
+    /// is told to stop, its request in flight is cancelled, and Trino gets one DELETE for it.
     public func cancel() {
-        let uri = lock.withLock { () -> String? in
-            _cancelled = true
-            return _currentNextUri
-        }
-        if let uri {
-            fireDelete(uri)
-        }
+        let statements = lock.withLock { Array(running.values) }
+        statements.forEach { $0.markCancelled() }
+        transport.cancelAll()
+        statements.compactMap { $0.claimRelease() }.forEach(fireDelete)
     }
 
     private struct StatementOutcome {
@@ -83,16 +82,31 @@ public final class TrinoStatementClient: @unchecked Sendable {
         onColumns: ([TrinoColumn]) -> Void,
         onPage: ([[TrinoValue]]) -> Void
     ) async throws -> StatementOutcome {
-        lock.withLock {
-            _cancelled = false
-            _currentNextUri = nil
-        }
         guard let statementURL = config.statementURL else {
             throw TrinoError.invalidConfiguration("Invalid Trino server URL")
         }
+        let statement = TrinoRunningStatement()
+        lock.withLock { running[ObjectIdentifier(statement)] = statement }
+        defer { lock.withLock { running[ObjectIdentifier(statement)] = nil } }
+
+        do {
+            return try await drive(statement, url: statementURL, sql: sql, onColumns: onColumns, onPage: onPage)
+        } catch {
+            releaseIfCancelled(statement)
+            throw error
+        }
+    }
 
+    private func drive(
+        _ statement: TrinoRunningStatement,
+        url statementURL: URL,
+        sql: String,
+        onColumns: ([TrinoColumn]) -> Void,
+        onPage: ([[TrinoValue]]) -> Void
+    ) async throws -> StatementOutcome {
         var httpResponse = try await sendWithRetry(
-            makeRequest(method: .post, url: statementURL, headers: initialHeaders(), body: Data(sql.utf8))
+            makeRequest(method: .post, url: statementURL, headers: initialHeaders(), body: Data(sql.utf8)),
+            for: statement
         )
         var results = try decode(httpResponse)
         session.apply(responseHeaders: httpResponse.headers, protocolHeaders: config.protocolHeaders)
@@ -113,12 +127,14 @@ public final class TrinoStatementClient: @unchecked Sendable {
         var nextUri = results.nextUri
 
         while let uri = nextUri {
-            try abortIfCancelled(currentUri: uri)
-            lock.withLock { _currentNextUri = uri }
+            statement.advance(to: uri)
             guard let nextURL = URL(string: uri) else {
                 throw TrinoError.invalidResponse("Trino returned an invalid nextUri")
             }
-            httpResponse = try await sendWithRetry(makeRequest(method: .get, url: nextURL, headers: followHeaders()))
+            httpResponse = try await sendWithRetry(
+                makeRequest(method: .get, url: nextURL, headers: followHeaders()),
+                for: statement
+            )
             results = try decode(httpResponse)
             session.apply(responseHeaders: httpResponse.headers, protocolHeaders: config.protocolHeaders)
             if let error = results.error {
@@ -141,7 +157,6 @@ public final class TrinoStatementClient: @unchecked Sendable {
             nextUri = results.nextUri
         }
 
-        lock.withLock { _currentNextUri = nil }
         return StatementOutcome(updateType: updateType, updateCount: updateCount, queryId: queryId)
     }
 
@@ -167,16 +182,23 @@ public final class TrinoStatementClient: @unchecked Sendable {
         }
     }
 
-    private func abortIfCancelled(currentUri: String) throws {
-        let cancelled = lock.withLock { _cancelled } || Task.isCancelled
-        guard cancelled else { return }
-        fireDelete(currentUri)
+    private func abortIfCancelled(_ statement: TrinoRunningStatement) throws {
+        guard statement.isCancelled || Task.isCancelled else { return }
         throw TrinoError.cancelled
     }
 
-    private func sendWithRetry(_ request: TrinoHTTPRequest) async throws -> TrinoHTTPResponse {
+    private func releaseIfCancelled(_ statement: TrinoRunningStatement) {
+        guard statement.isCancelled || Task.isCancelled, let uri = statement.claimRelease() else { return }
+        fireDelete(uri)
+    }
+
+    private func sendWithRetry(
+        _ request: TrinoHTTPRequest,
+        for statement: TrinoRunningStatement
+    ) async throws -> TrinoHTTPResponse {
         var attempt = 0
         while true {
+            try abortIfCancelled(statement)
             let response = try await transport.send(request)
             switch response.statusCode {
             case 200...299:
@@ -299,3 +321,30 @@ public final class TrinoStatementClient: @unchecked Sendable {
         String(data: response.body, encoding: .utf8)?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
     }
 }
+
+private final class TrinoRunningStatement: @unchecked Sendable {
+    private let lock = NSLock()
+    private var cancelled = false
+    private var released = false
+    private var nextUri: String?
+
+    var isCancelled: Bool {
+        lock.withLock { cancelled }
+    }
+
+    func markCancelled() {
+        lock.withLock { cancelled = true }
+    }
+
+    func advance(to uri: String) {
+        lock.withLock { nextUri = uri }
+    }
+
+    func claimRelease() -> String? {
+        lock.withLock {
+            guard !released, let nextUri else { return nil }
+            released = true
+            return nextUri
+        }
+    }
+}

Packages/TableProCore/Sources/TableProTrinoCore/TrinoTransport.swift

@@ -75,19 +75,37 @@ public struct TrinoHTTPResponse: Sendable {
 
 public protocol TrinoTransport: Sendable {
     func send(_ request: TrinoHTTPRequest) async throws -> TrinoHTTPResponse
+    func cancelAll()
 }
 
+/// Sends with `URLSession.data(for:delegate:)`, so cancelling the Swift task that awaits a request
+/// cancels its URL task, and keeps every request in flight so `cancelAll` stops each one. A DELETE
+/// is never tracked: it is how a statement tells Trino to stop, and a cancel must not cancel it.
 public final class URLSessionTrinoTransport: NSObject, TrinoTransport, @unchecked Sendable {
     private let session: URLSession
+    private let lock = NSLock()
+    private var inFlight: [ObjectIdentifier: URLSessionTask] = [:]
 
-    public init(tls: TrinoTLSOptions) {
-        let configuration = URLSessionConfiguration.ephemeral
+    public convenience init(tls: TrinoTLSOptions) {
+        self.init(tls: tls, configuration: .ephemeral)
+    }
+
+    init(tls: TrinoTLSOptions, configuration: URLSessionConfiguration) {
         configuration.requestCachePolicy = .reloadIgnoringLocalCacheData
         let delegateProxy = TrinoTLSDelegate(tls: tls)
         self.session = URLSession(configuration: configuration, delegate: delegateProxy, delegateQueue: nil)
         super.init()
     }
 
+    deinit {
+        session.invalidateAndCancel()
+    }
+
+    public func cancelAll() {
+        let tasks = lock.withLock { Array(inFlight.values) }
+        tasks.forEach { $0.cancel() }
+    }
+
     public func send(_ request: TrinoHTTPRequest) async throws -> TrinoHTTPResponse {
         var urlRequest = URLRequest(url: request.url)
         urlRequest.httpMethod = request.method.rawValue
@@ -97,24 +115,18 @@ public final class URLSessionTrinoTransport: NSObject, TrinoTransport, @unchecke
             urlRequest.setValue(value, forHTTPHeaderField: name)
         }
 
-        let (data, response) = try await withCheckedThrowingContinuation {
-            (continuation: CheckedContinuation<(Data, URLResponse), Error>) in
-            let task = session.dataTask(with: urlRequest) { data, response, error in
-                if let error {
-                    if (error as? URLError)?.code == .cancelled {
-                        continuation.resume(throwing: TrinoError.cancelled)
-                    } else {
-                        continuation.resume(throwing: TrinoError.transport(error.localizedDescription))
-                    }
-                    return
-                }
-                guard let data, let response else {
-                    continuation.resume(throwing: TrinoError.invalidResponse("Empty response from Trino"))
-                    return
-                }
-                continuation.resume(returning: (data, response))
-            }
-            task.resume()
+        let tracker = request.method == .delete ? nil : TrinoTaskTracker(transport: self)
+        defer { tracker?.finish() }
+        let data: Data
+        let response: URLResponse
+        do {
+            (data, response) = try await session.data(for: urlRequest, delegate: tracker)
+        } catch let error as URLError where error.code == .cancelled {
+            throw TrinoError.cancelled
+        } catch is CancellationError {
+            throw TrinoError.cancelled
+        } catch {
+            throw TrinoError.transport(error.localizedDescription)
         }
 
         guard let httpResponse = response as? HTTPURLResponse else {
@@ -126,6 +138,38 @@ public final class URLSessionTrinoTransport: NSObject, TrinoTransport, @unchecke
             body: data
         )
     }
+
+    var inFlightCount: Int {
+        lock.withLock { inFlight.count }
+    }
+
+    fileprivate func register(_ task: URLSessionTask) {
+        lock.withLock { inFlight[ObjectIdentifier(task)] = task }
+    }
+
+    fileprivate func unregister(_ task: URLSessionTask) {
+        lock.withLock { inFlight[ObjectIdentifier(task)] = nil }
+    }
+}
+
+private final class TrinoTaskTracker: NSObject, URLSessionTaskDelegate, @unchecked Sendable {
+    private weak var transport: URLSessionTrinoTransport?
+    private let lock = NSLock()
+    private var task: URLSessionTask?
+
+    init(transport: URLSessionTrinoTransport) {
+        self.transport = transport
+    }
+
+    func urlSession(_ session: URLSession, didCreateTask task: URLSessionTask) {
+        lock.withLock { self.task = task }
+        transport?.register(task)
+    }
+
+    func finish() {
+        guard let task = lock.withLock({ task }) else { return }
+        transport?.unregister(task)
+    }
 }
 
 private final class TrinoTLSDelegate: NSObject, URLSessionDelegate {

Packages/TableProCore/Tests/TableProTrinoCoreTests/TrinoStatementClientTests.swift

@@ -1,5 +1,5 @@
-import XCTest
 @testable import TableProTrinoCore
+import XCTest
 
 final class TrinoStatementClientTests: XCTestCase {
     private func makeClient(_ transport: StubTransport, session: TrinoSessionState = TrinoSessionState(catalog: "c", schema: "s")) -> TrinoStatementClient {
@@ -136,5 +136,6 @@ final class TrinoStatementClientTests: XCTestCase {
         } catch let error as TrinoError {
             XCTAssertEqual(error, .cancelled)
         }
+        XCTAssertEqual(transport.cancelAllCount, 1)
     }
 }

Packages/TableProCore/Tests/TableProTrinoCoreTests/TrinoTestSupport.swift

@@ -11,6 +11,7 @@ final class StubTransport: TrinoTransport, @unchecked Sendable {
     private let lock = NSLock()
     private var queue: [Canned]
     private var recorded: [TrinoHTTPRequest] = []
+    private var cancelAllCalls = 0
     var onSend: ((TrinoHTTPRequest, Int) -> Void)?
 
     init(_ responses: [Canned]) {
@@ -21,6 +22,14 @@ final class StubTransport: TrinoTransport, @unchecked Sendable {
         lock.withLock { recorded }
     }
 
+    var cancelAllCount: Int {
+        lock.withLock { cancelAllCalls }
+    }
+
+    func cancelAll() {
+        lock.withLock { cancelAllCalls += 1 }
+    }
+
     func send(_ request: TrinoHTTPRequest) async throws -> TrinoHTTPResponse {
         let index = lock.withLock { () -> Int in
             recorded.append(request)