Skip to content

Commit a789f2c

Browse files
authored
Custom command responses for HSCAN/ZSCAN/BLPOP (#247)
* ZSCAN response Signed-off-by: Adam Fowler <[email protected]> * HSCAN Signed-off-by: Adam Fowler <[email protected]> * Add custom responses for BLPOP and BRPOP Signed-off-by: Adam Fowler <[email protected]> * Revert KEYS change Signed-off-by: Adam Fowler <[email protected]> * Convert all command tests to use testCommandEncodesDecodes Add tests for HSCAN, ZSCAN and BLPOP Signed-off-by: Adam Fowler <[email protected]> --------- Signed-off-by: Adam Fowler <[email protected]>
1 parent 22651ce commit a789f2c

File tree

9 files changed

+571
-547
lines changed

9 files changed

+571
-547
lines changed
Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
//
2+
// This source file is part of the valkey-swift project
3+
// Copyright (c) 2025 the valkey-swift project authors
4+
//
5+
// See LICENSE.txt for license information
6+
// SPDX-License-Identifier: Apache-2.0
7+
//
8+
import NIOCore
9+
10+
/// Sorted set entry
11+
@_documentation(visibility: internal)
12+
public struct HashEntry: RESPTokenDecodable, Sendable {
13+
public let field: ByteBuffer
14+
public let value: ByteBuffer
15+
16+
init(field: ByteBuffer, value: ByteBuffer) {
17+
self.field = field
18+
self.value = value
19+
}
20+
21+
public init(fromRESP token: RESPToken) throws {
22+
switch token.value {
23+
case .array(let array):
24+
(self.field, self.value) = try array.decodeElements()
25+
default:
26+
throw RESPDecodeError.tokenMismatch(expected: [.array], token: token)
27+
}
28+
}
29+
}
30+
31+
extension HSCAN {
32+
public struct Response: RESPTokenDecodable, Sendable {
33+
public struct Members: RESPTokenDecodable, Sendable {
34+
/// List of members and possibly scores.
35+
public let elements: RESPToken.Array
36+
37+
public init(fromRESP token: RESPToken) throws {
38+
self.elements = try token.decode(as: RESPToken.Array.self)
39+
}
40+
41+
/// if HSCAN was called with the `NOVALUES` parameter use this
42+
/// function to get an array of fields
43+
public func withoutValues() throws -> [ByteBuffer] {
44+
try self.elements.decode(as: [ByteBuffer].self)
45+
}
46+
47+
/// if HSCAN was called without the `NOVALUES` parameter use this
48+
/// function to get an array of fields and values
49+
public func withValues() throws -> [HashEntry] {
50+
var array: [HashEntry] = []
51+
for respElement in try self.elements.asMap() {
52+
let field = try ByteBuffer(fromRESP: respElement.key)
53+
let value = try ByteBuffer(fromRESP: respElement.value)
54+
array.append(.init(field: field, value: value))
55+
}
56+
return array
57+
}
58+
}
59+
/// Cursor to use in next call to HSCAN
60+
public let cursor: Int
61+
/// Sorted set members
62+
public let members: Members
63+
64+
public init(fromRESP token: RESPToken) throws {
65+
(self.cursor, self.members) = try token.decodeArrayElements()
66+
}
67+
}
68+
}

Sources/Valkey/Commands/Custom/ListCustomCommands.swift

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,17 @@
88

99
import NIOCore
1010

11+
/// List entry
12+
@_documentation(visibility: internal)
13+
public struct ListEntry: RESPTokenDecodable, Sendable {
14+
public let key: ValkeyKey
15+
public let value: ByteBuffer
16+
17+
public init(fromRESP token: RESPToken) throws {
18+
(self.key, self.value) = try token.decodeArrayElements()
19+
}
20+
}
21+
1122
extension LMOVE {
1223
public typealias Response = ByteBuffer?
1324
}
@@ -38,3 +49,17 @@ extension BLMPOP {
3849
/// * [Array]: List key from which elements were popped.
3950
public typealias Response = LMPOP.Response
4051
}
52+
53+
extension BLPOP {
54+
/// - Response: One of the following
55+
/// * [Null]: No element could be popped and timeout expired
56+
/// * [Array]: The key from which the element was popped and the value of the popped element
57+
public typealias Response = ListEntry?
58+
}
59+
60+
extension BRPOP {
61+
/// - Response: One of the following
62+
/// * [Null]: No element could be popped and the timeout expired.
63+
/// * [Array]: The key from which the element was popped and the value of the popped element
64+
public typealias Response = ListEntry?
65+
}

Sources/Valkey/Commands/Custom/SortedSetCustomCommands.swift

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,11 @@ public struct SortedSetEntry: RESPTokenDecodable, Sendable {
1313
public let value: ByteBuffer
1414
public let score: Double
1515

16+
init(value: ByteBuffer, score: Double) {
17+
self.value = value
18+
self.score = score
19+
}
20+
1621
public init(fromRESP token: RESPToken) throws {
1722
switch token.value {
1823
case .array(let array):
@@ -76,3 +81,42 @@ extension ZPOPMIN {
7681
/// * [Array]: List of popped elements and scores when 'COUNT' is specified.
7782
public typealias Response = [SortedSetEntry]
7883
}
84+
85+
extension ZSCAN {
86+
public struct Response: RESPTokenDecodable, Sendable {
87+
public struct Members: RESPTokenDecodable, Sendable {
88+
/// List of members and possibly scores.
89+
public let elements: RESPToken.Array
90+
91+
public init(fromRESP token: RESPToken) throws {
92+
self.elements = try token.decode(as: RESPToken.Array.self)
93+
}
94+
95+
/// if ZSCAN was called with the `NOSCORES` parameter use this
96+
/// function to get an array of members
97+
public func withoutScores() throws -> [ByteBuffer] {
98+
try self.elements.decode(as: [ByteBuffer].self)
99+
}
100+
101+
/// if ZSCAN was called without the `NOSCORES` parameter use this
102+
/// function to get an array of members and scores
103+
public func withScores() throws -> [SortedSetEntry] {
104+
var array: [SortedSetEntry] = []
105+
for respElement in try self.elements.asMap() {
106+
let value = try ByteBuffer(fromRESP: respElement.key)
107+
let score = try Double(fromRESP: respElement.value)
108+
array.append(.init(value: value, score: score))
109+
}
110+
return array
111+
}
112+
}
113+
/// Cursor to use in next call to ZSCAN
114+
public let cursor: Int
115+
/// Sorted set members
116+
public let members: Members
117+
118+
public init(fromRESP token: RESPToken) throws {
119+
(self.cursor, self.members) = try token.decodeArrayElements()
120+
}
121+
}
122+
}

Sources/Valkey/Commands/GenericCommands.swift

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,8 @@ public struct EXPIRETIME: ValkeyCommand {
324324
/// Returns all key names that match a pattern.
325325
@_documentation(visibility: internal)
326326
public struct KEYS: ValkeyCommand {
327+
public typealias Response = RESPToken.Array
328+
327329
@inlinable public static var name: String { "KEYS" }
328330

329331
public var pattern: String
@@ -1164,7 +1166,7 @@ extension ValkeyClientProtocol {
11641166
/// - Complexity: O(N) with N being the number of keys in the database, under the assumption that the key names in the database and the given pattern have limited length.
11651167
/// - Response: [Array]: List of keys matching pattern.
11661168
@inlinable
1167-
public func keys(pattern: String) async throws -> KEYS.Response {
1169+
public func keys(pattern: String) async throws -> RESPToken.Array {
11681170
try await execute(KEYS(pattern: pattern))
11691171
}
11701172

Sources/Valkey/Commands/HashCommands.swift

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -831,8 +831,6 @@ public struct HRANDFIELD: ValkeyCommand {
831831
/// Iterates over fields and values of a hash.
832832
@_documentation(visibility: internal)
833833
public struct HSCAN: ValkeyCommand {
834-
public typealias Response = RESPToken.Array
835-
836834
@inlinable public static var name: String { "HSCAN" }
837835

838836
public var key: ValkeyKey
@@ -1410,13 +1408,8 @@ extension ValkeyClientProtocol {
14101408
/// - Complexity: O(1) for every call. O(N) for a complete iteration, including enough command calls for the cursor to return back to 0. N is the number of elements inside the collection.
14111409
/// - Response: [Array]: Cursor and scan response in array form.
14121410
@inlinable
1413-
public func hscan(
1414-
_ key: ValkeyKey,
1415-
cursor: Int,
1416-
pattern: String? = nil,
1417-
count: Int? = nil,
1418-
novalues: Bool = false
1419-
) async throws -> RESPToken.Array {
1411+
public func hscan(_ key: ValkeyKey, cursor: Int, pattern: String? = nil, count: Int? = nil, novalues: Bool = false) async throws -> HSCAN.Response
1412+
{
14201413
try await execute(HSCAN(key, cursor: cursor, pattern: pattern, count: count, novalues: novalues))
14211414
}
14221415

Sources/Valkey/Commands/ListCommands.swift

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -119,8 +119,6 @@ public struct BLMPOP: ValkeyCommand {
119119
/// Removes and returns the first element in a list. Blocks until an element is available otherwise. Deletes the list if the last element was popped.
120120
@_documentation(visibility: internal)
121121
public struct BLPOP: ValkeyCommand {
122-
public typealias Response = RESPToken.Array?
123-
124122
@inlinable public static var name: String { "BLPOP" }
125123

126124
public var keys: [ValkeyKey]
@@ -143,8 +141,6 @@ public struct BLPOP: ValkeyCommand {
143141
/// Removes and returns the last element in a list. Blocks until an element is available otherwise. Deletes the list if the last element was popped.
144142
@_documentation(visibility: internal)
145143
public struct BRPOP: ValkeyCommand {
146-
public typealias Response = RESPToken.Array?
147-
148144
@inlinable public static var name: String { "BRPOP" }
149145

150146
public var keys: [ValkeyKey]
@@ -701,7 +697,7 @@ extension ValkeyClientProtocol {
701697
/// * [Array]: The key from which the element was popped and the value of the popped element
702698
@inlinable
703699
@discardableResult
704-
public func blpop(keys: [ValkeyKey], timeout: Double) async throws -> RESPToken.Array? {
700+
public func blpop(keys: [ValkeyKey], timeout: Double) async throws -> BLPOP.Response {
705701
try await execute(BLPOP(keys: keys, timeout: timeout))
706702
}
707703

@@ -717,7 +713,7 @@ extension ValkeyClientProtocol {
717713
/// * [Array]: The name of the key where an element was popped
718714
@inlinable
719715
@discardableResult
720-
public func brpop(keys: [ValkeyKey], timeout: Double) async throws -> RESPToken.Array? {
716+
public func brpop(keys: [ValkeyKey], timeout: Double) async throws -> BRPOP.Response {
721717
try await execute(BRPOP(keys: keys, timeout: timeout))
722718
}
723719

Sources/Valkey/Commands/SortedSetCommands.swift

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1143,8 +1143,6 @@ public struct ZREVRANK<Member: RESPStringRenderable>: ValkeyCommand {
11431143
/// Iterates over members and scores of a sorted set.
11441144
@_documentation(visibility: internal)
11451145
public struct ZSCAN: ValkeyCommand {
1146-
public typealias Response = RESPToken.Array
1147-
11481146
@inlinable public static var name: String { "ZSCAN" }
11491147

11501148
public var key: ValkeyKey
@@ -1788,13 +1786,8 @@ extension ValkeyClientProtocol {
17881786
/// - Complexity: O(1) for every call. O(N) for a complete iteration, including enough command calls for the cursor to return back to 0. N is the number of elements inside the collection.
17891787
/// - Response: [Array]: Cursor and scan response in array form.
17901788
@inlinable
1791-
public func zscan(
1792-
_ key: ValkeyKey,
1793-
cursor: Int,
1794-
pattern: String? = nil,
1795-
count: Int? = nil,
1796-
noscores: Bool = false
1797-
) async throws -> RESPToken.Array {
1789+
public func zscan(_ key: ValkeyKey, cursor: Int, pattern: String? = nil, count: Int? = nil, noscores: Bool = false) async throws -> ZSCAN.Response
1790+
{
17981791
try await execute(ZSCAN(key, cursor: cursor, pattern: pattern, count: count, noscores: noscores))
17991792
}
18001793

Sources/_ValkeyCommandsBuilder/ValkeyCommandsRender.swift

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@ import Foundation
1010
/// List of functions where the Response calculation has been disabled because we want
1111
/// to override the response in the Valkey library
1212
private let disableResponseCalculationCommands: Set<String> = [
13+
"BLPOP",
1314
"BLMPOP",
15+
"BRPOP",
1416
"BZMPOP",
1517
"BZPOPMAX",
1618
"BZPOPMIN",
@@ -24,7 +26,8 @@ private let disableResponseCalculationCommands: Set<String> = [
2426
"GEODIST",
2527
"GEOPOS",
2628
"GEOSEARCH",
27-
"KEYS",
29+
"HSCAN",
30+
"ROLE",
2831
"LMOVE",
2932
"LMPOP",
3033
"ROLE",
@@ -43,6 +46,7 @@ private let disableResponseCalculationCommands: Set<String> = [
4346
"ZMPOP",
4447
"ZPOPMAX",
4548
"ZPOPMIN",
49+
"ZSCAN",
4650
]
4751

4852
/// List of subscribe commands, which have their own implementation in code

0 commit comments

Comments
 (0)