add postgres query metadata (#95)

Author: tanner0101

Sources/PostgresNIO/Message/PostgresMessage+CommandComplete.swift

@@ -2,7 +2,7 @@ import NIO
 
 extension PostgresMessage {
     /// Identifies the message as a Close command.
-    public struct CommandComplete {
+    public struct CommandComplete: PostgresMessageType {
         /// Parses an instance of this message type from a byte buffer.
         public static func parse(from buffer: inout ByteBuffer) throws -> CommandComplete {
             guard let string = buffer.readNullTerminatedString() else {

Sources/PostgresNIO/PostgresDatabase+Query.swift

@@ -2,43 +2,128 @@ import NIO
 import Logging
 
 extension PostgresDatabase {
-    public func query(_ string: String, _ binds: [PostgresData] = []) -> EventLoopFuture<[PostgresRow]> {
+    public func query(
+        _ string: String,
+        _ binds: [PostgresData] = []
+    ) -> EventLoopFuture<PostgresQueryResult> {
         var rows: [PostgresRow] = []
-        return query(string, binds) { rows.append($0) }.map { rows }
+        var metadata: PostgresQueryMetadata?
+        return self.query(string, binds, onMetadata: {
+            metadata = $0
+        }) {
+            rows.append($0)
+        }.map {
+            .init(metadata: metadata!, rows: rows)
+        }
     }
-    
-    public func query(_ string: String, _ binds: [PostgresData] = [], _ onRow: @escaping (PostgresRow) throws -> ()) -> EventLoopFuture<Void> {
-        let query = PostgresParameterizedQuery(query: string, binds: binds, onRow: onRow)
+
+    public func query(
+        _ string: String,
+        _ binds: [PostgresData] = [],
+        onMetadata: @escaping (PostgresQueryMetadata) -> () = { _ in },
+        onRow: @escaping (PostgresRow) throws -> ()
+    ) -> EventLoopFuture<Void> {
+        let query = PostgresParameterizedQuery(
+            query: string,
+            binds: binds,
+            onMetadata: onMetadata,
+            onRow: onRow
+        )
         return self.send(query, logger: self.logger)
     }
 }
 
+public struct PostgresQueryResult {
+    public let metadata: PostgresQueryMetadata
+    public let rows: [PostgresRow]
+}
+
+extension PostgresQueryResult: Collection {
+    public typealias Index = Int
+    public typealias Element = PostgresRow
+
+    public var startIndex: Int {
+        self.rows.startIndex
+    }
+
+    public var endIndex: Int {
+        self.rows.endIndex
+    }
+
+    public subscript(position: Int) -> PostgresRow {
+        self.rows[position]
+    }
+
+    public func index(after i: Int) -> Int {
+        self.rows.index(after: i)
+    }
+}
+
+public struct PostgresQueryMetadata {
+    public let command: String
+    public var oid: Int?
+    public var rows: Int?
+
+    init?(string: String) {
+        let parts = string.split(separator: " ")
+        guard parts.count >= 1 else {
+            return nil
+        }
+        switch parts[0] {
+        case "INSERT":
+            // INSERT oid rows
+            guard parts.count == 3 else {
+                return nil
+            }
+            self.command = .init(parts[0])
+            self.oid = Int(parts[1])
+            self.rows = Int(parts[2])
+        case "DELETE", "UPDATE", "SELECT", "MOVE", "FETCH", "COPY":
+            // <cmd> rows
+            guard parts.count == 2 else {
+                return nil
+            }
+            self.command = .init(parts[0])
+            self.oid = nil
+            self.rows = Int(parts[1])
+        default:
+            // <cmd>
+            self.command = string
+            self.oid = nil
+            self.rows = nil
+        }
+    }
+}
+
 // MARK: Private
 
 private final class PostgresParameterizedQuery: PostgresRequest {
     let query: String
     let binds: [PostgresData]
+    var onMetadata: (PostgresQueryMetadata) -> ()
     var onRow: (PostgresRow) throws -> ()
     var rowLookupTable: PostgresRow.LookupTable?
     var resultFormatCodes: [PostgresFormatCode]
     var logger: Logger?
-    
+
     init(
         query: String,
         binds: [PostgresData],
+        onMetadata: @escaping (PostgresQueryMetadata) -> (),
         onRow: @escaping (PostgresRow) throws -> ()
     ) {
         self.query = query
         self.binds = binds
+        self.onMetadata = onMetadata
         self.onRow = onRow
         self.resultFormatCodes = [.binary]
     }
-    
+
     func log(to logger: Logger) {
         self.logger = logger
         logger.debug("\(self.query) \(self.binds)")
     }
-    
+
     func respond(to message: PostgresMessage) throws -> [PostgresMessage]? {
         if case .error = message.identifier {
             // we should continue after errors
@@ -77,6 +162,11 @@ private final class PostgresParameterizedQuery: PostgresRequest {
             }
             return []
         case .commandComplete:
+            let complete = try PostgresMessage.CommandComplete(message: message)
+            guard let metadata = PostgresQueryMetadata(string: complete.tag) else {
+                throw PostgresError.protocol("Unexpected query metadata: \(complete.tag)")
+            }
+            self.onMetadata(metadata)
             return []
         case .notice:
             return []
@@ -87,7 +177,7 @@ private final class PostgresParameterizedQuery: PostgresRequest {
         default: throw PostgresError.protocol("Unexpected message during query: \(message)")
         }
     }
-    
+
     func start() throws -> [PostgresMessage] {
         let parse = PostgresMessage.Parse(
             statementName: "",
@@ -109,7 +199,7 @@ private final class PostgresParameterizedQuery: PostgresRequest {
             portalName: "",
             maxRows: 0
         )
-        
+
         let sync = PostgresMessage.Sync()
         return try [parse.message(), describe.message(), bind.message(), execute.message(), sync.message()]
     }

Tests/PostgresNIOTests/PostgresNIOTests.swift

@@ -855,6 +855,23 @@ final class PostgresNIOTests: XCTestCase {
         let res = try conn.query("SELECT $1::text as foo", [String?.none.postgresData!]).wait()
         XCTAssertEqual(res[0].column("foo")?.string, nil)
     }
+
+    func testUpdateMetadata() throws {
+        let conn = try PostgresConnection.test(on: eventLoop).wait()
+        defer { try! conn.close().wait() }
+        _ = try conn.simpleQuery("DROP TABLE IF EXISTS test_table").wait()
+        _ = try conn.simpleQuery("CREATE TABLE test_table(pk int PRIMARY KEY)").wait()
+        _ = try conn.simpleQuery("INSERT INTO test_table VALUES(1)").wait()
+        try conn.query("DELETE FROM test_table", onMetadata: { metadata in
+            XCTAssertEqual(metadata.command, "DELETE")
+            XCTAssertEqual(metadata.oid, nil)
+            XCTAssertEqual(metadata.rows, 1)
+        }, onRow: { _ in }).wait()
+        let rows = try conn.query("DELETE FROM test_table").wait()
+        XCTAssertEqual(rows.metadata.command, "DELETE")
+        XCTAssertEqual(rows.metadata.oid, nil)
+        XCTAssertEqual(rows.metadata.rows, 0)
+    }
 }
 
 let isLoggingConfigured: Bool = {