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

Author: MahdiBM

Sources/PostgresNIO/New/PostgresQuery.swift

@@ -172,6 +172,16 @@ public struct PostgresBindings: Sendable, Hashable {
         try self.append(value, context: .default)
     }
 
+    @inlinable
+    public mutating func append<Value: PostgresThrowingDynamicTypeEncodable>(_ value: Optional<Value>) throws {
+        switch value {
+        case .none:
+            self.appendNull()
+        case let .some(value):
+            try self.append(value)
+        }
+    }
+
     @inlinable
     public mutating func append<Value: PostgresThrowingDynamicTypeEncodable, JSONEncoder: PostgresJSONEncoder>(
         _ value: Value,
@@ -181,11 +191,34 @@ public struct PostgresBindings: Sendable, Hashable {
         self.metadata.append(.init(value: value, protected: true))
     }
 
+    @inlinable
+    public mutating func append<Value: PostgresThrowingDynamicTypeEncodable, JSONEncoder: PostgresJSONEncoder>(
+        _ value: Optional<Value>,
+        context: PostgresEncodingContext<JSONEncoder>
+    ) throws {
+        switch value {
+        case .none:
+            self.appendNull()
+        case let .some(value):
+            try self.append(value, context: context)
+        }
+    }
+
     @inlinable
     public mutating func append<Value: PostgresDynamicTypeEncodable>(_ value: Value) {
         self.append(value, context: .default)
     }
 
+    @inlinable
+    public mutating func append<Value: PostgresDynamicTypeEncodable>(_ value: Optional<Value>) {
+        switch value {
+        case .none:
+            self.appendNull()
+        case let .some(value):
+            self.append(value)
+        }
+    }
+
     @inlinable
     public mutating func append<Value: PostgresDynamicTypeEncodable, JSONEncoder: PostgresJSONEncoder>(
         _ value: Value,
@@ -195,6 +228,19 @@ public struct PostgresBindings: Sendable, Hashable {
         self.metadata.append(.init(value: value, protected: true))
     }
 
+    @inlinable
+    public mutating func append<Value: PostgresDynamicTypeEncodable, JSONEncoder: PostgresJSONEncoder>(
+        _ value: Optional<Value>,
+        context: PostgresEncodingContext<JSONEncoder>
+    ) {
+        switch value {
+        case .none:
+            self.appendNull()
+        case let .some(value):
+            self.append(value, context: context)
+        }
+    }
+
     @inlinable
     mutating func appendUnprotected<Value: PostgresEncodable, JSONEncoder: PostgresJSONEncoder>(
         _ value: Value,

Sources/PostgresNIO/Pool/ConnectionFactory.swift

@@ -89,6 +89,7 @@ final class ConnectionFactory: Sendable {
         connectionConfig.options.connectTimeout = TimeAmount(config.options.connectTimeout)
         connectionConfig.options.tlsServerName = config.options.tlsServerName
         connectionConfig.options.requireBackendKeyData = config.options.requireBackendKeyData
+        connectionConfig.options.additionalStartupParameters = config.options.additionalStartupParameters
 
         return connectionConfig
     }

Sources/PostgresNIO/Pool/PostgresClient.swift

@@ -106,6 +106,10 @@ public final class PostgresClient: Sendable, ServiceLifecycle.Service {
             /// If you are not using Amazon RDS Proxy, you should leave this set to `true` (the default).
             public var requireBackendKeyData: Bool = true
 
+            /// Additional parameters to send to the server on startup. The name value pairs are added to the initial
+            /// startup message that the client sends to the server.
+            public var additionalStartupParameters: [(String, String)] = []
+
             /// The minimum number of connections that the client shall keep open at any time, even if there is no
             /// demand. Default to `0`.
             ///

Tests/IntegrationTests/AsyncTests.swift

@@ -525,6 +525,87 @@ final class AsyncPostgresConnectionTests: XCTestCase {
             XCTFail("Unexpected error: \(String(describing: error))")
         }
     }
+
+    static let preparedStatementWithOptionalTestTable = "AsyncTestPreparedStatementWithOptionalTestTable"
+    func testPreparedStatementWithOptionalBinding() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        struct InsertPreparedStatement: PostgresPreparedStatement {
+            static let name = "INSERT-AsyncTestPreparedStatementWithOptionalTestTable"
+
+            static let sql = #"INSERT INTO "\#(AsyncPostgresConnectionTests.preparedStatementWithOptionalTestTable)" (uuid) VALUES ($1);"#
+            typealias Row = ()
+
+            var uuid: UUID?
+
+            func makeBindings() -> PostgresBindings {
+                var bindings = PostgresBindings()
+                bindings.append(self.uuid)
+                return bindings
+            }
+
+            func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row {
+                ()
+            }
+        }
+
+        struct SelectPreparedStatement: PostgresPreparedStatement {
+            static let name = "SELECT-AsyncTestPreparedStatementWithOptionalTestTable"
+
+            static let sql = #"SELECT id, uuid FROM "\#(AsyncPostgresConnectionTests.preparedStatementWithOptionalTestTable)" WHERE id <= $1;"#
+            typealias Row = (Int, UUID?)
+
+            var id: Int
+
+            func makeBindings() -> PostgresBindings {
+                var bindings = PostgresBindings()
+                bindings.append(self.id)
+                return bindings
+            }
+
+            func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row {
+                try row.decode((Int, UUID?).self)
+            }
+        }
+
+        do {
+            try await withTestConnection(on: eventLoop) { connection in
+                try await connection.query("""
+                    CREATE TABLE IF NOT EXISTS "\(unescaped: Self.preparedStatementWithOptionalTestTable)" (
+                        id SERIAL PRIMARY KEY,
+                        uuid UUID
+                    )
+                    """,
+                    logger: .psqlTest
+                )
+
+                _ = try await connection.execute(InsertPreparedStatement(uuid: nil), logger: .psqlTest)
+                _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest)
+                _ = try await connection.execute(InsertPreparedStatement(uuid: nil), logger: .psqlTest)
+                _ = try await connection.execute(InsertPreparedStatement(uuid: .init()), logger: .psqlTest)
+                _ = try await connection.execute(InsertPreparedStatement(uuid: nil), logger: .psqlTest)
+
+                let rows = try await connection.execute(SelectPreparedStatement(id: 3), logger: .psqlTest)
+                var counter = 0
+                for try await (id, uuid) in rows {
+                    Logger.psqlTest.info("Received row", metadata: [
+                        "id": "\(id)", "uuid": "\(String(describing: uuid))"
+                    ])
+                    counter += 1
+                }
+
+                try await connection.query("""
+                    DROP TABLE "\(unescaped: Self.preparedStatementWithOptionalTestTable)";
+                    """,
+                    logger: .psqlTest
+                )
+            }
+        } catch {
+            XCTFail("Unexpected error: \(String(describing: error))")
+        }
+    }
 }
 
 extension XCTestCase {

Tests/IntegrationTests/PostgresClientTests.swift

@@ -43,6 +43,41 @@ final class PostgresClientTests: XCTestCase {
         }
     }
 
+    func testApplicationNameIsForwardedCorrectly() async throws {
+        var mlogger = Logger(label: "test")
+        mlogger.logLevel = .debug
+        let logger = mlogger
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 8)
+        self.addTeardownBlock {
+            try await eventLoopGroup.shutdownGracefully()
+        }
+
+        var clientConfig = PostgresClient.Configuration.makeTestConfiguration()
+        let applicationName = "postgres_nio_test_run"
+        clientConfig.options.additionalStartupParameters = [("application_name", applicationName)]
+        let client = PostgresClient(configuration: clientConfig, eventLoopGroup: eventLoopGroup, backgroundLogger: logger)
+
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup in
+            taskGroup.addTask {
+                await client.run()
+            }
+
+            let rows = try await client.query("select * from pg_stat_activity;");
+            var applicationNameFound = 0
+            for try await row in rows {
+                let randomAccessRow = row.makeRandomAccess()
+                if try randomAccessRow["application_name"].decode(String?.self) == applicationName {
+                    applicationNameFound += 1
+                }
+            }
+
+            XCTAssertGreaterThanOrEqual(applicationNameFound, 1)
+
+            taskGroup.cancelAll()
+        }
+    }
+
+
     func testQueryDirectly() async throws {
         var mlogger = Logger(label: "test")
         mlogger.logLevel = .debug