Don't modify the format of columns in a SimpleQuery

vapor · Aug 26, 2024 · 39de666 · 39de666
1 parent ffa0dc6
commit 39de666
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 13 deletions.
diff --git a/Sources/PostgresNIO/New/Connection State Machine/SimpleQueryStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/SimpleQueryStateMachine.swift
@@ -108,23 +108,14 @@ struct SimpleQueryStateMachine {
             return self.setAndFireError(.unexpectedBackendMessage(.rowDescription(rowDescription)))
         }
 
-        // In Postgres extended queries we always request the response rows to be returned in
-        // `.binary` format.
-        // However, this is a simple query and almost all responses will be in text format anyway.
-        let columns = rowDescription.columns.map { column -> RowDescription.Column in
-            var column = column
-            // FIXME: .binary is not valid in a simple-query
-            column.format = .binary
-            return column
-        }
-
         guard !self.isCancelled else {
             self.state = .drain(rowDescription.columns)
             return .failQuery(queryContext.promise, with: .queryCancelled)
         }
 
         self.avoidingStateMachineCoW { state in
-            state = .rowDescriptionReceived(queryContext, columns)
+            // In a simple query almost all responses/columns will be in text format.
+            state = .rowDescriptionReceived(queryContext, rowDescription.columns)
         }
 
         return .wait

diff --git a/Tests/IntegrationTests/AsyncTests.swift b/Tests/IntegrationTests/AsyncTests.swift
@@ -79,6 +79,27 @@ final class AsyncPostgresConnectionTests: XCTestCase {
         }
     }
 
+    func testSelect10kRowsSimpleQuery() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        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.__simpleQuery("SELECT generate_series(\(start), \(end));", logger: .psqlTest)
+            var counter = 0
+            for try await row in rows {
+                let element = try row.decode(Int.self)
+                XCTAssertEqual(element, counter + 1)
+                counter += 1
+            }
+
+            XCTAssertEqual(counter, end)
+        }
+    }
+
     func testSelectActiveConnection() async throws {
         let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }

diff --git a/Tests/PostgresNIOTests/New/Connection State Machine/SimpleQueryStateMachineTests.swift b/Tests/PostgresNIOTests/New/Connection State Machine/SimpleQueryStateMachineTests.swift
@@ -59,7 +59,7 @@ class SimpleQueryStateMachineTests: XCTestCase {
         ]
         let expected: [RowDescription.Column] = input.map {
             .init(name: $0.name, tableOID: $0.tableOID, columnAttributeNumber: $0.columnAttributeNumber, dataType: $0.dataType,
-                  dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .binary)
+                  dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .text)
         }
 
         XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
@@ -202,7 +202,7 @@ class SimpleQueryStateMachineTests: XCTestCase {
         ]
         let expected: [RowDescription.Column] = input.map {
             .init(name: $0.name, tableOID: $0.tableOID, columnAttributeNumber: $0.columnAttributeNumber, dataType: $0.dataType,
-                  dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .binary)
+                  dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .text)
         }
 
         XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)