Skip to content

Commit 9cc7158

Browse files
adam-fowleradam-fowler
authored
Role command custom response (#141)
* Add parsing of role command response Signed-off-by: Adam Fowler <[email protected]> * minor edit Signed-off-by: Adam Fowler <[email protected]> * Add test for replica role Signed-off-by: Adam Fowler <[email protected]> --------- Signed-off-by: Adam Fowler <[email protected]>
1 parent f7a2599 commit 9cc7158

File tree

6 files changed

+211
-6
lines changed

6 files changed

+211
-6
lines changed

‎Sources/Valkey/Commands/Custom/ServerCustomCommands.swift

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the valkey-swift project
4+
//
5+
// Copyright (c) 2025 the valkey-swift authors
6+
// Licensed under Apache License v2.0
7+
//
8+
// See LICENSE.txt for license information
9+
// See valkey-swift/CONTRIBUTORS.txt for the list of valkey-swift authors
10+
//
11+
// SPDX-License-Identifier: Apache-2.0
12+
//
13+
//===----------------------------------------------------------------------===//
14+
15+
extension ROLE {
16+
public enum Response: RESPTokenDecodable, Sendable {
17+
struct DecodeError: Error {}
18+
public struct Primary: Sendable {
19+
public struct Replica: RESPTokenDecodable, Sendable {
20+
public let ip: String
21+
public let port: Int
22+
public let replicationOffset: Int
23+
24+
public init(fromRESP token: RESPToken) throws {
25+
(self.ip, self.port, self.replicationOffset) = try token.decodeArrayElements()
26+
}
27+
}
28+
public let replicationOffset: Int
29+
public let replicas: [Replica]
30+
31+
init(arrayIterator: inout RESPToken.Array.Iterator) throws {
32+
guard let replicationOffsetToken = arrayIterator.next(), let replicasToken = arrayIterator.next() else {
33+
throw DecodeError()
34+
}
35+
self.replicationOffset = try .init(fromRESP: replicationOffsetToken)
36+
self.replicas = try .init(fromRESP: replicasToken)
37+
}
38+
}
39+
public struct Replica: Sendable {
40+
public enum State: String, RESPTokenDecodable, Sendable {
41+
case connect
42+
case connecting
43+
case sync
44+
case connected
45+
46+
public init(fromRESP token: RESPToken) throws {
47+
let string = try String(fromRESP: token)
48+
guard let state = State(rawValue: string) else {
49+
throw RESPParsingError(code: .unexpectedType, buffer: token.base)
50+
}
51+
self = state
52+
}
53+
}
54+
public let primaryIP: String
55+
public let primaryPort: Int
56+
public let state: State
57+
public let replicationOffset: Int
58+
59+
init(arrayIterator: inout RESPToken.Array.Iterator) throws {
60+
guard let primaryIPToken = arrayIterator.next(),
61+
let primaryPortToken = arrayIterator.next(),
62+
let stateToken = arrayIterator.next(),
63+
let replicationToken = arrayIterator.next()
64+
else {
65+
throw DecodeError()
66+
}
67+
self.primaryIP = try .init(fromRESP: primaryIPToken)
68+
self.primaryPort = try .init(fromRESP: primaryPortToken)
69+
self.state = try .init(fromRESP: stateToken)
70+
self.replicationOffset = try .init(fromRESP: replicationToken)
71+
}
72+
}
73+
public struct Sentinel: Sendable {
74+
public let primaryNames: [String]
75+
76+
init(arrayIterator: inout RESPToken.Array.Iterator) throws {
77+
guard let primaryNamesToken = arrayIterator.next() else { throw DecodeError() }
78+
self.primaryNames = try .init(fromRESP: primaryNamesToken)
79+
}
80+
}
81+
case primary(Primary)
82+
case replica(Replica)
83+
case sentinel(Sentinel)
84+
85+
public init(fromRESP token: RESPToken) throws {
86+
switch token.value {
87+
case .array(let array):
88+
do {
89+
var iterator = array.makeIterator()
90+
guard let roleToken = iterator.next() else {
91+
throw RESPParsingError(code: .unexpectedType, buffer: token.base)
92+
}
93+
let role = try String(fromRESP: roleToken)
94+
switch role {
95+
case "master":
96+
let primary = try Primary(arrayIterator: &iterator)
97+
self = .primary(primary)
98+
case "slave":
99+
let replica = try Replica(arrayIterator: &iterator)
100+
self = .replica(replica)
101+
case "sentinel":
102+
let sentinel = try Sentinel(arrayIterator: &iterator)
103+
self = .sentinel(sentinel)
104+
default:
105+
throw DecodeError()
106+
}
107+
} catch {
108+
throw RESPParsingError(code: .unexpectedType, buffer: token.base)
109+
}
110+
default:
111+
throw RESPParsingError(code: .unexpectedType, buffer: token.base)
112+
}
113+
}
114+
}
115+
}

‎Sources/Valkey/Commands/ServerCommands.swift

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1271,8 +1271,6 @@ public struct REPLICAOF: ValkeyCommand {
12711271
/// Returns the replication role.
12721272
@_documentation(visibility: internal)
12731273
public struct ROLE: ValkeyCommand {
1274-
public typealias Response = RESPToken.Array
1275-
12761274
@inlinable public init() {
12771275
}
12781276

@@ -2205,8 +2203,7 @@ extension ValkeyConnectionProtocol {
22052203
/// - Available: 2.8.12
22062204
/// - Complexity: O(1)
22072205
@inlinable
2208-
@discardableResult
2209-
public func role() async throws -> RESPToken.Array {
2206+
public func role() async throws -> ROLE.Response {
22102207
try await send(command: ROLE())
22112208
}
22122209

‎Sources/Valkey/RESP/RESPTokenDecodable.swift

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -146,8 +146,13 @@ extension Int64: RESPTokenDecodable {
146146
case .number(let value):
147147
self = value
148148

149-
case .bulkString,
150-
.simpleString,
149+
case .bulkString(let buffer):
150+
guard let value = Int64(String(buffer: buffer)) else {
151+
throw RESPParsingError(code: .canNotParseInteger, buffer: token.base)
152+
}
153+
self = value
154+
155+
case .simpleString,
151156
.bulkError,
152157
.simpleError,
153158
.verbatimString,

‎Sources/_ValkeyCommandsBuilder/ValkeyCommandsRender.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ func renderValkeyCommands(_ commands: [String: ValkeyCommand], fullCommandList:
2121
"GEODIST",
2222
"GEOPOS",
2323
"GEOSEARCH",
24+
"ROLE",
2425
"LMPOP",
2526
"SPOP",
2627
"SSCAN",

‎Tests/IntegrationTests/ValkeyTests.swift

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,19 @@ struct GeneratedCommands {
302302
}
303303
}
304304

305+
@Test
306+
@available(valkeySwift 1.0, *)
307+
func testRole() async throws {
308+
var logger = Logger(label: "Valkey")
309+
logger.logLevel = .debug
310+
try await withValkeyConnection(.hostname(valkeyHostname, port: 6379), logger: logger) { connection in
311+
try await withKey(connection: connection) { key in
312+
let role = try await connection.role()
313+
print(role)
314+
}
315+
}
316+
}
317+
305318
@available(valkeySwift 1.0, *)
306319
@Test("Array with count using LMPOP")
307320
func testArrayWithCount() async throws {

‎Tests/ValkeyTests/CommandTests.swift

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,80 @@ import Valkey
2323
/// Generally the commands being tested here are ones we have written custom responses for
2424
struct CommandTests {
2525
struct ServerCommands {
26+
@Test
27+
@available(valkeySwift 1.0, *)
28+
func role() async throws {
29+
let channel = NIOAsyncTestingChannel()
30+
let logger = Logger(label: "test")
31+
let connection = try await ValkeyConnection.setupChannelAndConnect(channel, configuration: .init(), logger: logger)
32+
try await channel.processHello()
33+
34+
try await withThrowingTaskGroup(of: Void.self) { group in
35+
group.addTask {
36+
var role = try await connection.role()
37+
guard case .primary(let primary) = role else {
38+
Issue.record()
39+
return
40+
}
41+
#expect(primary.replicationOffset == 10)
42+
#expect(primary.replicas.count == 2)
43+
#expect(primary.replicas[0].ip == "127.0.0.1")
44+
#expect(primary.replicas[0].port == 9001)
45+
#expect(primary.replicas[0].replicationOffset == 1)
46+
#expect(primary.replicas[1].ip == "127.0.0.1")
47+
#expect(primary.replicas[1].port == 9002)
48+
#expect(primary.replicas[1].replicationOffset == 6)
49+
50+
role = try await connection.role()
51+
guard case .replica(let replica) = role else {
52+
Issue.record()
53+
return
54+
}
55+
#expect(replica.primaryIP == "127.0.0.1")
56+
#expect(replica.primaryPort == 9000)
57+
#expect(replica.state == .connected)
58+
#expect(replica.replicationOffset == 6)
59+
}
60+
group.addTask {
61+
var outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
62+
#expect(outbound == RESPToken(.command(["ROLE"])).base)
63+
try await channel.writeInbound(
64+
RESPToken(
65+
.array([
66+
.bulkString("master"),
67+
.number(10),
68+
.array([
69+
.array([
70+
.bulkString("127.0.0.1"),
71+
.bulkString("9001"),
72+
.bulkString("1"),
73+
]),
74+
.array([
75+
.bulkString("127.0.0.1"),
76+
.bulkString("9002"),
77+
.bulkString("6"),
78+
]),
79+
]),
80+
])
81+
).base
82+
)
83+
outbound = try await channel.waitForOutboundWrite(as: ByteBuffer.self)
84+
#expect(outbound == RESPToken(.command(["ROLE"])).base)
85+
try await channel.writeInbound(
86+
RESPToken(
87+
.array([
88+
.bulkString("slave"),
89+
.bulkString("127.0.0.1"),
90+
.number(9000),
91+
.bulkString("connected"),
92+
.number(6),
93+
])
94+
).base
95+
)
96+
}
97+
try await group.waitForAll()
98+
}
99+
}
26100
/// Test non-optional tokens render correctly
27101
@Test
28102
@available(valkeySwift 1.0, *)

0 commit comments

Comments
 (0)