Add BufferedMessageEncoder, fix PSQLFrontendMessageEncoder

Author: fabianfett

Sources/PostgresNIO/New/BufferedMessageEncoder.swift

@@ -0,0 +1,39 @@
+import NIOCore
+
+struct BufferedMessageEncoder<Encoder: MessageToByteEncoder> {
+    private enum State {
+        case flushed
+        case writable
+    }
+    
+    private var buffer: ByteBuffer
+    private var state: State = .writable
+    private var encoder: Encoder
+    
+    init(buffer: ByteBuffer, encoder: Encoder) {
+        self.buffer = buffer
+        self.encoder = encoder
+    }
+    
+    mutating func encode(_ message: Encoder.OutboundIn) throws {
+        switch self.state {
+        case .flushed:
+            self.state = .writable
+            self.buffer.clear()
+            
+        case .writable:
+            break
+        }
+        
+        try self.encoder.encode(data: message, out: &self.buffer)
+    }
+    
+    mutating func flush() -> ByteBuffer? {
+        guard self.buffer.readableBytes > 0 else {
+            return nil
+        }
+        
+        self.state = .flushed
+        return self.buffer
+    }
+}

Sources/PostgresNIO/New/PSQLConnection.swift

@@ -231,7 +231,7 @@ final class PSQLConnection {
 
                         return channel.pipeline.addHandlers([
                             decoder,
-                            MessageToByteHandler(PSQLFrontendMessage.Encoder(jsonEncoder: configuration.coders.jsonEncoder)),
+                            MessageToByteHandler(PSQLFrontendMessageEncoder(jsonEncoder: configuration.coders.jsonEncoder)),
                             PSQLChannelHandler(
                                 authentification: configuration.authentication,
                                 logger: logger,

Sources/PostgresNIO/New/PSQLFrontendMessage.swift

@@ -129,82 +129,6 @@ extension PSQLFrontendMessage {
     }
 }
 
-extension PSQLFrontendMessage {
-    struct Encoder: MessageToByteEncoder {
-        typealias OutboundIn = PSQLFrontendMessage
-        
-        let jsonEncoder: PSQLJSONEncoder
-        
-        init(jsonEncoder: PSQLJSONEncoder) {
-            self.jsonEncoder = jsonEncoder
-        }
-        
-        func encode(data message: PSQLFrontendMessage, out buffer: inout ByteBuffer) throws {
-            struct EmptyPayload: PSQLMessagePayloadEncodable {
-                func encode(into buffer: inout ByteBuffer) {}
-            }
-            
-            func encode<Payload: PSQLMessagePayloadEncodable>(_ payload: Payload, into buffer: inout ByteBuffer) {
-                let startIndex = buffer.writerIndex
-                buffer.writeInteger(Int32(0)) // placeholder for length
-                payload.encode(into: &buffer)
-                let length = Int32(buffer.writerIndex - startIndex)
-                buffer.setInteger(length, at: startIndex)
-            }
-            
-            switch message {
-            case .bind(let bind):
-                buffer.writeInteger(message.id.rawValue)
-                let startIndex = buffer.writerIndex
-                buffer.writeInteger(Int32(0)) // placeholder for length
-                try bind.encode(into: &buffer, using: self.jsonEncoder)
-                let length = Int32(buffer.writerIndex - startIndex)
-                buffer.setInteger(length, at: startIndex)
-                
-            case .cancel(let cancel):
-                // cancel requests don't have an identifier
-                encode(cancel, into: &buffer)
-            case .close(let close):
-                buffer.writeFrontendMessageID(message.id)
-                encode(close, into: &buffer)
-            case .describe(let describe):
-                buffer.writeFrontendMessageID(message.id)
-                encode(describe, into: &buffer)
-            case .execute(let execute):
-                buffer.writeFrontendMessageID(message.id)
-                encode(execute, into: &buffer)
-            case .flush:
-                buffer.writeFrontendMessageID(message.id)
-                encode(EmptyPayload(), into: &buffer)
-            case .parse(let parse):
-                buffer.writeFrontendMessageID(message.id)
-                encode(parse, into: &buffer)
-            case .password(let password):
-                buffer.writeFrontendMessageID(message.id)
-                encode(password, into: &buffer)
-            case .saslInitialResponse(let saslInitialResponse):
-                buffer.writeFrontendMessageID(message.id)
-                encode(saslInitialResponse, into: &buffer)
-            case .saslResponse(let saslResponse):
-                buffer.writeFrontendMessageID(message.id)
-                encode(saslResponse, into: &buffer)
-            case .sslRequest(let request):
-                // sslRequests don't have an identifier
-                encode(request, into: &buffer)
-            case .startup(let startup):
-                // startup requests don't have an identifier
-                encode(startup, into: &buffer)
-            case .sync:
-                buffer.writeFrontendMessageID(message.id)
-                encode(EmptyPayload(), into: &buffer)
-            case .terminate:
-                buffer.writeFrontendMessageID(message.id)
-                encode(EmptyPayload(), into: &buffer)
-            }
-        }
-    }
-}
-
 protocol PSQLMessagePayloadEncodable {
     func encode(into buffer: inout ByteBuffer)
 }

Sources/PostgresNIO/New/PSQLFrontendMessageEncoder.swift

@@ -0,0 +1,88 @@
+
+struct PSQLFrontendMessageEncoder: MessageToByteEncoder {
+    typealias OutboundIn = PSQLFrontendMessage
+    
+    let jsonEncoder: PSQLJSONEncoder
+    
+    init(jsonEncoder: PSQLJSONEncoder) {
+        self.jsonEncoder = jsonEncoder
+    }
+    
+    func encode(data message: PSQLFrontendMessage, out buffer: inout ByteBuffer) throws {
+        switch message {
+        case .bind(let bind):
+            buffer.writeInteger(message.id.rawValue)
+            let startIndex = buffer.writerIndex
+            buffer.writeInteger(Int32(0)) // placeholder for length
+            try bind.encode(into: &buffer, using: self.jsonEncoder)
+            let length = Int32(buffer.writerIndex - startIndex)
+            buffer.setInteger(length, at: startIndex)
+            
+        case .cancel(let cancel):
+            // cancel requests don't have an identifier
+            self.encode(payload: cancel, into: &buffer)
+            
+        case .close(let close):
+            self.encode(messageID: message.id, payload: close, into: &buffer)
+
+        case .describe(let describe):
+            self.encode(messageID: message.id, payload: describe, into: &buffer)
+
+        case .execute(let execute):
+            self.encode(messageID: message.id, payload: execute, into: &buffer)
+
+        case .flush:
+            self.encode(messageID: message.id, payload: EmptyPayload(), into: &buffer)
+
+        case .parse(let parse):
+            self.encode(messageID: message.id, payload: parse, into: &buffer)
+
+        case .password(let password):
+            self.encode(messageID: message.id, payload: password, into: &buffer)
+
+        case .saslInitialResponse(let saslInitialResponse):
+            self.encode(messageID: message.id, payload: saslInitialResponse, into: &buffer)
+
+        case .saslResponse(let saslResponse):
+            self.encode(messageID: message.id, payload: saslResponse, into: &buffer)
+
+        case .sslRequest(let request):
+            // sslRequests don't have an identifier
+            self.encode(payload: request, into: &buffer)
+            
+        case .startup(let startup):
+            // startup requests don't have an identifier
+            self.encode(payload: startup, into: &buffer)
+            
+        case .sync:
+            self.encode(messageID: message.id, payload: EmptyPayload(), into: &buffer)
+
+        case .terminate:
+            self.encode(messageID: message.id, payload: EmptyPayload(), into: &buffer)
+        }
+    }
+    
+    private struct EmptyPayload: PSQLMessagePayloadEncodable {
+        func encode(into buffer: inout ByteBuffer) {}
+    }
+    
+    private func encode<Payload: PSQLMessagePayloadEncodable>(
+        messageID: PSQLFrontendMessage.ID,
+        payload: Payload,
+        into buffer: inout ByteBuffer)
+    {
+        buffer.writeFrontendMessageID(messageID)
+        self.encode(payload: payload, into: &buffer)
+    }
+    
+    private func encode<Payload: PSQLMessagePayloadEncodable>(
+        payload: Payload,
+        into buffer: inout ByteBuffer)
+    {
+        let startIndex = buffer.writerIndex
+        buffer.writeInteger(Int32(0)) // placeholder for length
+        payload.encode(into: &buffer)
+        let length = Int32(buffer.writerIndex - startIndex)
+        buffer.setInteger(length, at: startIndex)
+    }
+}

Tests/PostgresNIOTests/New/Extensions/PSQLCoding+TestUtils.swift

@@ -1,7 +1,7 @@
 @testable import PostgresNIO
 import Foundation
 
-extension PSQLFrontendMessage.Encoder {
+extension PSQLFrontendMessageEncoder {
     static var forTests: Self {
         Self(jsonEncoder: JSONEncoder())
     }

Tests/PostgresNIOTests/New/Messages/BindTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class BindTests: XCTestCase {
 
     func testEncodeBind() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let bind = PSQLFrontendMessage.Bind(portalName: "", preparedStatementName: "", parameters: ["Hello", "World"])
         let message = PSQLFrontendMessage.bind(bind)

Tests/PostgresNIOTests/New/Messages/CancelTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class CancelTests: XCTestCase {
 
     func testEncodeCancel() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let cancel = PSQLFrontendMessage.Cancel(processID: 1234, secretKey: 4567)
         let message = PSQLFrontendMessage.cancel(cancel)

Tests/PostgresNIOTests/New/Messages/CloseTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class CloseTests: XCTestCase {
 
     func testEncodeClosePortal() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let message = PSQLFrontendMessage.close(.portal("Hello"))
         XCTAssertNoThrow(try encoder.encode(data: message, out: &byteBuffer))
@@ -19,7 +19,7 @@ class CloseTests: XCTestCase {
     }
 
     func testEncodeCloseUnnamedStatement() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let message = PSQLFrontendMessage.close(.preparedStatement(""))
         XCTAssertNoThrow(try encoder.encode(data: message, out: &byteBuffer))

Tests/PostgresNIOTests/New/Messages/DescribeTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class DescribeTests: XCTestCase {
 
     func testEncodeDescribePortal() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let message = PSQLFrontendMessage.describe(.portal("Hello"))
         XCTAssertNoThrow(try encoder.encode(data: message, out: &byteBuffer))
@@ -19,7 +19,7 @@ class DescribeTests: XCTestCase {
     }
 
     func testEncodeDescribeUnnamedStatement() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let message = PSQLFrontendMessage.describe(.preparedStatement(""))
         XCTAssertNoThrow(try encoder.encode(data: message, out: &byteBuffer))

Tests/PostgresNIOTests/New/Messages/ExecuteTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class ExecuteTests: XCTestCase {
 
     func testEncodeExecute() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let message = PSQLFrontendMessage.execute(.init(portalName: "", maxNumberOfRows: 0))
         XCTAssertNoThrow(try encoder.encode(data: message, out: &byteBuffer))

Tests/PostgresNIOTests/New/Messages/ParseTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class ParseTests: XCTestCase {
 
     func testEncode() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let parse = PSQLFrontendMessage.Parse(
             preparedStatementName: "test",

Tests/PostgresNIOTests/New/Messages/PasswordTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class PasswordTests: XCTestCase {
 
     func testEncodePassword() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         // md522d085ed8dc3377968dc1c1a40519a2a = "abc123" with salt 1, 2, 3, 4
         let message = PSQLFrontendMessage.password(.init(value: "md522d085ed8dc3377968dc1c1a40519a2a"))

Tests/PostgresNIOTests/New/Messages/SASLInitialResponseTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class SASLInitialResponseTests: XCTestCase {
 
     func testEncodeWithData() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let sasl = PSQLFrontendMessage.SASLInitialResponse(
             saslMechanism: "hello", initialData: [0, 1, 2, 3, 4, 5, 6, 7])
@@ -30,7 +30,7 @@ class SASLInitialResponseTests: XCTestCase {
     }
 
     func testEncodeWithoutData() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let sasl = PSQLFrontendMessage.SASLInitialResponse(
             saslMechanism: "hello", initialData: [])

Tests/PostgresNIOTests/New/Messages/SASLResponseTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class SASLResponseTests: XCTestCase {
 
     func testEncodeWithData() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let sasl = PSQLFrontendMessage.SASLResponse(data: [0, 1, 2, 3, 4, 5, 6, 7])
         let message = PSQLFrontendMessage.saslResponse(sasl)
@@ -21,7 +21,7 @@ class SASLResponseTests: XCTestCase {
     }
 
     func testEncodeWithoutData() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let sasl = PSQLFrontendMessage.SASLResponse(data: [])
         let message = PSQLFrontendMessage.saslResponse(sasl)

Tests/PostgresNIOTests/New/Messages/SSLRequestTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class SSLRequestTests: XCTestCase {
 
     func testSSLRequest() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         let request = PSQLFrontendMessage.SSLRequest()
         let message = PSQLFrontendMessage.sslRequest(request)

Tests/PostgresNIOTests/New/Messages/StartupTests.swift

@@ -5,7 +5,7 @@ import NIOCore
 class StartupTests: XCTestCase {
 
     func testStartupMessage() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
 
         let replicationValues: [PSQLFrontendMessage.Startup.Parameters.Replication] = [

Tests/PostgresNIOTests/New/PSQLFrontendMessageTests.swift

@@ -23,7 +23,7 @@ class PSQLFrontendMessageTests: XCTestCase {
     // MARK: Encoder
 
     func testEncodeFlush() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         XCTAssertNoThrow(try encoder.encode(data: .flush, out: &byteBuffer))
 
@@ -33,7 +33,7 @@ class PSQLFrontendMessageTests: XCTestCase {
     }
 
     func testEncodeSync() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         XCTAssertNoThrow(try encoder.encode(data: .sync, out: &byteBuffer))
 
@@ -43,7 +43,7 @@ class PSQLFrontendMessageTests: XCTestCase {
     }
 
     func testEncodeTerminate() {
-        let encoder = PSQLFrontendMessage.Encoder.forTests
+        let encoder = PSQLFrontendMessageEncoder.forTests
         var byteBuffer = ByteBuffer()
         XCTAssertNoThrow(try encoder.encode(data: .terminate, out: &byteBuffer))