Fix Missing Connection Metadata

Package.swift

@@ -14,7 +14,7 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/vapor/fluent-kit.git", from: "1.52.2"),
-        .package(url: "https://github.com/vapor/postgres-kit.git", from: "2.14.1"),
+        .package(url: "https://github.com/vapor/postgres-kit.git", from: "2.17.0"),
         .package(url: "https://github.com/vapor/async-kit.git", from: "1.21.0"),
         .package(url: "https://github.com/vapor/postgres-nio.git", from: "1.33.1"),
         .package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.12.0"),

Sources/FluentPostgresDriver/FluentPostgresClientDatabase.swift

@@ -90,9 +90,12 @@ extension _FluentPostgresClientDatabase: Database {
 
     private func scoped(to conn: PostgresConnection, inTransaction: Bool? = nil) -> Self {
         .init(
-            source: .connection(conn
-                .logging(to: self.logger)
-                .sql(encodingContext: self.encodingContext, decodingContext: self.decodingContext, queryLogLevel: self.sqlLogLevel)),
+            source: .connection(conn.sql(
+                encodingContext: self.encodingContext,
+                decodingContext: self.decodingContext,
+                queryLogLevel: self.sqlLogLevel,
+                logger: self.logger.withConnectionID(of: conn)
+            )),
             context: self.context,
             encodingContext: self.encodingContext,
             decodingContext: self.decodingContext,

Sources/FluentPostgresDriver/FluentPostgresDatabase.swift

@@ -38,10 +38,10 @@ extension _FluentPostgresDatabase: Database {
     }
 
     func withConnection<T>(_ closure: @escaping @Sendable (any Database) -> EventLoopFuture<T>) -> EventLoopFuture<T> {
-        self.withConnection { (underlying: any PostgresDatabase) in
+        self.withConnection { (connection: PostgresConnection) in
             closure(
                 _FluentPostgresDatabase(
-                    database: underlying.logging(to: self.logger).sql(
+                    database: connection.logging(to: self.logger.withConnectionID(of: connection)).sql(
                         encodingContext: self.encodingContext,
                         decodingContext: self.decodingContext,
                         queryLogLevel: self.database.queryLogLevel
@@ -106,3 +106,15 @@ extension _FluentPostgresDatabase: PostgresDatabase {
         return psqlDb.withConnection(closure)
     }
 }
+
+extension Logger {
+    /// This logger, with `connection`'s ID attached the same way PostgresNIO tags its own logs.
+    ///
+    /// Only the ID is taken from the connection. The rest of the connection's own logger metadata (such as a request ID)
+    /// belongs to whichever request happened to open the connection, so it must not be used for anyone else's queries.
+    func withConnectionID(of connection: PostgresConnection) -> Logger {
+        var logger = self
+        logger[metadataKey: "psql_connection_id"] = "\(connection.id)"
+        return logger
+    }
+}

Tests/FluentPostgresDriverTests/FluentPostgresDriverTests.swift

@@ -226,6 +226,20 @@ struct FluentPostgresDriverTests {
             try await db.schema(Seq.schema).delete()
         }
     }
+
+    /// Only the AsyncKit driver's databases are castable to `PostgresDatabase`, and that must hold once they're scoped to a connection.
+    @Test
+    func connectionScopedDatabaseIsStillPostgresDatabase() async throws {
+        try await withDbs(.asyncKit) { _, db in
+            try await db.transaction { db in
+                #expect(try await (db as! any PostgresDatabase).simpleQuery("SELECT 1").get().count == 1)
+
+                try await db.withConnection { db in
+                    try await (db as! any SQLDatabase).raw("SELECT 1").run()
+                }
+            }
+        }
+    }
 }
 }
 

Tests/FluentPostgresDriverTests/FluentPostgresLoggingTests.swift

@@ -39,13 +39,37 @@ extension AllSuites {
             }
         }
 
+        @Test(.bug("https://github.com/vapor/fluent-postgres-driver/issues/230"), arguments: TestDriver.allCases)
+        func transactionLogsIncludeConnectionID(_ driver: TestDriver) async throws {
+            try await withDbs(driver) { dbs, _ in
+                let (_, second) = try await Self.warmedUpDatabases(dbs)
+
+                try await second.db.transaction { db in
+                    try await (db as! any SQLDatabase).raw("SELECT 'inside transaction'").run()
+                }
+
+                let metadata = ["BEGIN", "SELECT 'inside transaction'", "COMMIT"].map { second.logs.metadata(forQuery: $0) }
+
+                // Every query in the transaction ran on the same connection
+                let connectionIDs = metadata.map { $0?["psql_connection_id"] }
+                #expect(connectionIDs[0] != nil)
+                #expect(connectionIDs.allSatisfy { $0 == connectionIDs[0] })
+
+                // It should have it's own context though
+                #expect(metadata.allSatisfy { $0?["request-id"] == "second" })
+            }
+        }
+
         private static func warmedUpDatabases(
             _ dbs: Databases
         ) async throws -> (first: (db: any Database, logs: InMemoryLogHandler), second: (db: any Database, logs: InMemoryLogHandler)) {
             let eventLoop = dbs.eventLoopGroup.any()
             let firstLogs = InMemoryLogHandler.capturingAllLevels(), secondLogs = InMemoryLogHandler.capturingAllLevels()
-            let first = dbs.database(.a, logger: Logger(label: "first") { _ in firstLogs }, on: eventLoop)!
-            let second = dbs.database(.a, logger: Logger(label: "second") { _ in secondLogs }, on: eventLoop)!
+            var firstLogger = Logger(label: "first") { _ in firstLogs }, secondLogger = Logger(label: "second") { _ in secondLogs }
+            firstLogger[metadataKey: "request-id"] = "first"
+            secondLogger[metadataKey: "request-id"] = "second"
+            let first = dbs.database(.a, logger: firstLogger, on: eventLoop)!
+            let second = dbs.database(.a, logger: secondLogger, on: eventLoop)!
 
             try await first.withConnection { db in
                 try await (db as! any SQLDatabase).raw("SELECT 'warm up'").run()
@@ -70,4 +94,9 @@ extension InMemoryLogHandler {
             return sql
         }
     }
+
+    /// The metadata that was logged alongside the given query, if it was logged to this handler.
+    func metadata(forQuery sql: String) -> Logger.Metadata? {
+        self.entries.first { $0.metadata["sql"] == .string(sql) }?.metadata
+    }
 }