Skip to content

Commit 90c0302

Browse files
fabianfettfabianfett
committed
Support additional connection parameters
1 parent dd1b740 commit 90c0302

File tree

7 files changed

+88
-42
lines changed

7 files changed

+88
-42
lines changed

‎Sources/PostgresNIO/Connection/PostgresConnection+Configuration.swift

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@ import NIOSSL
44

55
extension PostgresConnection {
66
/// A configuration object for a connection
7-
public struct Configuration {
8-
7+
public struct Configuration: Sendable {
8+
99
// MARK: - TLS
1010

1111
/// The possible modes of operation for TLS encapsulation of a connection.
12-
public struct TLS {
12+
public struct TLS: Sendable {
1313
// MARK: Initializers
1414

1515
/// Do not try to create a TLS connection to the server.
@@ -63,7 +63,7 @@ extension PostgresConnection {
6363
// MARK: - Connection options
6464

6565
/// Describes options affecting how the underlying connection is made.
66-
public struct Options {
66+
public struct Options: Sendable {
6767
/// A timeout for connection attempts. Defaults to ten seconds.
6868
///
6969
/// Ignored when using a preexisting communcation channel. (See
@@ -85,14 +85,19 @@ extension PostgresConnection {
8585
/// This property is provided for compatibility with Amazon RDS Proxy, which requires it to be `false`.
8686
/// If you are not using Amazon RDS Proxy, you should leave this set to `true` (the default).
8787
public var requireBackendKeyData: Bool
88-
88+
89+
/// Additional parameters to send to the server on startup. The name value pairs are added to the initial
90+
/// startup message that the client sends to the server.
91+
public var additionalStartupParameters: [(String, String)]
92+
8993
/// Create an options structure with default values.
9094
///
9195
/// Most users should not need to adjust the defaults.
9296
public init() {
9397
self.connectTimeout = .seconds(10)
9498
self.tlsServerName = nil
9599
self.requireBackendKeyData = true
100+
self.additionalStartupParameters = []
96101
}
97102
}
98103

‎Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1249,11 +1249,19 @@ struct SendPrepareStatement {
12491249
let query: String
12501250
}
12511251

1252-
struct AuthContext: Equatable, CustomDebugStringConvertible {
1253-
let username: String
1254-
let password: String?
1255-
let database: String?
1256-
1252+
struct AuthContext: CustomDebugStringConvertible {
1253+
var username: String
1254+
var password: String?
1255+
var database: String?
1256+
var additionalParameters: [(String, String)]
1257+
1258+
init(username: String, password: String? = nil, database: String? = nil, additionalParameters: [(String, String)] = []) {
1259+
self.username = username
1260+
self.password = password
1261+
self.database = database
1262+
self.additionalParameters = additionalParameters
1263+
}
1264+
12571265
var debugDescription: String {
12581266
"""
12591267
AuthContext(username: \(String(reflecting: self.username)), \
@@ -1263,6 +1271,28 @@ struct AuthContext: Equatable, CustomDebugStringConvertible {
12631271
}
12641272
}
12651273

1274+
extension AuthContext: Equatable {
1275+
static func ==(lhs: Self, rhs: Self) -> Bool {
1276+
guard lhs.username == rhs.username
1277+
&& lhs.password == rhs.password
1278+
&& lhs.database == rhs.database
1279+
&& lhs.additionalParameters.count == rhs.additionalParameters.count
1280+
else {
1281+
return false
1282+
}
1283+
1284+
var lhsIterator = lhs.additionalParameters.makeIterator()
1285+
var rhsIterator = rhs.additionalParameters.makeIterator()
1286+
1287+
while let lhsNext = lhsIterator.next(), let rhsNext = rhsIterator.next() {
1288+
guard lhsNext.0 == rhsNext.0 && lhsNext.1 == rhsNext.1 else {
1289+
return false
1290+
}
1291+
}
1292+
return true
1293+
}
1294+
}
1295+
12661296
enum PasswordAuthencationMode: Equatable {
12671297
case cleartext
12681298
case md5(salt: (UInt8, UInt8, UInt8, UInt8))

‎Sources/PostgresNIO/New/Messages/Startup.swift

Lines changed: 29 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ extension PostgresFrontendMessage {
1717
/// name and value strings. A zero byte is required as a terminator after
1818
/// the last name/value pair. `user` is required, others are optional.
1919
struct Parameters: Equatable {
20-
enum Replication {
20+
enum Replication: Equatable {
2121
case `true`
2222
case `false`
2323
case database
@@ -29,16 +29,32 @@ extension PostgresFrontendMessage {
2929
/// The database to connect to. Defaults to the user name.
3030
var database: String?
3131

32-
/// Command-line arguments for the backend. (This is deprecated in favor
33-
/// of setting individual run-time parameters.) Spaces within this string are
34-
/// considered to separate arguments, unless escaped with a
35-
/// backslash (\); write \\ to represent a literal backslash.
36-
var options: String?
37-
3832
/// Used to connect in streaming replication mode, where a small set of
3933
/// replication commands can be issued instead of SQL statements. Value
4034
/// can be true, false, or database, and the default is false.
4135
var replication: Replication
36+
37+
var options: [(String, String)]
38+
39+
static func ==(lhs: Self, rhs: Self) -> Bool {
40+
guard lhs.user == rhs.user
41+
&& lhs.database == rhs.database
42+
&& lhs.replication == rhs.replication
43+
&& lhs.options.count == rhs.options.count
44+
else {
45+
return false
46+
}
47+
48+
var lhsIterator = lhs.options.makeIterator()
49+
var rhsIterator = rhs.options.makeIterator()
50+
51+
while let lhsNext = lhsIterator.next(), let rhsNext = rhsIterator.next() {
52+
guard lhsNext.0 == rhsNext.0 && lhsNext.1 == rhsNext.1 else {
53+
return false
54+
}
55+
}
56+
return true
57+
}
4258
}
4359
var parameters: Parameters
4460

@@ -59,11 +75,6 @@ extension PostgresFrontendMessage {
5975
buffer.writeNullTerminatedString(database)
6076
}
6177

62-
if let options = self.parameters.options {
63-
buffer.writeNullTerminatedString("options")
64-
buffer.writeNullTerminatedString(options)
65-
}
66-
6778
switch self.parameters.replication {
6879
case .database:
6980
buffer.writeNullTerminatedString("replication")
@@ -74,9 +85,13 @@ extension PostgresFrontendMessage {
7485
case .false:
7586
break
7687
}
77-
88+
89+
for (key, value) in self.parameters.options {
90+
buffer.writeNullTerminatedString(key)
91+
buffer.writeNullTerminatedString(value)
92+
}
93+
7894
buffer.writeInteger(UInt8(0))
7995
}
8096
}
81-
8297
}

‎Sources/PostgresNIO/New/PostgresChannelHandler.swift

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -521,8 +521,8 @@ extension AuthContext {
521521
PostgresFrontendMessage.Startup.Parameters(
522522
user: self.username,
523523
database: self.database,
524-
options: nil,
525-
replication: .false
524+
replication: .false,
525+
options: self.additionalParameters
526526
)
527527
}
528528
}

‎Tests/PostgresNIOTests/New/Extensions/PSQLFrontendMessageDecoder.swift

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,8 @@ struct PSQLFrontendMessageDecoder: NIOSingleStepByteToMessageDecoder {
3939
case 196608:
4040
var user: String?
4141
var database: String?
42-
var options: String?
43-
42+
var options = [(String, String)]()
43+
4444
while let name = messageSlice.readNullTerminatedString(), messageSlice.readerIndex < finalIndex {
4545
let value = messageSlice.readNullTerminatedString()
4646

@@ -51,19 +51,18 @@ struct PSQLFrontendMessageDecoder: NIOSingleStepByteToMessageDecoder {
5151
case "database":
5252
database = value
5353

54-
case "options":
55-
options = value
56-
5754
default:
58-
break
55+
if let value = value {
56+
options.append((name, value))
57+
}
5958
}
6059
}
6160

6261
let parameters = PostgresFrontendMessage.Startup.Parameters(
6362
user: user!,
6463
database: database,
65-
options: options,
66-
replication: .false
64+
replication: .false,
65+
options: options
6766
)
6867

6968
let startup = PostgresFrontendMessage.Startup(

‎Tests/PostgresNIOTests/New/Messages/StartupTests.swift

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@ class StartupTests: XCTestCase {
1818
let parameters = PostgresFrontendMessage.Startup.Parameters(
1919
user: "test",
2020
database: "abc123",
21-
options: "some options",
22-
replication: replication
21+
replication: replication,
22+
options: [("some", "options")]
2323
)
2424

2525
let startup = PostgresFrontendMessage.Startup.versionThree(parameters: parameters)
@@ -33,12 +33,12 @@ class StartupTests: XCTestCase {
3333
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "test")
3434
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "database")
3535
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "abc123")
36-
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "options")
37-
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "some options")
3836
if replication != .false {
3937
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "replication")
4038
XCTAssertEqual(byteBuffer.readNullTerminatedString(), replication.stringValue)
4139
}
40+
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "some")
41+
XCTAssertEqual(byteBuffer.readNullTerminatedString(), "options")
4242
XCTAssertEqual(byteBuffer.readInteger(), UInt8(0))
4343

4444
XCTAssertEqual(byteBuffer.readableBytes, 0)

‎Tests/PostgresNIOTests/New/PostgresChannelHandlerTests.swift

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,7 @@ class PostgresChannelHandlerTests: XCTestCase {
2626
return XCTFail("Unexpected message")
2727
}
2828

29-
XCTAssertEqual(startup.parameters.user, config.username)
30-
XCTAssertEqual(startup.parameters.database, config.database)
31-
XCTAssertEqual(startup.parameters.options, nil)
32-
XCTAssertEqual(startup.parameters.replication, .false)
29+
XCTAssertEqual(startup.parameters, try .init(user: XCTUnwrap(config.username), database: config.database, replication: .false, options: config.options.additionalStartupParameters))
3330

3431
XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.authentication(.ok)))
3532
XCTAssertNoThrow(try embedded.writeInbound(PostgresBackendMessage.backendKeyData(.init(processID: 1234, secretKey: 5678))))

0 commit comments

Comments
 (0)