Add support for Network.framework (#253)

Author: fabianfett

Package.swift

@@ -14,6 +14,7 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-nio.git", from: "2.35.0"),
+        .package(url: "https://github.com/apple/swift-nio-transport-services.git", from: "1.11.4"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.14.1"),
         .package(url: "https://github.com/apple/swift-crypto.git", "1.0.0" ..< "3.0.0"),
         .package(url: "https://github.com/apple/swift-metrics.git", from: "2.0.0"),
@@ -27,6 +28,7 @@ let package = Package(
             .product(name: "NIO", package: "swift-nio"),
             .product(name: "NIOCore", package: "swift-nio"),
             .product(name: "NIOPosix", package: "swift-nio"),
+            .product(name: "NIOTransportServices", package: "swift-nio-transport-services"),
             .product(name: "NIOTLS", package: "swift-nio"),
             .product(name: "NIOSSL", package: "swift-nio-ssl"),
             .product(name: "NIOFoundationCompat", package: "swift-nio"),

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -1,5 +1,8 @@
 import NIOCore
 import NIOConcurrencyHelpers
+#if canImport(Network)
+import NIOTransportServices
+#endif
 import NIOSSL
 import Logging
 import NIOPosix
@@ -249,12 +252,13 @@ public final class PostgresConnection {
         // thread and the EventLoop.
         return eventLoop.flatSubmit { () -> EventLoopFuture<PostgresConnection> in
             let connectFuture: EventLoopFuture<Channel>
+            let bootstrap = self.makeBootstrap(on: eventLoop, configuration: configuration)
 
             switch configuration.connection {
             case .resolved(let address, _):
-                connectFuture = ClientBootstrap(group: eventLoop).connect(to: address)
+                connectFuture = bootstrap.connect(to: address)
             case .unresolved(let host, let port):
-                connectFuture = ClientBootstrap(group: eventLoop).connect(host: host, port: port)
+                connectFuture = bootstrap.connect(host: host, port: port)
             }
 
             return connectFuture.flatMap { channel -> EventLoopFuture<PostgresConnection> in
@@ -271,6 +275,23 @@ public final class PostgresConnection {
         }
     }
 
+    static func makeBootstrap(
+        on eventLoop: EventLoop,
+        configuration: PostgresConnection.InternalConfiguration
+    ) -> NIOClientTCPBootstrapProtocol {
+        #if canImport(Network)
+        if let tsBootstrap = NIOTSConnectionBootstrap(validatingGroup: eventLoop) {
+            return tsBootstrap
+        }
+        #endif
+
+        if let nioBootstrap = ClientBootstrap(validatingGroup: eventLoop) {
+            return nioBootstrap
+        }
+
+        fatalError("No matching bootstrap found")
+    }
+
     // MARK: Query
 
     func query(_ query: PostgresQuery, logger: Logger) -> EventLoopFuture<PSQLRowStream> {

Tests/IntegrationTests/AsyncTests.swift

@@ -1,6 +1,9 @@
 import Logging
 import XCTest
 import PostgresNIO
+#if canImport(Network)
+import NIOTransportServices
+#endif
 
 #if swift(>=5.5.2)
 final class AsyncPostgresConnectionTests: XCTestCase {
@@ -41,6 +44,28 @@ final class AsyncPostgresConnectionTests: XCTestCase {
             XCTAssertEqual(counter, end + 1)
         }
     }
+
+    #if canImport(Network)
+    func testSelect10kRowsNetworkFramework() async throws {
+        let eventLoopGroup = NIOTSEventLoopGroup()
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        let start = 1
+        let end = 10000
+
+        try await withTestConnection(on: eventLoop) { connection in
+            let rows = try await connection.query("SELECT generate_series(\(start), \(end));", logger: .psqlTest)
+            var counter = 1
+            for try await element in rows.decode(Int.self, context: .default) {
+                XCTAssertEqual(element, counter)
+                counter += 1
+            }
+
+            XCTAssertEqual(counter, end + 1)
+        }
+    }
+    #endif
 }
 
 extension XCTestCase {