From b07eeb5d38cf2dc5421a9563f5a8d266ab1427b4 Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Tue, 11 Feb 2025 14:00:08 +0100 Subject: [PATCH 1/5] Improve transaction handling --- .../Connection/PostgresConnection.swift | 50 +++++++++++++++++++ .../New/PostgresTransactionError.swift | 21 ++++++++ Sources/PostgresNIO/Pool/PostgresClient.swift | 38 +++++++------- 3 files changed, 92 insertions(+), 17 deletions(-) create mode 100644 Sources/PostgresNIO/New/PostgresTransactionError.swift diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/Sources/PostgresNIO/Connection/PostgresConnection.swift index 229cd647..7f52ae1c 100644 --- a/Sources/PostgresNIO/Connection/PostgresConnection.swift +++ b/Sources/PostgresNIO/Connection/PostgresConnection.swift @@ -530,6 +530,56 @@ extension PostgresConnection { throw error // rethrow with more metadata } } + + /// Puts the connection into an open transaction state, for the provided `closure`'s lifetime. + /// + /// The function starts a transaction by running a `BEGIN` query on the connection against the database. It then + /// lends the connection to the user provided closure. The user can then modify the database as they wish. If the user + /// provided closure returns successfully, the function will attempt to commit the changes by running a `COMMIT` + /// query against the database. If the user provided closure throws an error, the function will attempt to rollback the + /// changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! + /// - Returns: The closure's return value. + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + _ process: (PostgresConnection) async throws -> Result + ) async throws -> Result { + do { + try await self.query("BEGIN;", logger: logger) + } catch { + throw PostgresTransactionError(file: file, line: line, beginError: error) + } + + var closureHasFinished: Bool = false + do { + let value = try await process(self) + closureHasFinished = true + try await self.query("COMMIT;", logger: logger) + return value + } catch { + var transactionError = PostgresTransactionError(file: file, line: line) + if !closureHasFinished { + transactionError.closureError = error + do { + try await self.query("ROLLBACK;", logger: logger) + } catch { + transactionError.rollbackError = error + } + } else { + transactionError.commitError = error + } + + throw transactionError + } + } } // MARK: EventLoopFuture interface diff --git a/Sources/PostgresNIO/New/PostgresTransactionError.swift b/Sources/PostgresNIO/New/PostgresTransactionError.swift new file mode 100644 index 00000000..35038446 --- /dev/null +++ b/Sources/PostgresNIO/New/PostgresTransactionError.swift @@ -0,0 +1,21 @@ +/// A wrapper around the errors that can occur during a transaction. +public struct PostgresTransactionError: Error { + + /// The file in which the transaction was started + public var file: String + /// The line in which the transaction was started + public var line: Int + + /// The error thrown when running the `BEGIN` query + public var beginError: Error? + /// The error thrown in the transaction closure + public var closureError: Error? + + /// The error thrown while rolling the transaction back. If the ``closureError`` is set, + /// but the ``rollbackError`` is empty, the rollback was successful. If the ``rollbackError`` + /// is set, the rollback failed. + public var rollbackError: Error? + + /// The error thrown while commiting the transaction. + public var commitError: Error? +} diff --git a/Sources/PostgresNIO/Pool/PostgresClient.swift b/Sources/PostgresNIO/Pool/PostgresClient.swift index e9e947ef..bfa67a60 100644 --- a/Sources/PostgresNIO/Pool/PostgresClient.swift +++ b/Sources/PostgresNIO/Pool/PostgresClient.swift @@ -308,25 +308,29 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { return try await closure(connection) } - /// Lease a connection for the provided `closure`'s lifetime. - /// A transation starts with call to withConnection - /// A transaction should end with a call to COMMIT or ROLLBACK - /// COMMIT is called upon successful completion and ROLLBACK is called should any steps fail + /// Lease a connection, which is in an open transaction state, for the provided `closure`'s lifetime. /// - /// - Parameter closure: A closure that uses the passed `PostgresConnection`. The closure **must not** capture - /// the provided `PostgresConnection`. + /// The function leases a connection from the underlying connection pool and starts a transaction by running a `BEGIN` + /// query on the leased connection against the database. It then lends the connection to the user provided closure. + /// The user can then modify the database as they wish. If the user provided closure returns successfully, the function + /// will attempt to commit the changes by running a `COMMIT` query against the database. If the user provided closure + /// throws an error, the function will attempt to rollback the changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! /// - Returns: The closure's return value. - public func withTransaction(_ process: (PostgresConnection) async throws -> Result) async throws -> Result { - try await withConnection { connection in - try await connection.query("BEGIN;", logger: self.backgroundLogger) - do { - let value = try await process(connection) - try await connection.query("COMMIT;", logger: self.backgroundLogger) - return value - } catch { - try await connection.query("ROLLBACK;", logger: self.backgroundLogger) - throw error - } + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + _ closure: (PostgresConnection) async throws -> Result + ) async throws -> Result { + try await self.withConnection { connection in + try await connection.withTransaction(logger: logger, file: file, line: line, closure) } } From 604a1b1d99f1f657d088e65239bf09eeb49e3ddc Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Tue, 11 Feb 2025 16:51:28 +0100 Subject: [PATCH 2/5] Don't jump in actors. --- .../Connection/PostgresConnection.swift | 5 ++-- Sources/PostgresNIO/Pool/PostgresClient.swift | 24 ++++++++++++++++--- .../PostgresClientTests.swift | 4 ++-- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/Sources/PostgresNIO/Connection/PostgresConnection.swift index 7f52ae1c..fc48fa31 100644 --- a/Sources/PostgresNIO/Connection/PostgresConnection.swift +++ b/Sources/PostgresNIO/Connection/PostgresConnection.swift @@ -550,8 +550,9 @@ extension PostgresConnection { logger: Logger, file: String = #file, line: Int = #line, - _ process: (PostgresConnection) async throws -> Result - ) async throws -> Result { + isolation: isolated (any Actor)? = #isolation, + _ process: (PostgresConnection) async throws -> sending Result + ) async throws -> sending Result { do { try await self.query("BEGIN;", logger: logger) } catch { diff --git a/Sources/PostgresNIO/Pool/PostgresClient.swift b/Sources/PostgresNIO/Pool/PostgresClient.swift index bfa67a60..db4c46fd 100644 --- a/Sources/PostgresNIO/Pool/PostgresClient.swift +++ b/Sources/PostgresNIO/Pool/PostgresClient.swift @@ -300,6 +300,7 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { /// - Parameter closure: A closure that uses the passed `PostgresConnection`. The closure **must not** capture /// the provided `PostgresConnection`. /// - Returns: The closure's return value. + @_disfavoredOverload public func withConnection(_ closure: (PostgresConnection) async throws -> Result) async throws -> Result { let connection = try await self.leaseConnection() @@ -307,7 +308,23 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { return try await closure(connection) } - + + /// Lease a connection for the provided `closure`'s lifetime. + /// + /// - Parameter closure: A closure that uses the passed `PostgresConnection`. The closure **must not** capture + /// the provided `PostgresConnection`. + /// - Returns: The closure's return value. + public func withConnection( + isolation: isolated (any Actor)? = #isolation, + _ closure: (PostgresConnection) async throws -> sending Result + ) async throws -> sending Result { + let connection = try await self.leaseConnection() + + defer { self.pool.releaseConnection(connection) } + + return try await closure(connection) + } + /// Lease a connection, which is in an open transaction state, for the provided `closure`'s lifetime. /// /// The function leases a connection from the underlying connection pool and starts a transaction by running a `BEGIN` @@ -327,8 +344,9 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { logger: Logger, file: String = #file, line: Int = #line, - _ closure: (PostgresConnection) async throws -> Result - ) async throws -> Result { + isolation: isolated (any Actor)? = #isolation, + _ closure: (PostgresConnection) async throws -> sending Result + ) async throws -> sending Result { try await self.withConnection { connection in try await connection.withTransaction(logger: logger, file: file, line: line, closure) } diff --git a/Tests/IntegrationTests/PostgresClientTests.swift b/Tests/IntegrationTests/PostgresClientTests.swift index 167ba298..b4e0c964 100644 --- a/Tests/IntegrationTests/PostgresClientTests.swift +++ b/Tests/IntegrationTests/PostgresClientTests.swift @@ -77,7 +77,7 @@ final class PostgresClientTests: XCTestCase { for _ in 0.. Date: Thu, 13 Feb 2025 10:23:49 +0100 Subject: [PATCH 3/5] Compile pre 6.0 --- .../Connection/PostgresConnection.swift | 52 +++++++++++++++++++ Sources/PostgresNIO/Pool/PostgresClient.swift | 29 +++++++++++ 2 files changed, 81 insertions(+) diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/Sources/PostgresNIO/Connection/PostgresConnection.swift index fc48fa31..36b9a355 100644 --- a/Sources/PostgresNIO/Connection/PostgresConnection.swift +++ b/Sources/PostgresNIO/Connection/PostgresConnection.swift @@ -531,6 +531,7 @@ extension PostgresConnection { } } + #if compiler(>=6.0) /// Puts the connection into an open transaction state, for the provided `closure`'s lifetime. /// /// The function starts a transaction by running a `BEGIN` query on the connection against the database. It then @@ -581,6 +582,57 @@ extension PostgresConnection { throw transactionError } } + #else + /// Puts the connection into an open transaction state, for the provided `closure`'s lifetime. + /// + /// The function starts a transaction by running a `BEGIN` query on the connection against the database. It then + /// lends the connection to the user provided closure. The user can then modify the database as they wish. If the user + /// provided closure returns successfully, the function will attempt to commit the changes by running a `COMMIT` + /// query against the database. If the user provided closure throws an error, the function will attempt to rollback the + /// changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! + /// - Returns: The closure's return value. + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + _ process: (PostgresConnection) async throws -> Result + ) async throws -> Result { + do { + try await self.query("BEGIN;", logger: logger) + } catch { + throw PostgresTransactionError(file: file, line: line, beginError: error) + } + + var closureHasFinished: Bool = false + do { + let value = try await process(self) + closureHasFinished = true + try await self.query("COMMIT;", logger: logger) + return value + } catch { + var transactionError = PostgresTransactionError(file: file, line: line) + if !closureHasFinished { + transactionError.closureError = error + do { + try await self.query("ROLLBACK;", logger: logger) + } catch { + transactionError.rollbackError = error + } + } else { + transactionError.commitError = error + } + + throw transactionError + } + } + #endif } // MARK: EventLoopFuture interface diff --git a/Sources/PostgresNIO/Pool/PostgresClient.swift b/Sources/PostgresNIO/Pool/PostgresClient.swift index db4c46fd..0ccd969b 100644 --- a/Sources/PostgresNIO/Pool/PostgresClient.swift +++ b/Sources/PostgresNIO/Pool/PostgresClient.swift @@ -309,6 +309,7 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { return try await closure(connection) } + #if compiler(>=6.0) /// Lease a connection for the provided `closure`'s lifetime. /// /// - Parameter closure: A closure that uses the passed `PostgresConnection`. The closure **must not** capture @@ -351,6 +352,34 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { try await connection.withTransaction(logger: logger, file: file, line: line, closure) } } + #else + + /// Lease a connection, which is in an open transaction state, for the provided `closure`'s lifetime. + /// + /// The function leases a connection from the underlying connection pool and starts a transaction by running a `BEGIN` + /// query on the leased connection against the database. It then lends the connection to the user provided closure. + /// The user can then modify the database as they wish. If the user provided closure returns successfully, the function + /// will attempt to commit the changes by running a `COMMIT` query against the database. If the user provided closure + /// throws an error, the function will attempt to rollback the changes made within the closure. + /// + /// - Parameters: + /// - logger: The `Logger` to log into for the transaction. + /// - file: The file, the transaction was started in. Used for better error reporting. + /// - line: The line, the transaction was started in. Used for better error reporting. + /// - closure: The user provided code to modify the database. Use the provided connection to run queries. + /// The connection must stay in the transaction mode. Otherwise this method will throw! + /// - Returns: The closure's return value. + public func withTransaction( + logger: Logger, + file: String = #file, + line: Int = #line, + _ closure: (PostgresConnection) async throws -> Result + ) async throws -> Result { + try await self.withConnection { connection in + try await connection.withTransaction(logger: logger, file: file, line: line, closure) + } + } + #endif /// Run a query on the Postgres server the client is connected to. /// From ee2fdb19e4e4c3c06a0d7d8ab1c4b8130ccbb597 Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Thu, 13 Feb 2025 10:52:05 +0100 Subject: [PATCH 4/5] Fix 5.10? --- .../PostgresNIO/Connection/PostgresConnection.swift | 5 +++-- Sources/PostgresNIO/Pool/PostgresClient.swift | 11 ++++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/Sources/PostgresNIO/Connection/PostgresConnection.swift index 36b9a355..e267d8f9 100644 --- a/Sources/PostgresNIO/Connection/PostgresConnection.swift +++ b/Sources/PostgresNIO/Connection/PostgresConnection.swift @@ -552,8 +552,9 @@ extension PostgresConnection { file: String = #file, line: Int = #line, isolation: isolated (any Actor)? = #isolation, - _ process: (PostgresConnection) async throws -> sending Result - ) async throws -> sending Result { + // DO NOT FIX THE WHITESPACE IN THE NEXT LINE UNTIL 5.10 IS UNSUPPORTED + // https://github.com/swiftlang/swift/issues/79285 + _ process: (PostgresConnection) async throws -> sending Result) async throws -> sending Result { do { try await self.query("BEGIN;", logger: logger) } catch { diff --git a/Sources/PostgresNIO/Pool/PostgresClient.swift b/Sources/PostgresNIO/Pool/PostgresClient.swift index 0ccd969b..d54e34eb 100644 --- a/Sources/PostgresNIO/Pool/PostgresClient.swift +++ b/Sources/PostgresNIO/Pool/PostgresClient.swift @@ -293,7 +293,6 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { return ConnectionAndMetadata(connection: connection, maximalStreamsOnConnection: 1) } } - /// Lease a connection for the provided `closure`'s lifetime. /// @@ -317,8 +316,9 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { /// - Returns: The closure's return value. public func withConnection( isolation: isolated (any Actor)? = #isolation, - _ closure: (PostgresConnection) async throws -> sending Result - ) async throws -> sending Result { + // DO NOT FIX THE WHITESPACE IN THE NEXT LINE UNTIL 5.10 IS UNSUPPORTED + // https://github.com/swiftlang/swift/issues/79285 + _ closure: (PostgresConnection) async throws -> sending Result) async throws -> sending Result { let connection = try await self.leaseConnection() defer { self.pool.releaseConnection(connection) } @@ -346,8 +346,9 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service { file: String = #file, line: Int = #line, isolation: isolated (any Actor)? = #isolation, - _ closure: (PostgresConnection) async throws -> sending Result - ) async throws -> sending Result { + // DO NOT FIX THE WHITESPACE IN THE NEXT LINE UNTIL 5.10 IS UNSUPPORTED + // https://github.com/swiftlang/swift/issues/79285 + _ closure: (PostgresConnection) async throws -> sending Result) async throws -> sending Result { try await self.withConnection { connection in try await connection.withTransaction(logger: logger, file: file, line: line, closure) } From 938bedf56257be9448f70c931bcd3a4eb8585abb Mon Sep 17 00:00:00 2001 From: Fabian Fett Date: Thu, 13 Feb 2025 10:59:27 +0100 Subject: [PATCH 5/5] Fix test --- Tests/IntegrationTests/PostgresClientTests.swift | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Tests/IntegrationTests/PostgresClientTests.swift b/Tests/IntegrationTests/PostgresClientTests.swift index b4e0c964..34a8ad2a 100644 --- a/Tests/IntegrationTests/PostgresClientTests.swift +++ b/Tests/IntegrationTests/PostgresClientTests.swift @@ -120,10 +120,10 @@ final class PostgresClientTests: XCTestCase { } } catch { XCTAssertNotNil(error) - guard let error = error as? PSQLError else { return XCTFail("Unexpected error type") } - - XCTAssertEqual(error.code, .server) - XCTAssertEqual(error.serverInfo?[.severity], "ERROR") + guard let error = error as? PostgresTransactionError else { return XCTFail("Unexpected error type: \(error)") } + + XCTAssertEqual((error.closureError as? PSQLError)?.code, .server) + XCTAssertEqual((error.closureError as? PSQLError)?.serverInfo?[.severity], "ERROR") } }