async/await prepared statement API #390

mariosangiorgio · 2023-08-09T10:13:14Z

Implementation of async/await API for prepared statements (See also #385).

This PR adds a new PreparedStatement protocol to represent prepared statements and an execute function on PostgresConnection to prepare and execute statements.

To implement the features the PR also introduces a PreparedStatementStateMachine that keeps track of the state of a prepared statement at the connection level. This ensures that, for each connection, each statement is prepared once at time of first use and then subsequent uses are going to skip the preparation step and just execute it.

Example usage

First define the struct to represent the prepared statement:

 struct ExamplePreparedStatement: 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)
    }
}

then, assuming you already have a PostgresConnection you can execute it:

let preparedStatement = ExamplePreparedStatement(state: "active")
let results = try await connection.execute(preparedStatement, logger: logger)
for (pid, database) in results {
    print("PID: \(pid), database: \(database)")
}

mariosangiorgio · 2023-08-09T10:14:37Z

Sources/PostgresNIO/Connection/PostgresConnection.swift

+    public func execute<P: PreparedStatement>(
+        _ preparedStatement: P,
+        logger: Logger
+    ) async throws -> AsyncThrowingMapSequence<PostgresRowSequence, P.Row>


This return type isn't exactly nice. Any advice on how I could improve it?

some AsyncSequence where Element == P.Row?

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

mariosangiorgio · 2023-08-09T10:17:40Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

+                        context: context
+                    )
+                }
+                promise.futureResult.whenFailure { error in


With this I want to propagate failures happening during preparation to all the pending queries we have.
I am not 100% convinced by the way it looks (especially by the need to do error as! PSQLError) but I cannot find any better alternative.

Do you have any advice?

First I think we should use whenComplete here instead of whenSuccess and whenFailure. Reason for this is, we can save an allocation...

Second, we must check if the error we get is actually a PSQLError, if it is not, we need to wrap error with a PSQLError.

mariosangiorgio · 2023-08-09T10:19:16Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

+                        logger: statement.logger,
+                        promise: statement.promise
+                    )))
+                    self.run(action, with: context)


Here (and in a couple other places) I need to execute all the pending statement executions.
To do that I enqueue and run many actions at the same time. Is this okay?

fabianfett

First review pass

fabianfett · 2023-08-09T12:48:22Z

Sources/PostgresNIO/Connection/PostgresConnection.swift

+    public func execute<P: PreparedStatement>(
+        _ preparedStatement: P,
+        logger: Logger
+    ) async throws -> AsyncThrowingMapSequence<PostgresRowSequence, P.Row>


some AsyncSequence where Element == P.Row?

fabianfett · 2023-08-09T12:48:45Z

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -448,6 +448,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>(


Lets rename P to Statement

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

fabianfett · 2023-08-09T12:51:13Z

Sources/PostgresNIO/Connection/PostgresConnection.swift

+    public func execute<P: PreparedStatement>(
+        _ preparedStatement: P,
+        logger: Logger
+    ) async throws -> AsyncThrowingMapSequence<PostgresRowSequence, P.Row>


also let's overload this method with a method where Statement.Row == () where we only return the command tag?

fabianfett · 2023-08-09T12:52:16Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

+            switch self.preparedStatementState.lookup(
+                name: preparedStatement.name,
+                context: preparedStatement
+            ) {


can we assign a variable here and then switch over the variable?

fabianfett · 2023-08-09T12:54:01Z

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

+    enum Action {
+        case prepareStatement
+        case waitForAlreadyInFlightPreparation
+        case executePendingStatements([PreparedStatementContext], RowDescription?)
+        case returnError([PreparedStatementContext], PSQLError)
+    }


returning the same Action for all mutations is what I did in the past, but what we generally frown upon by now. Can we get specific Actions for each mutation?

fabianfett · 2023-08-09T12:54:52Z

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

+                self.preparedStatements[name] = .preparing(statements)
+                return .waitForAlreadyInFlightPreparation
+            case .prepared(let rowDescription):
+                return .executePendingStatements([context], rowDescription)


we shouldn't allocate here :P

fabianfett · 2023-08-09T12:55:53Z

Sources/PostgresNIO/New/PreparedStatement.swift

+///
+/// 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 {


we should require that this is Sendable :)

Also we should prefix this one with Postgres. Reason: https://github.com/apple/swift-nio/blob/main/docs/public-api.md#1-no-global-namespace-additions-that-arent-prefixed

fabianfett · 2023-08-09T14:05:48Z

Fixes #385

codecov-commenter · 2023-08-10T14:02:50Z

Codecov Report

Merging #390 (da87ce0) into main (5217ba7) will increase coverage by 1.44%.
The diff coverage is 83.33%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #390      +/-   ##
==========================================
+ Coverage   47.53%   48.98%   +1.44%     
==========================================
  Files         109      110       +1     
  Lines        8764     8954     +190     
==========================================
+ Hits         4166     4386     +220     
+ Misses       4598     4568      -30

Files Changed	Coverage Δ
Sources/PostgresNIO/New/PostgresQuery.swift	`81.29% <0.00%> (+2.77%)`	⬆️
...es/PostgresNIO/Connection/PostgresConnection.swift	`42.35% <52.08%> (+1.23%)`	⬆️
... State Machine/PreparedStatementStateMachine.swift	`91.83% <91.83%> (ø)`
...urces/PostgresNIO/New/PostgresChannelHandler.swift	`84.45% <98.82%> (+6.23%)`	⬆️
Sources/PostgresNIO/New/PSQLTask.swift	`78.12% <100.00%> (+16.58%)`	⬆️

... and 6 files with indirect coverage changes

fabianfett

Looks better and better. For unit test examples have a look at PostgresConnectionTests.testSimpleListen

Sources/PostgresNIO/Connection/PostgresConnection.swift

fabianfett · 2023-08-11T06:59:43Z

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

+        case prepareStatement
+        case waitForAlreadyInFlightPreparation
+        case executeStatement(RowDescription?)
+        case executePendingStatements([PreparedStatementContext], RowDescription?)


we can drop this one, as it isn't used, can't we?

fabianfett · 2023-08-11T07:00:10Z

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

+    enum PreparationCompleteAction {
+        case executePendingStatements([PreparedStatementContext], RowDescription?)
+    }
+
+    mutating func preparationComplete(
+        name: String,
+        rowDescription: RowDescription?
+    ) -> PreparationCompleteAction {


would this work better as a simple struct?

fabianfett · 2023-08-11T07:01:44Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

        self.listenState = ListenStateMachine()
+        self.preparedStatementState = PreparedStatementStateMachine()


I wonder if we should just init those two directly in the declaration...

fabianfett · 2023-08-11T07:13:28Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

+                        context: context
+                    )
+                }
+                promise.futureResult.whenFailure { error in


First I think we should use whenComplete here instead of whenSuccess and whenFailure. Reason for this is, we can save an allocation...

Second, we must check if the error we get is actually a PSQLError, if it is not, we need to wrap error with a PSQLError.

fabianfett · 2023-08-11T07:14:19Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

+            let action = self.preparedStatementState.lookup(
+                name: preparedStatement.name,
+                context: preparedStatement
+            )
+            switch action {
+            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 .executeStatement(let rowDescription):
+                psqlTask = .extendedQuery(.init(
+                    executeStatement: .init(
+                        name: preparedStatement.name,
+                        binds: preparedStatement.bindings,
+                        rowDescription: rowDescription),
+                    logger: preparedStatement.logger,
+                    promise: preparedStatement.promise
+                ))
+            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
+            }


can we move this into its own method? we should likely do the same with all the other longer ones above.

fabianfett · 2023-08-11T07:14:53Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

+        guard case .executePendingStatements(let statements, let rowDescription) = action else {
+            preconditionFailure("Expected to have pending statements to execute")
+        }


we should reconsider the return type from the statemachine here :). No enum needed.

fabianfett · 2023-08-11T07:15:03Z

Sources/PostgresNIO/New/PostgresChannelHandler.swift

+        guard case .returnError(let statements, let error) = action else {
+            preconditionFailure("Expected to have pending statements to execute")
+        }


fabianfett · 2023-08-11T09:56:55Z

Sources/PostgresNIO/New/PreparedStatement.swift

+///
+///     func makeBindings() -> PostgresBindings {
+///         var bindings = PostgresBindings()
+///         bindings.append(.init(string: self.state))


Suggested change

/// bindings.append(.init(string: self.state))

/// bindings.append(self.state)

does this work as well?

No, I get back an error: Cannot convert value of type 'String' to expected argument type 'PostgresData'

Oh, we shouldn't create PostgresData. The problem here is that we better names for those methods, so that users can use them without Interpolation as well.

struct PostgresBindings { @inlinable public mutating func appendInterpolation<Value: PostgresEncodable>(_ value: Value) throws { try self.binds.append(value, context: .default) self.sql.append(contentsOf: "$\(self.binds.count)") } @inlinable public mutating func appendInterpolation<Value: PostgresEncodable>(_ value: Optional<Value>) throws { switch value { case .none: self.binds.appendNull() case .some(let value): try self.binds.append(value, context: .default) } self.sql.append(contentsOf: "$\(self.binds.count)") } }

Sources/PostgresNIO/New/PostgresChannelHandler.swift

fabianfett · 2023-08-14T08:00:52Z

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

+        name: String,
+        rowDescription: RowDescription?
+    ) -> PreparationCompleteAction {
+        guard case .preparing(let statements) = self.preparedStatements[name] else {


can we switch over all cases here?

fabianfett · 2023-08-14T08:01:04Z

Sources/PostgresNIO/New/Connection State Machine/PreparedStatementStateMachine.swift

+    }
+
+    mutating func errorHappened(name: String, error: PSQLError) -> ErrorHappenedAction {
+        guard case .preparing(let statements) = self.preparedStatements[name] else {


can we switch over all cases here?

fabianfett · 2023-08-14T08:23:33Z

Sources/PostgresNIO/New/PreparedStatement.swift

+///
+///     func makeBindings() -> PostgresBindings {
+///         var bindings = PostgresBindings()
+///         bindings.append(.init(string: self.state))


Oh, we shouldn't create PostgresData. The problem here is that we better names for those methods, so that users can use them without Interpolation as well.

struct PostgresBindings { @inlinable public mutating func appendInterpolation<Value: PostgresEncodable>(_ value: Value) throws { try self.binds.append(value, context: .default) self.sql.append(contentsOf: "$\(self.binds.count)") } @inlinable public mutating func appendInterpolation<Value: PostgresEncodable>(_ value: Optional<Value>) throws { switch value { case .none: self.binds.appendNull() case .some(let value): try self.binds.append(value, context: .default) } self.sql.append(contentsOf: "$\(self.binds.count)") } }

Co-authored-by: Fabian Fett <fabianfett@apple.com>

fabianfett

This looks so good. A few nits.

Sources/PostgresNIO/Connection/PostgresConnection.swift

Sources/PostgresNIO/New/PostgresChannelHandler.swift

Sources/PostgresNIO/New/PreparedStatement.swift

fabianfett · 2023-08-17T07:10:50Z

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

+        }
+    }
+
+    func testPreparedStatement() async throws {


I know that these tests are somewhat painful to write but we should embrace them, as we can produce hard to reproduce scenarios here.

For this reason, please lets try to create those scenarios in extra tests here:

Sending the a prepared query twice, before the preparation step has succeeded. Is the order of the send statements still correct? Is the query just prepared once?

Sending the a prepared, after the the query of the same type has already succeeded. Is the query executed right away?

Failing a prepared statement in preparation are other executes failed in the same way right away?

Happy to help.

Co-authored-by: Fabian Fett <fabianfett@apple.com>

fabianfett

I'm extremely happy with this. Let's wait on the go from @gwynne and then let's merge it! Thanks @mariosangiorgio!

fabianfett

Super nits

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

Co-authored-by: Fabian Fett <fabianfett@apple.com>

mariosangiorgio commented Aug 9, 2023

View reviewed changes

fabianfett reviewed Aug 9, 2023

View reviewed changes

fabianfett added this to the 2.0.0 milestone Aug 9, 2023

mariosangiorgio force-pushed the async-prepared-statement branch from f90b4e5 to 4cf7f48 Compare August 10, 2023 14:02

fabianfett reviewed Aug 11, 2023

View reviewed changes

fabianfett added the semver-minor Adds new public API. label Aug 11, 2023

fabianfett mentioned this pull request Aug 11, 2023

Feature/typed queries #375

Draft

fabianfett reviewed Aug 11, 2023

View reviewed changes

mariosangiorgio commented Aug 11, 2023

View reviewed changes

Sources/PostgresNIO/New/PostgresChannelHandler.swift Show resolved Hide resolved

fabianfett reviewed Aug 14, 2023

View reviewed changes

mariosangiorgio and others added 6 commits August 16, 2023 13:36

async/await prepared statements

62b61ab

Address PR feedbacks

8283d8b

Apply suggestions from code review

9574725

Co-authored-by: Fabian Fett <fabianfett@apple.com>

PR feedback and state machine tests

b03c251

Add test to PostgresConnectionTests to test prepared statements

800ceaf

Explicit switch over all the state machine states

a426205

mariosangiorgio force-pushed the async-prepared-statement branch from 759587c to a426205 Compare August 16, 2023 12:37

mariosangiorgio added 3 commits August 16, 2023 14:33

PostgresBindings.append that use the default context

abe1cd1

More assertions in testPreparedStatement

ff0e139

Fix a build issue with Swift 5.6

725e778

mariosangiorgio marked this pull request as ready for review August 16, 2023 15:24

fabianfett reviewed Aug 17, 2023

View reviewed changes

mariosangiorgio and others added 3 commits August 17, 2023 11:34

Apply suggestions from code review

44309e3

Co-authored-by: Fabian Fett <fabianfett@apple.com>

Fix build issues

0740293

Add more PostgresConnectionTests cases

660e78d

fabianfett requested a review from gwynne August 18, 2023 08:23

fabianfett approved these changes Aug 18, 2023

View reviewed changes

Apply suggestions from code review

da87ce0

Co-authored-by: Fabian Fett <fabianfett@apple.com>

gwynne approved these changes Aug 18, 2023

View reviewed changes

fabianfett merged commit d5c5258 into vapor:main Aug 18, 2023
13 checks passed

mariosangiorgio deleted the async-prepared-statement branch August 18, 2023 10:56

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

async/await prepared statement API #390

async/await prepared statement API #390

mariosangiorgio commented Aug 9, 2023

mariosangiorgio Aug 9, 2023

fabianfett Aug 9, 2023

mariosangiorgio Aug 9, 2023

fabianfett Aug 11, 2023

mariosangiorgio Aug 9, 2023

fabianfett left a comment

fabianfett Aug 9, 2023

fabianfett Aug 9, 2023

fabianfett Aug 9, 2023

fabianfett Aug 9, 2023

fabianfett Aug 9, 2023

fabianfett Aug 9, 2023

fabianfett Aug 9, 2023

fabianfett Aug 9, 2023

fabianfett commented Aug 9, 2023 •

edited

codecov-commenter commented Aug 10, 2023 •

edited

fabianfett left a comment

fabianfett Aug 11, 2023

fabianfett Aug 11, 2023

fabianfett Aug 11, 2023

fabianfett Aug 11, 2023

fabianfett Aug 11, 2023

fabianfett Aug 11, 2023

fabianfett Aug 11, 2023

fabianfett Aug 11, 2023

mariosangiorgio Aug 11, 2023

fabianfett Aug 14, 2023

fabianfett Aug 14, 2023

fabianfett Aug 14, 2023

fabianfett Aug 14, 2023

fabianfett left a comment

fabianfett Aug 17, 2023

fabianfett left a comment

fabianfett left a comment

		self.listenState = ListenStateMachine()
		self.preparedStatementState = PreparedStatementStateMachine()

	/// bindings.append(.init(string: self.state))
	/// bindings.append(self.state)

async/await prepared statement API #390

async/await prepared statement API #390

Conversation

mariosangiorgio commented Aug 9, 2023

Example usage

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fabianfett left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fabianfett commented Aug 9, 2023 • edited

codecov-commenter commented Aug 10, 2023 • edited

Codecov Report

fabianfett left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fabianfett left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fabianfett left a comment

Choose a reason for hiding this comment

fabianfett left a comment

Choose a reason for hiding this comment

fabianfett commented Aug 9, 2023 •

edited

codecov-commenter commented Aug 10, 2023 •

edited