From 563512ad09785f050300f2e08f9afe8ff1809d9f Mon Sep 17 00:00:00 2001
From: Steve Lau <stevelauc@outlook.com>
Date: Thu, 14 Aug 2025 13:25:43 +0800
Subject: [PATCH] feat: add PostgresRowSequence.getColumns() to get column
 metadata

This commit implements `PostgresRowSequence.getColumns()` to enable users
to retrieve column metadata of their query results.

Discussion: https://github.com/vapor/postgres-nio/issues/576
---
 Sources/PostgresNIO/New/PostgresColumn.swift  | 46 +++++++++++++
 .../PostgresNIO/New/PostgresRowSequence.swift | 15 +++++
 .../New/PostgresRowSequenceTests.swift        | 65 +++++++++++++++++++
 3 files changed, 126 insertions(+)
 create mode 100644 Sources/PostgresNIO/New/PostgresColumn.swift

diff --git a/Sources/PostgresNIO/New/PostgresColumn.swift b/Sources/PostgresNIO/New/PostgresColumn.swift
new file mode 100644
index 00000000..06e5d5b6
--- /dev/null
+++ b/Sources/PostgresNIO/New/PostgresColumn.swift
@@ -0,0 +1,46 @@
+/// Information of a column.
+//
+// This type has the same definition as `RowDescription.column`, we need to keep 
+// that type private so we defines this type.
+public struct PostgresColumn: Hashable, Sendable {
+    /// The column name.
+    public let name: String
+    
+    /// If the field can be identified as a column of a specific table, the object ID of the table; otherwise zero.
+    public let tableOID: Int32
+    
+    /// If the field can be identified as a column of a specific table, the attribute number of the column; otherwise zero.
+    public let columnAttributeNumber: Int16
+    
+    /// The object ID of the field's data type.
+    public let dataType: PostgresDataType
+    
+    /// The data type size (see pg_type.typlen). Note that negative values denote variable-width types.
+    public let dataTypeSize: Int16
+    
+    /// The type modifier (see pg_attribute.atttypmod). The meaning of the modifier is type-specific.
+    public let dataTypeModifier: Int32
+    
+    /// The format being used for the field. Currently will be text or binary. In a RowDescription returned
+    /// from the statement variant of Describe, the format code is not yet known and will always be text.
+    public let format: PostgresFormat
+
+
+    internal init(
+        name: String,
+        tableOID: Int32,
+        columnAttributeNumber: Int16,
+        dataType: PostgresDataType,
+        dataTypeSize: Int16,
+        dataTypeModifier: Int32,
+        format: PostgresFormat
+    ) {
+        self.name = name
+        self.tableOID = tableOID
+        self.columnAttributeNumber = columnAttributeNumber
+        self.dataType = dataType
+        self.dataTypeSize = dataTypeSize
+        self.dataTypeModifier = dataTypeModifier
+        self.format = format
+    }
+}
\ No newline at end of file
diff --git a/Sources/PostgresNIO/New/PostgresRowSequence.swift b/Sources/PostgresNIO/New/PostgresRowSequence.swift
index 3936b51e..e1159a06 100644
--- a/Sources/PostgresNIO/New/PostgresRowSequence.swift
+++ b/Sources/PostgresNIO/New/PostgresRowSequence.swift
@@ -26,6 +26,21 @@ public struct PostgresRowSequence: AsyncSequence, Sendable {
             columns: self.columns
         )
     }
+
+    /// Get the column information of the query results.
+    public func getColumns() -> [PostgresColumn] {
+        self.columns.map { column in
+            PostgresColumn(
+                name: column.name,
+                tableOID: column.tableOID,
+                columnAttributeNumber: column.columnAttributeNumber,
+                dataType: column.dataType,
+                dataTypeSize: column.dataTypeSize,
+                dataTypeModifier: column.dataTypeModifier,
+                format: column.format
+            )
+        }
+    }
 }
 
 extension PostgresRowSequence {
diff --git a/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift b/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift
index 9d662252..35c143e6 100644
--- a/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift
+++ b/Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift
@@ -433,6 +433,71 @@ final class PostgresRowSequenceTests: XCTestCase {
         let emptyRow = try await rowIterator.next()
         XCTAssertNil(emptyRow)
     }
+
+    func testGetColumnsReturnsCorrectColumnInformation() async throws {
+        let dataSource = MockRowDataSource()
+        let embeddedEventLoop = EmbeddedEventLoop()
+        
+        let sourceColumns = [
+            RowDescription.Column(
+                name: "id",
+                tableOID: 12345,
+                columnAttributeNumber: 1,
+                dataType: .int8,
+                dataTypeSize: 8,
+                dataTypeModifier: -1,
+                format: .binary
+            ),
+            RowDescription.Column(
+                name: "name",
+                tableOID: 12345,
+                columnAttributeNumber: 2,
+                dataType: .text,
+                dataTypeSize: -1,
+                dataTypeModifier: -1,
+                format: .text
+            )
+        ]
+        
+        let expectedColumns = sourceColumns.map { column in
+            PostgresColumn(
+                name: column.name,
+                tableOID: column.tableOID,
+                columnAttributeNumber: column.columnAttributeNumber,
+                dataType: column.dataType,
+                dataTypeSize: column.dataTypeSize,
+                dataTypeModifier: column.dataTypeModifier,
+                format: column.format
+            )
+        }
+        
+        let stream = PSQLRowStream(
+            source: .stream(sourceColumns, dataSource),
+            eventLoop: embeddedEventLoop,
+            logger: self.logger
+        )
+        
+        let rowSequence = stream.asyncSequence()
+        let actualColumns = rowSequence.getColumns()
+        
+        XCTAssertEqual(actualColumns, expectedColumns)
+    }
+
+    func testGetColumnsWithEmptyColumns() async throws {
+        let dataSource = MockRowDataSource()
+        let embeddedEventLoop = EmbeddedEventLoop()
+        
+        let stream = PSQLRowStream(
+            source: .stream([], dataSource),
+            eventLoop: embeddedEventLoop,
+            logger: self.logger
+        )
+        
+        let rowSequence = stream.asyncSequence()
+        let columns = rowSequence.getColumns()
+        
+        XCTAssertTrue(columns.isEmpty)
+    }
 }
 
 final class MockRowDataSource: PSQLRowsDataSource {