From 39de6661aa0bd0ac0de0793047c95b3c1cb38267 Mon Sep 17 00:00:00 2001
From: MahdiBM <github@mahdibm.com>
Date: Mon, 26 Aug 2024 13:41:32 +0330
Subject: [PATCH] Don't modify the format of columns in a SimpleQuery

---
 .../SimpleQueryStateMachine.swift             | 13 ++----------
 Tests/IntegrationTests/AsyncTests.swift       | 21 +++++++++++++++++++
 .../SimpleQueryStateMachineTests.swift        |  4 ++--
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/Sources/PostgresNIO/New/Connection State Machine/SimpleQueryStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/SimpleQueryStateMachine.swift
index f0340679..15b4dc6b 100644
--- 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
index 9bde652b..be117027 100644
--- 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
index 0f51afd9..f98e64bc 100644
--- 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)