Feature/typed queries

Package.swift

@@ -1,5 +1,7 @@
-// swift-tools-version:5.6
+// swift-tools-version:5.9
+
 import PackageDescription
+import CompilerPluginSupport
 
 let package = Package(
     name: "postgres-nio",
@@ -20,6 +22,7 @@ let package = Package(
         .package(url: "https://github.com/apple/swift-crypto.git", "1.0.0" ..< "3.0.0"),
         .package(url: "https://github.com/apple/swift-metrics.git", from: "2.0.0"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.5.2"),
+        .package(url: "https://github.com/apple/swift-syntax.git", branch: "main")
     ],
     targets: [
         .target(
@@ -36,6 +39,14 @@ let package = Package(
                 .product(name: "NIOTLS", package: "swift-nio"),
                 .product(name: "NIOSSL", package: "swift-nio-ssl"),
                 .product(name: "NIOFoundationCompat", package: "swift-nio"),
+                .target(name: "PostgresNIOMacros")
+            ]
+        ),
+        .macro(
+            name: "PostgresNIOMacros",
+            dependencies: [
+                .product(name: "SwiftSyntaxMacros", package: "swift-syntax"),
+                .product(name: "SwiftCompilerPlugin", package: "swift-syntax")
             ]
         ),
         .testTarget(
@@ -44,6 +55,7 @@ let package = Package(
                 .target(name: "PostgresNIO"),
                 .product(name: "NIOEmbedded", package: "swift-nio"),
                 .product(name: "NIOTestUtils", package: "swift-nio"),
+                .product(name: "SwiftSyntaxMacrosTestSupport", package: "swift-syntax"),
             ]
         ),
         .testTarget(

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -428,6 +428,25 @@ extension PostgresConnection {
             throw error // rethrow with more metadata
         }
     }
+
+    /// Run a query on the Postgres server the connection is connected to.
+    ///
+    /// - Parameters:
+    ///   - query: A ``PostgresTypedQuery`` to run
+    ///   - logger: The `Logger` to log into for the query
+    ///   - file: The file, the query was started in. Used for better error reporting.
+    ///   - line: The line, the query was started in. Used for better error reporting.
+    /// - Returns: A ``PostgresTypedSequence`` containing typed rows the server sent as the query result.
+    @discardableResult
+    public func query<T: PostgresTypedQuery>(
+        _ query: T,
+        logger: Logger,
+        file: String = #fileID,
+        line: Int = #line
+    ) async throws -> PostgresTypedSequence<T.Row> {
+        let rowSequence = try await self.query(query.sql, logger: logger, file: file, line: line)
+        return PostgresTypedSequence(rowSequence: rowSequence)
+    }
 }
 
 // MARK: EventLoopFuture interface

Sources/PostgresNIO/Data/PostgresRow.swift

@@ -22,6 +22,10 @@ public struct PostgresRow: Sendable {
     }
 }
 
+public protocol PostgresTypedRow {
+    init(from row: PostgresRow) throws
+}
+
 extension PostgresRow: Equatable {
     public static func ==(lhs: Self, rhs: Self) -> Bool {
         // we don't need to compare the lookup table here, as the looup table is only derived

Sources/PostgresNIO/New/PostgresQuery.swift

@@ -13,6 +13,43 @@ public struct PostgresQuery: Sendable, Hashable {
     }
 }
 
+struct PostgresMacroQuery: ExpressibleByStringInterpolation {
+    var sql: String
+
+    public init(stringInterpolation: StringInterpolation) {
+        sql = stringInterpolation.sql
+    }
+
+    public init(stringLiteral value: String) {
+        sql = value
+    }
+
+    struct StringInterpolation: StringInterpolationProtocol {
+        typealias StringLiteralType = String
+
+        var sql: String
+
+        init(literalCapacity: Int, interpolationCount: Int) {
+            sql = ""
+        }
+
+        mutating func appendLiteral(_ literal: String) {
+            sql.append(contentsOf: literal)
+        }
+
+        mutating func appendInterpolation<T: PostgresDecodable>(_ sql: String, type: T.Type) {}
+    }
+}
+
+@Query("SELECT \("id", type: Int.self) FROM users")
+struct GetAllUsersQuery {}
+
+public protocol PostgresTypedQuery {
+    associatedtype Row: PostgresTypedRow
+
+    var sql: PostgresQuery { get }
+}
+
 extension PostgresQuery: ExpressibleByStringInterpolation {
     public init(stringInterpolation: StringInterpolation) {
         self.sql = stringInterpolation.sql

Sources/PostgresNIO/New/PostgresRowSequence.swift

@@ -43,7 +43,7 @@ extension PostgresRowSequence {
             self.columns = columns
         }
 
-        public mutating func next() async throws -> PostgresRow? {
+        public func next() async throws -> PostgresRow? {
             if let dataRow = try await self.backing.next() {
                 return PostgresRow(
                     data: dataRow,
@@ -56,6 +56,33 @@ extension PostgresRowSequence {
     }
 }
 
+public struct PostgresTypedSequence<T: PostgresTypedRow>: AsyncSequence {
+    public typealias Element = T
+
+    let rowSequence: PostgresRowSequence
+
+    init(rowSequence: PostgresRowSequence) {
+        self.rowSequence = rowSequence
+    }
+
+    public func makeAsyncIterator() -> AsyncIterator {
+        AsyncIterator(rowSequence: rowSequence.makeAsyncIterator())
+    }
+}
+
+extension PostgresTypedSequence {
+    public struct AsyncIterator: AsyncIteratorProtocol {
+        let rowSequence: PostgresRowSequence.AsyncIterator
+
+        public func next() async throws -> T? {
+            guard let row = try await self.rowSequence.next() else {
+                return nil
+            }
+            return try T.init(from: row)
+        }
+    }
+}
+
 extension PostgresRowSequence {
     public func collect() async throws -> [PostgresRow] {
         var result = [PostgresRow]()

Sources/PostgresNIO/Utilities/Macros.swift

@@ -0,0 +1,5 @@
+@attached(member)
+macro Query(_ query: PostgresMacroQuery) = #externalMacro(
+    module: "PostgresNIOMacros",
+    type: "PostgresTypedQueryMacro"
+)

Sources/PostgresNIOMacros/PostgresNIODiagnostic.swift

@@ -0,0 +1,21 @@
+import SwiftDiagnostics
+
+enum PostgresNIODiagnostic: String, DiagnosticMessage {
+    case wrongArgument
+    case notAStruct
+
+    var message: String {
+        switch self {
+        case .wrongArgument:
+            return "Invalid argument"
+        case .notAStruct:
+            return "Macro only works with structs"
+        }
+    }
+
+    var diagnosticID: SwiftDiagnostics.MessageID {
+        MessageID(domain: "PostgresNIOMacros", id: rawValue)
+    }
+
+    var severity: SwiftDiagnostics.DiagnosticSeverity { .error }
+}

Sources/PostgresNIOMacros/PostgresNIOMacro.swift

@@ -0,0 +1,76 @@
+import SwiftCompilerPlugin
+import SwiftSyntax
+import SwiftSyntaxBuilder
+import SwiftSyntaxMacros
+import SwiftDiagnostics
+
+public struct PostgresTypedQueryMacro: MemberMacro {
+    public static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        in context: some MacroExpansionContext
+    ) throws -> [DeclSyntax] {
+        guard declaration.is(StructDeclSyntax.self) else {
+            context.diagnose(Diagnostic(node: Syntax(node), message: PostgresNIODiagnostic.notAStruct))
+            return []
+        }
+
+        guard let elements = node.argument?.as(TupleExprElementListSyntax.self)?
+            .first?.as(TupleExprElementSyntax.self)?
+            .expression.as(StringLiteralExprSyntax.self)?.segments else {
+            // TODO: Be more specific about this error
+            context.diagnose(Diagnostic(node: Syntax(node), message: PostgresNIODiagnostic.wrongArgument))
+            return []
+        }
+
+
+
+        var outputTypes: [(String, String)] = []
+        for tup in elements {
+            if let expression = tup.as(ExpressionSegmentSyntax.self) {
+                outputTypes.append(extractColumnTypes(expression))
+            }
+        }
+
+        let rowStruct = try StructDeclSyntax("struct Row") {
+            for outputType in outputTypes {
+                MemberDeclListItemSyntax.init(decl: DeclSyntax(stringLiteral: "let \(outputType.0): \(outputType.1)"))
+            }
+            try InitializerDeclSyntax("init(from: PostgresRow) throws") {
+                // TODO: (id, name) = try row.decode((Int, String).self, context: .default)
+            }
+        }
+
+        return [
+//            DeclSyntax(rowStruct)
+        ]
+    }
+
+    /// Returns ("name", "String")
+    private static func extractColumnTypes(_ node: ExpressionSegmentSyntax) -> (String, String) {
+        let tupleElements = node.expressions
+        guard tupleElements.count == 2 else {
+            fatalError("Expected tuple with exactly two elements")
+        }
+
+        // First element needs to be the column name
+        var iterator = tupleElements.makeIterator()
+        guard let columnName = iterator.next()?.expression.as(StringLiteralExprSyntax.self)?
+            .segments.first?.as(StringSegmentSyntax.self)?.content
+            .text else {
+            fatalError("Expected column name")
+        }
+
+        guard let columnType = iterator.next()?.expression.as(MemberAccessExprSyntax.self)?.base?.as(IdentifierExprSyntax.self)?.identifier.text else {
+            fatalError("Expected column type")
+        }
+        return (columnName, columnType)
+    }
+}
+
+@main
+struct PostgresNIOMacros: CompilerPlugin {
+    let providingMacros: [Macro.Type] = [
+        PostgresTypedQueryMacro.self
+    ]
+}

Tests/IntegrationTests/TypedQueriesTests.swift

@@ -0,0 +1,47 @@
+import Logging
+import XCTest
+import PostgresNIO
+
+final class TypedQueriesTests: XCTestCase {
+    func testTypedPostgresQuery() async throws {
+        struct MyQuery: PostgresTypedQuery {
+            struct Row: PostgresTypedRow {
+                let id: Int
+                let name: String
+
+                init(from row: PostgresRow) throws {
+                    (id, name) = try row.decode((Int, String).self, context: .default)
+                }
+            }
+
+            var sql: PostgresQuery {
+                "SELECT id, name FROM users"
+            }
+        }
+
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        try await withTestConnection(on: eventLoop) { connection in
+            let createTableQuery = PostgresQuery(unsafeSQL: """
+            CREATE TABLE users (
+                id SERIAL PRIMARY KEY,
+                name character varying(255) NOT NULL
+            );
+            """)
+            let name = "foobar"
+
+            try await connection.query(createTableQuery, logger: .psqlTest)
+            try await connection.query("INSERT INTO users (name) VALUES (\(name));", logger: .psqlTest)
+
+            let rows = try await connection.query(MyQuery(), logger: .psqlTest)
+            for try await row in rows {
+                XCTAssertEqual(row.name, name)
+            }
+
+            let dropQuery = PostgresQuery(unsafeSQL: "DROP TABLE users")
+            try await connection.query(dropQuery, logger: .psqlTest)
+        }
+    }
+}

Tests/PostgresNIOTests/Macros/MacroTests.swift

@@ -0,0 +1,45 @@
+import SwiftSyntaxMacros
+import SwiftSyntaxMacrosTestSupport
+import XCTest
+import PostgresNIOMacros
+
+let testMacros: [String: Macro.Type] = [
+    "Query": PostgresTypedQueryMacro.self,
+]
+
+final class MacrosTests: XCTestCase {
+    func testMacro() {
+        assertMacroExpansion(
+            #"""
+            @Query("SELECT \("id", Int.self), \("name", String.self) FROM users")
+            struct MyQuery {}
+            """#,
+            expandedSource: #"""
+            struct MyQuery {
+                struct Row: PostgresTypedRow {
+                    let id: Int
+                    let name: String
+                }
+            }
+            """#,
+//            expandedSource: #"""
+//            struct MyQuery: PostgresTypedQuery {
+//                struct Row: PostgresTypedRow {
+//                    let id: Int
+//                    let name: String
+//
+//                    init(from row: PostgresRow) throws {
+//                        (id, name) = try row.decode((Int, String).self, context: .default)
+//                    }
+//                }
+//
+//                var sql: PostgresQuery {
+//                    "SELECT id, name FROM users"
+//                }
+//            }
+//            """#,
+            macros: testMacros
+        )
+    }
+}
+