Skip to content

feat: add PostgresRowSequence.getColumns() to get column metadata #577

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions Sources/PostgresNIO/New/PostgresColumn.swift
Original file line number Diff line number Diff line change
@@ -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
}
}
15 changes: 15 additions & 0 deletions Sources/PostgresNIO/New/PostgresRowSequence.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
65 changes: 65 additions & 0 deletions Tests/PostgresNIOTests/New/PostgresRowSequenceTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down