From f2a6394a2e7157d547727b975fc0328b92f89fb1 Mon Sep 17 00:00:00 2001
From: Fabian Fett <fabianfett@apple.com>
Date: Mon, 21 Oct 2024 10:57:36 +0200
Subject: [PATCH] Support `additionalStartupParameters` in PostgresClient
 (#521)

---
 .../PostgresNIO/Pool/ConnectionFactory.swift  |  1 +
 Sources/PostgresNIO/Pool/PostgresClient.swift |  4 +++
 .../PostgresClientTests.swift                 | 35 +++++++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/Sources/PostgresNIO/Pool/ConnectionFactory.swift b/Sources/PostgresNIO/Pool/ConnectionFactory.swift
index 77a0c047..319b86c4 100644
--- a/Sources/PostgresNIO/Pool/ConnectionFactory.swift
+++ b/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
     }
diff --git a/Sources/PostgresNIO/Pool/PostgresClient.swift b/Sources/PostgresNIO/Pool/PostgresClient.swift
index 0907f1f8..ad8a4bf1 100644
--- a/Sources/PostgresNIO/Pool/PostgresClient.swift
+++ b/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`.
             ///
diff --git a/Tests/IntegrationTests/PostgresClientTests.swift b/Tests/IntegrationTests/PostgresClientTests.swift
index d6d89dc3..579c92cd 100644
--- a/Tests/IntegrationTests/PostgresClientTests.swift
+++ b/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