Merge pull request #273 from datlechin/feat/mssql-cancel-query

feat: implement MSSQL cancelQuery and applyQueryTimeout

Author: datlechin

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Copy as INSERT/UPDATE SQL statements from data grid context menu
 - Plugin download count display in Browse Plugins — fetched from GitHub Releases API and cached for 1 hour
+- MSSQL query cancellation (`cancelQuery`) and lock timeout (`applyQueryTimeout`) support
 
 ### Fixed
 

Plugins/MSSQLDriverPlugin/MSSQLPlugin.swift

@@ -92,6 +92,7 @@ private final class FreeTDSConnection: @unchecked Sendable {
     private let database: String
     private let lock = NSLock()
     private var _isConnected = false
+    private var _isCancelled = false
 
     var isConnected: Bool {
         lock.lock()
@@ -172,6 +173,16 @@ private final class FreeTDSConnection: @unchecked Sendable {
         }
     }
 
+    func cancelCurrentQuery() {
+        lock.lock()
+        _isCancelled = true
+        let proc = dbproc
+        lock.unlock()
+
+        guard let proc else { return }
+        dbcancel(proc)
+    }
+
     func executeQuery(_ query: String) async throws -> FreeTDSQueryResult {
         let queryToRun = String(query)
         return try await pluginDispatchAsync(on: queue) { [self] in
@@ -186,6 +197,10 @@ private final class FreeTDSConnection: @unchecked Sendable {
 
         _ = dbcanquery(proc)
 
+        lock.lock()
+        _isCancelled = false
+        lock.unlock()
+
         freetdsLastError = ""
         if dbcmd(proc, query) == FAIL {
             throw MSSQLPluginError.queryFailed("Failed to prepare query")
@@ -202,6 +217,14 @@ private final class FreeTDSConnection: @unchecked Sendable {
         var truncated = false
 
         while true {
+            lock.lock()
+            let cancelledBetweenResults = _isCancelled
+            if cancelledBetweenResults { _isCancelled = false }
+            lock.unlock()
+            if cancelledBetweenResults {
+                throw MSSQLPluginError.queryFailed("Query cancelled")
+            }
+
             let resCode = dbresults(proc)
             if resCode == FAIL {
                 throw MSSQLPluginError.queryFailed("Query execution failed")
@@ -232,6 +255,14 @@ private final class FreeTDSConnection: @unchecked Sendable {
                 if rowCode == Int32(NO_MORE_ROWS) { break }
                 if rowCode == FAIL { break }
 
+                lock.lock()
+                let cancelled = _isCancelled
+                if cancelled { _isCancelled = false }
+                lock.unlock()
+                if cancelled {
+                    throw MSSQLPluginError.queryFailed("Query cancelled")
+                }
+
                 var row: [String?] = []
                 for i in 1...numCols {
                     let len = dbdatlen(proc, Int32(i))
@@ -386,6 +417,16 @@ final class MSSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
         )
     }
 
+    func cancelQuery() throws {
+        freeTDSConn?.cancelCurrentQuery()
+    }
+
+    func applyQueryTimeout(_ seconds: Int) async throws {
+        guard seconds > 0 else { return }
+        let ms = seconds * 1_000
+        _ = try await execute(query: "SET LOCK_TIMEOUT \(ms)")
+    }
+
     func executeParameterized(query: String, parameters: [String?]) async throws -> PluginQueryResult {
         guard !parameters.isEmpty else {
             return try await execute(query: query)

TableProTests/Core/Database/MSSQLDriverTests.swift

@@ -14,6 +14,10 @@ import Testing
 
 private final class MockMSSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
     private var schema: String?
+    var cancelQueryCallCount = 0
+    var applyQueryTimeoutValues: [Int] = []
+    var executedQueries: [String] = []
+    var shouldFailExecute = true
 
     init(initialSchema: String?) {
         schema = initialSchema
@@ -29,12 +33,25 @@ private final class MockMSSQLPluginDriver: PluginDatabaseDriver, @unchecked Send
     func connect() async throws {}
     func disconnect() {}
 
+    func cancelQuery() throws {
+        cancelQueryCallCount += 1
+    }
+
+    func applyQueryTimeout(_ seconds: Int) async throws {
+        applyQueryTimeoutValues.append(seconds)
+        executedQueries.append("SET LOCK_TIMEOUT \(seconds * 1_000)")
+    }
+
     func execute(query: String) async throws -> PluginQueryResult {
-        throw NSError(
-            domain: "MockMSSQLPluginDriver",
-            code: -1,
-            userInfo: [NSLocalizedDescriptionKey: "Not connected"]
-        )
+        executedQueries.append(query)
+        if shouldFailExecute {
+            throw NSError(
+                domain: "MockMSSQLPluginDriver",
+                code: -1,
+                userInfo: [NSLocalizedDescriptionKey: "Not connected"]
+            )
+        }
+        return PluginQueryResult(columns: [], columnTypeNames: [], rows: [], rowsAffected: 0, executionTime: 0)
     }
 
     func fetchTables(schema: String?) async throws -> [PluginTableInfo] { [] }
@@ -64,10 +81,16 @@ struct MSSQLDriverTests {
     }
 
     private func makeAdapter(mssqlSchema: String? = nil) -> PluginDriverAdapter {
+        let (adapter, _) = makeAdapterWithMock(mssqlSchema: mssqlSchema)
+        return adapter
+    }
+
+    private func makeAdapterWithMock(mssqlSchema: String? = nil) -> (PluginDriverAdapter, MockMSSQLPluginDriver) {
         let conn = makeConnection(mssqlSchema: mssqlSchema)
         let effectiveSchema: String? = if let s = mssqlSchema, !s.isEmpty { s } else { "dbo" }
-        let pluginDriver = MockMSSQLPluginDriver(initialSchema: effectiveSchema)
-        return PluginDriverAdapter(connection: conn, pluginDriver: pluginDriver)
+        let mock = MockMSSQLPluginDriver(initialSchema: effectiveSchema)
+        let adapter = PluginDriverAdapter(connection: conn, pluginDriver: mock)
+        return (adapter, mock)
     }
 
     // MARK: - Initialization Tests
@@ -147,4 +170,49 @@ struct MSSQLDriverTests {
             _ = try await adapter.execute(query: "SELECT 1")
         }
     }
+
+    // MARK: - cancelQuery Tests
+
+    @Test("cancelQuery delegates to plugin driver")
+    func cancelQueryDelegatesToPlugin() throws {
+        let (adapter, mock) = makeAdapterWithMock()
+        try adapter.cancelQuery()
+        #expect(mock.cancelQueryCallCount == 1)
+    }
+
+    @Test("cancelQuery can be called multiple times")
+    func cancelQueryMultipleCalls() throws {
+        let (adapter, mock) = makeAdapterWithMock()
+        try adapter.cancelQuery()
+        try adapter.cancelQuery()
+        try adapter.cancelQuery()
+        #expect(mock.cancelQueryCallCount == 3)
+    }
+
+    // MARK: - applyQueryTimeout Tests
+
+    @Test("applyQueryTimeout delegates to plugin driver with correct value")
+    func applyQueryTimeoutDelegates() async throws {
+        let (adapter, mock) = makeAdapterWithMock()
+        mock.shouldFailExecute = false
+        try await adapter.applyQueryTimeout(30)
+        #expect(mock.applyQueryTimeoutValues == [30])
+    }
+
+    @Test("applyQueryTimeout with zero is handled by plugin")
+    func applyQueryTimeoutZero() async throws {
+        let (adapter, mock) = makeAdapterWithMock()
+        mock.shouldFailExecute = false
+        try await adapter.applyQueryTimeout(0)
+        #expect(mock.applyQueryTimeoutValues == [0])
+    }
+
+    @Test("applyQueryTimeout with different values records each call")
+    func applyQueryTimeoutMultipleCalls() async throws {
+        let (adapter, mock) = makeAdapterWithMock()
+        mock.shouldFailExecute = false
+        try await adapter.applyQueryTimeout(10)
+        try await adapter.applyQueryTimeout(60)
+        #expect(mock.applyQueryTimeoutValues == [10, 60])
+    }
 }