async/await prepared statements

mariosangiorgio · mariosangiorgio · commit d2d34600cc52 · 2023-08-10T14:57:53.000+02:00
diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/Sources/PostgresNIO/Connection/PostgresConnection.swift
@@ -460,6 +460,27 @@ extension PostgresConnection {
             self.channel.write(task, promise: nil)
         }
     }
+
+    /// Execute a prepared statement, taking care of the preparation when necessary
+    public func execute<P: PreparedStatement>(
+        _ preparedStatement: P,
+        logger: Logger
+    ) async throws -> AsyncThrowingMapSequence<PostgresRowSequence, P.Row>
+    {
+        let promise = self.channel.eventLoop.makePromise(of: PSQLRowStream.self)
+        let task = HandlerTask.executePreparedStatement(.init(
+            name: String(reflecting: P.self),
+            sql: P.sql,
+            bindings: preparedStatement.makeBindings(),
+            logger: logger,
+            promise: promise
+        ))
+        self.channel.write(task, promise: nil)
+        return try await promise.futureResult
+            .map { $0.asyncSequence() }
+            .get()
+            .map { try preparedStatement.decodeRow($0) }
+    }
 }
 
 // MARK: EventLoopFuture interface
diff --git a/Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift b/Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift
@@ -0,0 +1,68 @@
+import NIOCore
+
+struct PreparedStatementStateMachine {
+    enum State {
+        case preparing([PreparedStatementContext])
+        case prepared(RowDescription?)
+        case error(PSQLError)
+    }
+    
+    enum Action {
+        case prepareStatement
+        case waitForAlreadyInFlightPreparation
+        case executePendingStatements([PreparedStatementContext], RowDescription?)
+        case returnError([PreparedStatementContext], PSQLError)
+    }
+    
+    var preparedStatements: [String: State]
+    
+    init() {
+        self.preparedStatements = [:]
+    }
+    
+    mutating func lookup(name: String, context: PreparedStatementContext) -> Action {
+        if let state = self.preparedStatements[name] {
+            switch state {
+            case .preparing(var statements):
+                statements.append(context)
+                self.preparedStatements[name] = .preparing(statements)
+                return .waitForAlreadyInFlightPreparation
+            case .prepared(let rowDescription):
+                return .executePendingStatements([context], rowDescription)
+            case .error(let error):
+                return .returnError([context], error)
+            }
+        } else {
+            self.preparedStatements[name] = .preparing([context])
+            return .prepareStatement
+        }
+    }
+    
+    mutating func preparationComplete(
+        name: String,
+        rowDescription: RowDescription?
+    ) -> Action {
+        guard case .preparing(let statements) = self.preparedStatements[name] else {
+            preconditionFailure("Preparation completed for an unexpected statement")
+        }
+        // When sending the bindings we are going to ask for binary data.
+        if var rowDescription {
+            for i in 0..<rowDescription.columns.count {
+                rowDescription.columns[i].format = .binary
+            }
+            self.preparedStatements[name] = .prepared(rowDescription)
+            return .executePendingStatements(statements, rowDescription)
+        } else {
+            self.preparedStatements[name] = .prepared(nil)
+            return .executePendingStatements(statements, nil)
+        }
+    }
+
+    mutating func errorHappened(name: String, error: PSQLError) -> Action {
+        guard case .preparing(let statements) = self.preparedStatements[name] else {
+            preconditionFailure("Preparation completed for an unexpected statement")
+        }
+        self.preparedStatements[name] = .error(error)
+        return .returnError(statements, error)
+    }
+}
diff --git a/Sources/PostgresNIO/New/PSQLTask.swift b/Sources/PostgresNIO/New/PSQLTask.swift
@@ -6,6 +6,7 @@ enum HandlerTask {
     case closeCommand(CloseCommandContext)
     case startListening(NotificationListener)
     case cancelListening(String, Int)
+    case executePreparedStatement(PreparedStatementContext)
 }
 
 enum PSQLTask {
@@ -69,6 +70,28 @@ final class ExtendedQueryContext {
     }
 }
 
+final class PreparedStatementContext{
+    let name: String
+    let sql: String
+    let bindings: PostgresBindings
+    let logger: Logger
+    let promise: EventLoopPromise<PSQLRowStream>
+
+    init(
+        name: String,
+        sql: String,
+        bindings: PostgresBindings,
+        logger: Logger,
+        promise: EventLoopPromise<PSQLRowStream>
+    ) {
+        self.name = name
+        self.sql = sql
+        self.bindings = bindings
+        self.logger = logger
+        self.promise = promise
+    }
+}
+
 final class CloseCommandContext {
     let target: CloseTarget
     let logger: Logger
diff --git a/Sources/PostgresNIO/New/PostgresChannelHandler.swift b/Sources/PostgresNIO/New/PostgresChannelHandler.swift
@@ -23,6 +23,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
     private let configureSSLCallback: ((Channel) throws -> Void)?
 
     private var listenState: ListenStateMachine
+    private var preparedStatementState: PreparedStatementStateMachine
 
     init(
         configuration: PostgresConnection.InternalConfiguration,
@@ -33,6 +34,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
         self.state = ConnectionStateMachine(requireBackendKeyData: configuration.options.requireBackendKeyData)
         self.eventLoop = eventLoop
         self.listenState = ListenStateMachine()
+        self.preparedStatementState = PreparedStatementStateMachine()
         self.configuration = configuration
         self.configureSSLCallback = configureSSLCallback
         self.logger = logger
@@ -51,6 +53,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
         self.state = state
         self.eventLoop = eventLoop
         self.listenState = ListenStateMachine()
+        self.preparedStatementState = PreparedStatementStateMachine()
         self.configuration = configuration
         self.configureSSLCallback = configureSSLCallback
         self.logger = logger
@@ -233,6 +236,56 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
                 listener.failed(CancellationError())
                 return
             }
+        case .executePreparedStatement(let preparedStatement):
+            switch self.preparedStatementState.lookup(
+                name: preparedStatement.name,
+                context: preparedStatement
+            ) {
+            case .prepareStatement:
+                let promise = self.eventLoop.makePromise(of: RowDescription?.self)
+                promise.futureResult.whenSuccess { rowDescription in
+                    self.prepareStatementComplete(
+                        name: preparedStatement.name,
+                        rowDescription: rowDescription,
+                        context: context
+                    )
+                }
+                promise.futureResult.whenFailure { error in
+                    self.prepareStatementFailed(
+                        name: preparedStatement.name,
+                        error: error as! PSQLError,
+                        context: context
+                    )
+                }
+                psqlTask = .extendedQuery(.init(
+                    name: preparedStatement.name,
+                    query: preparedStatement.sql,
+                    logger: preparedStatement.logger,
+                    promise: promise
+                ))
+            case .waitForAlreadyInFlightPreparation:
+                // The state machine already keeps track of this
+                // and will execute the statement as soon as it's prepared
+                return
+            case .executePendingStatements(let pendingStatements, let rowDescription):
+                for statement in pendingStatements {
+                    let action = self.state.enqueue(task: .extendedQuery(.init(
+                        executeStatement: .init(
+                            name: statement.name,
+                            binds: statement.bindings,
+                            rowDescription: rowDescription),
+                        logger: statement.logger,
+                        promise: statement.promise
+                    )))
+                    self.run(action, with: context)
+                }
+                return
+            case .returnError(let pendingStatements, let error):
+                for statement in pendingStatements {
+                    statement.promise.fail(error)
+                }
+                return
+            }
         }
 
         let action = self.state.enqueue(task: psqlTask)
@@ -664,6 +717,49 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
         }
     }
 
+    private func prepareStatementComplete(
+        name: String,
+        rowDescription: RowDescription?,
+        context: ChannelHandlerContext
+    ) {
+        let action = self.preparedStatementState.preparationComplete(
+            name: name,
+            rowDescription: rowDescription
+        )
+        guard case .executePendingStatements(let statements, let rowDescription) = action else {
+            preconditionFailure("Expected to have pending statements to execute")
+        }
+        for preparedStatement in statements {
+            let action = self.state.enqueue(task: .extendedQuery(.init(
+                executeStatement: .init(
+                    name: preparedStatement.name,
+                    binds: preparedStatement.bindings,
+                    rowDescription: rowDescription
+                ),
+                logger: preparedStatement.logger,
+                promise: preparedStatement.promise
+            ))
+            )
+            self.run(action, with: context)
+        }
+    }
+
+    private func prepareStatementFailed(
+        name: String,
+        error: PSQLError,
+        context: ChannelHandlerContext
+    ) {
+        let action = self.preparedStatementState.errorHappened(
+            name: name,
+            error: error
+        )
+        guard case .returnError(let statements, let error) = action else {
+            preconditionFailure("Expected to have pending statements to execute")
+        }
+        for statement in statements {
+            statement.promise.fail(error)
+        }
+    }
 }
 
 extension PostgresChannelHandler: PSQLRowsDataSource {
diff --git a/Sources/PostgresNIO/New/PreparedStatement.swift b/Sources/PostgresNIO/New/PreparedStatement.swift
@@ -0,0 +1,40 @@
+/// A prepared statement.
+///
+/// Structs conforming to this protocol will need to provide the SQL statement to
+/// send to the server and a way of creating bindings are decoding the result.
+///
+/// As an example, consider this struct:
+/// ```swift
+/// struct Example: PreparedStatement {
+///     static var sql = "SELECT pid, datname FROM pg_stat_activity WHERE state = $1"
+///     typealias Row = (Int, String)
+///
+///     var state: String
+///
+///     func makeBindings() -> PostgresBindings {
+///         var bindings = PostgresBindings()
+///         bindings.append(.init(string: self.state))
+///         return bindings
+///     }
+///
+///     func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row {
+///         try row.decode(Row.self)
+///     }
+/// }
+/// ```
+///
+/// Structs conforming to this protocol can then be used with `PostgresConnection.execute(_ preparedStatement:, logger:)`,
+/// which will take care of preparing the statement on the server side and executing it.
+public protocol PreparedStatement {
+    /// The type rows returned by the statement will be decoded into
+    associatedtype Row
+
+    /// The SQL statement to prepare on the database server.
+    static var sql: String { get }
+
+    /// Make the bindings to provided concrete values to use when executing the prepared SQL statement
+    func makeBindings() -> PostgresBindings
+    
+    /// Decode a row returned by the database into an instance of `Row`
+    func decodeRow(_ row: PostgresRow) throws -> Row
+}
diff --git a/Tests/IntegrationTests/AsyncTests.swift b/Tests/IntegrationTests/AsyncTests.swift
@@ -315,6 +315,48 @@ final class AsyncPostgresConnectionTests: XCTestCase {
             try await connection.query("SELECT 1;", logger: .psqlTest)
         }
     }
+
+    func testPreparedStatement() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        struct TestPreparedStatement: PreparedStatement {
+            static var sql = "SELECT pid, datname FROM pg_stat_activity WHERE state = $1"
+            typealias Row = (Int, String)
+
+            var state: String
+
+            func makeBindings() -> PostgresBindings {
+                var bindings = PostgresBindings()
+                bindings.append(.init(string: self.state))
+                return bindings
+            }
+
+            func decodeRow(_ row: PostgresNIO.PostgresRow) throws -> Row {
+                try row.decode(Row.self)
+            }
+        }
+        let preparedStatement = TestPreparedStatement(state: "active")
+        try await withTestConnection(on: eventLoop) { connection in
+            var results = try await connection.execute(preparedStatement, logger: .psqlTest)
+            var counter = 0
+
+            for try await element in results {
+                XCTAssertEqual(element.1, env("POSTGRES_DB") ?? "test_database")
+                counter += 1
+            }
+
+            XCTAssertGreaterThanOrEqual(counter, 1)
+
+            // Second execution, which reuses the existing prepared statement
+            results = try await connection.execute(preparedStatement, logger: .psqlTest)
+            for try await element in results {
+                XCTAssertEqual(element.1, env("POSTGRES_DB") ?? "test_database")
+                counter += 1
+            }
+        }
+    }
 }
 
 extension XCTestCase {
