Merge branch 'main' into mmbm-row-seq-expose-metadata

Author: MahdiBM

Package.swift

@@ -20,7 +20,7 @@ let package = Package(
     dependencies: [
         .package(url: "https://github.com/apple/swift-atomics.git", from: "1.2.0"),
         .package(url: "https://github.com/apple/swift-collections.git", from: "1.0.4"),
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.59.0"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.81.0"),
         .package(url: "https://github.com/apple/swift-nio-transport-services.git", from: "1.19.0"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.25.0"),
         .package(url: "https://github.com/apple/swift-crypto.git", "2.0.0" ..< "4.0.0"),

Sources/ConnectionPoolModule/NIOLock.swift

@@ -24,6 +24,13 @@ import WinSDK
 import Glibc
 #elseif canImport(Musl)
 import Musl
+#elseif canImport(Bionic)
+import Bionic
+#elseif canImport(WASILibc)
+import WASILibc
+#if canImport(wasi_pthread)
+import wasi_pthread
+#endif
 #else
 #error("The concurrency NIOLock module was unable to identify your C library.")
 #endif
@@ -37,16 +44,16 @@ typealias LockPrimitive = pthread_mutex_t
 #endif
 
 @usableFromInline
-enum LockOperations { }
+enum LockOperations {}
 
 extension LockOperations {
     @inlinable
     static func create(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
         mutex.assertValidAlignment()
 
-#if os(Windows)
+        #if os(Windows)
         InitializeSRWLock(mutex)
-#else
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
         var attr = pthread_mutexattr_t()
         pthread_mutexattr_init(&attr)
         debugOnly {
@@ -55,43 +62,43 @@ extension LockOperations {
 
         let err = pthread_mutex_init(mutex, &attr)
         precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
-#endif
+        #endif
     }
 
     @inlinable
     static func destroy(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
         mutex.assertValidAlignment()
 
-#if os(Windows)
+        #if os(Windows)
         // SRWLOCK does not need to be free'd
-#else
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
         let err = pthread_mutex_destroy(mutex)
         precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
-#endif
+        #endif
     }
 
     @inlinable
     static func lock(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
         mutex.assertValidAlignment()
 
-#if os(Windows)
+        #if os(Windows)
         AcquireSRWLockExclusive(mutex)
-#else
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
         let err = pthread_mutex_lock(mutex)
         precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
-#endif
+        #endif
     }
 
     @inlinable
     static func unlock(_ mutex: UnsafeMutablePointer<LockPrimitive>) {
         mutex.assertValidAlignment()
 
-#if os(Windows)
+        #if os(Windows)
         ReleaseSRWLockExclusive(mutex)
-#else
+        #elseif (compiler(<6.1) && !os(WASI)) || (compiler(>=6.1) && _runtime(_multithreaded))
         let err = pthread_mutex_unlock(mutex)
         precondition(err == 0, "\(#function) failed in pthread_mutex with error \(err)")
-#endif
+        #endif
     }
 }
 
@@ -129,9 +136,11 @@ final class LockStorage<Value>: ManagedBuffer<Value, LockPrimitive> {
     @inlinable
     static func create(value: Value) -> Self {
         let buffer = Self.create(minimumCapacity: 1) { _ in
-            return value
+            value
         }
-        // Avoid 'unsafeDowncast' as there is a miscompilation on 5.10.
+        // Intentionally using a force cast here to avoid a miss compiliation in 5.10.
+        // This is as fast as an unsafeDownCast since ManagedBuffer is inlined and the optimizer
+        // can eliminate the upcast/downcast pair
         let storage = buffer as! Self
 
         storage.withUnsafeMutablePointers { _, lockPtr in
@@ -165,7 +174,7 @@ final class LockStorage<Value>: ManagedBuffer<Value, LockPrimitive> {
     @inlinable
     func withLockPrimitive<T>(_ body: (UnsafeMutablePointer<LockPrimitive>) throws -> T) rethrows -> T {
         try self.withUnsafeMutablePointerToElements { lockPtr in
-            return try body(lockPtr)
+            try body(lockPtr)
         }
     }
 
@@ -179,17 +188,14 @@ final class LockStorage<Value>: ManagedBuffer<Value, LockPrimitive> {
     }
 }
 
-extension LockStorage: @unchecked Sendable { }
-
 /// A threading lock based on `libpthread` instead of `libdispatch`.
 ///
-/// - note: ``NIOLock`` has reference semantics.
+/// - Note: ``NIOLock`` has reference semantics.
 ///
 /// This object provides a lock on top of a single `pthread_mutex_t`. This kind
 /// of lock is safe to use with `libpthread`-based threading models, such as the
 /// one used by NIO. On Windows, the lock is based on the substantially similar
 /// `SRWLOCK` type.
-@usableFromInline
 struct NIOLock {
     @usableFromInline
     internal let _storage: LockStorage<Void>
@@ -220,7 +226,7 @@ struct NIOLock {
 
     @inlinable
     internal func withLockPrimitive<T>(_ body: (UnsafeMutablePointer<LockPrimitive>) throws -> T) rethrows -> T {
-        return try self._storage.withLockPrimitive(body)
+        try self._storage.withLockPrimitive(body)
     }
 }
 
@@ -243,12 +249,12 @@ extension NIOLock {
     }
 
     @inlinable
-    func withLockVoid(_ body: () throws -> Void) rethrows -> Void {
+    func withLockVoid(_ body: () throws -> Void) rethrows {
         try self.withLock(body)
     }
 }
 
-extension NIOLock: Sendable {}
+extension NIOLock: @unchecked Sendable {}
 
 extension UnsafeMutablePointer {
     @inlinable
@@ -264,6 +270,10 @@ extension UnsafeMutablePointer {
 /// https://forums.swift.org/t/support-debug-only-code/11037 for a discussion.
 @inlinable
 internal func debugOnly(_ body: () -> Void) {
-    // FIXME: duplicated with NIO.
-    assert({ body(); return true }())
+    assert(
+        {
+            body()
+            return true
+        }()
+    )
 }

Sources/ConnectionPoolModule/NIOLockedValueBox.swift

@@ -17,7 +17,7 @@
 
 /// Provides locked access to `Value`.
 ///
-/// - note: ``NIOLockedValueBox`` has reference semantics and holds the `Value`
+/// - Note: ``NIOLockedValueBox`` has reference semantics and holds the `Value`
 ///         alongside a lock behind a reference.
 ///
 /// This is no different than creating a ``Lock`` and protecting all
@@ -39,8 +39,48 @@ struct NIOLockedValueBox<Value> {
     /// Access the `Value`, allowing mutation of it.
     @inlinable
     func withLockedValue<T>(_ mutate: (inout Value) throws -> T) rethrows -> T {
-        return try self._storage.withLockedValue(mutate)
+        try self._storage.withLockedValue(mutate)
+    }
+
+    /// Provides an unsafe view over the lock and its value.
+    ///
+    /// This can be beneficial when you require fine grained control over the lock in some
+    /// situations but don't want lose the benefits of ``withLockedValue(_:)`` in others by
+    /// switching to ``NIOLock``.
+    var unsafe: Unsafe {
+        Unsafe(_storage: self._storage)
+    }
+
+    /// Provides an unsafe view over the lock and its value.
+    struct Unsafe {
+        @usableFromInline
+        let _storage: LockStorage<Value>
+
+        /// Manually acquire the lock.
+        @inlinable
+        func lock() {
+            self._storage.lock()
+        }
+
+        /// Manually release the lock.
+        @inlinable
+        func unlock() {
+            self._storage.unlock()
+        }
+
+        /// Mutate the value, assuming the lock has been acquired manually.
+        ///
+        /// - Parameter mutate: A closure with scoped access to the value.
+        /// - Returns: The result of the `mutate` closure.
+        @inlinable
+        func withValueAssumingLockIsAcquired<Result>(
+            _ mutate: (_ value: inout Value) throws -> Result
+        ) rethrows -> Result {
+            try self._storage.withUnsafeMutablePointerToHeader { value in
+                try mutate(&value.pointee)
+            }
+        }
     }
 }
 
-extension NIOLockedValueBox: Sendable where Value: Sendable {}
+extension NIOLockedValueBox: @unchecked Sendable where Value: Sendable {}

Sources/PostgresNIO/Connection/PostgresConnection+Configuration.swift

@@ -192,9 +192,22 @@ extension PostgresConnection {
         /// - Parameters:
         ///   - channel: The `NIOCore/Channel` to use. The channel must already be active and connected to an
         ///     endpoint (i.e. `NIOCore/Channel/isActive` must be `true`).
-        ///   - tls: The TLS mode to use. Defaults to ``TLS-swift.struct/disable``.
+        ///   - tls: The TLS mode to use.
+        public init(establishedChannel channel: Channel, tls: PostgresConnection.Configuration.TLS, username: String, password: String?, database: String?) {
+            self.init(endpointInfo: .configureChannel(channel), tls: tls, username: username, password: password, database: database)
+        }
+        
+        /// Create a configuration for establishing a connection to a Postgres server over a preestablished
+        /// `NIOCore/Channel`.
+        ///
+        /// This is provided for calling code which wants to manage the underlying connection transport on its
+        /// own, such as when tunneling a connection through SSH.
+        ///
+        /// - Parameters:
+        ///   - channel: The `NIOCore/Channel` to use. The channel must already be active and connected to an
+        ///     endpoint (i.e. `NIOCore/Channel/isActive` must be `true`).
         public init(establishedChannel channel: Channel, username: String, password: String?, database: String?) {
-            self.init(endpointInfo: .configureChannel(channel), tls: .disable, username: username, password: password, database: database)
+            self.init(establishedChannel: channel, tls: .disable, username: username, password: password, database: database)
         }
 
         // MARK: - Implementation details

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -60,18 +60,18 @@ public final class PostgresConnection: @unchecked Sendable {
     func start(configuration: InternalConfiguration) -> EventLoopFuture<Void> {
         // 1. configure handlers
 
-        let configureSSLCallback: ((Channel) throws -> ())?
+        let configureSSLCallback: ((Channel, PostgresChannelHandler) throws -> ())?
 
         switch configuration.tls.base {
         case .prefer(let context), .require(let context):
-            configureSSLCallback = { channel in
+            configureSSLCallback = { channel, postgresChannelHandler in
                 channel.eventLoop.assertInEventLoop()
 
                 let sslHandler = try NIOSSLClientHandler(
                     context: context,
                     serverHostname: configuration.serverNameForTLS
                 )
-                try channel.pipeline.syncOperations.addHandler(sslHandler, position: .first)
+                try channel.pipeline.syncOperations.addHandler(sslHandler, position: .before(postgresChannelHandler))
             }
         case .disable:
             configureSSLCallback = nil
@@ -530,6 +530,110 @@ extension PostgresConnection {
             throw error // rethrow with more metadata
         }
     }
+
+    #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
+    /// 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<Result>(
+        logger: Logger,
+        file: String = #file,
+        line: Int = #line,
+        isolation: isolated (any Actor)? = #isolation,
+        // 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 {
+            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
+        }
+    }
+    #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<Result>(
+        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