60 -- Provide Strong Option Types in SortedSet Commands

Motivation:

While working through issue #59, it was noticed just how "stringly" the SortedSet command options for `zadd`, `zinterstore`, and `zunionstore` were, and Swift provides ways of having strong type safety for these options.

Modifications:

- Add `RedisSortedSetAddOption` and `RedisSortedSetAggregateMethod` to replace the String API in `zadd`, `zinterstore`, and `zunionstore`
- Fix an implication of how `overestimatedCountBeingAdded` documentation for `Array where Element == RESPValue` for `add(contentsOf:overestimatedCountBeingAdded:_:)`

Result:

Users should have a more discoverable and straightforward way that isn't error prone for calling `zadd`, `zinterstore`, and `zunionstore` with Redis supported options.

Author: Mordil

Sources/RedisNIO/Commands/SortedSetCommands.swift

@@ -46,6 +46,22 @@ extension RedisClient {
 
 // MARK: General
 
+/// The supported options for the `zadd` command with Redis SortedSet types.
+/// - Important: Per Redis documentation, `.onlyUpdateExistingElements` and `.onlyAddNewElements` are mutually exclusive!
+/// - Note: `INCR` is not supported by this library in `zadd`. Use the `zincrby(:element:in:)` method instead.
+/// See [https://redis.io/commands/zadd#zadd-options-redis-302-or-greater](https://redis.io/commands/zadd#zadd-options-redis-302-or-greater)
+public enum RedisSortedSetAddOption: String {
+    /// When adding elements, any that do not already exist in the SortedSet will be ignored and the score of the existing element will be updated.
+    case onlyUpdateExistingElements = "XX"
+    /// When adding elements, any that already exist in the SortedSet will be ignored and the score of the existing element will not be updated.
+    case onlyAddNewElements = "NX"
+    /// `zadd` normally returns the number of new elements added to the set,
+    /// but this option will instead have the command return the number of elements changed.
+    ///
+    /// "Changed" in this context are new elements added, and elements that had their score updated.
+    case returnChangedCount = "CH"
+}
+
 extension RedisClient {
     /// Adds elements to a sorted set, assigning their score to the values provided.
     ///
@@ -59,28 +75,21 @@ extension RedisClient {
     public func zadd<Value: RESPValueConvertible>(
         _ elements: [(element: Value, score: Double)],
         to key: String,
-        options: Set<String> = []
+        options: Set<RedisSortedSetAddOption> = []
     ) -> EventLoopFuture<Int> {
-        guard !options.contains("INCR") else {
-            return self.eventLoop.makeFailedFuture(RedisNIOError.unsupportedOperation(
-                method: #function,
-                message: "INCR option is unsupported. Use zincrby(_:element:in:) instead."
-            ))
-        }
-
         assert(options.count <= 2, "Invalid number of options provided.")
-        assert(options.allSatisfy(["XX", "NX", "CH"].contains), "Unsupported option provided!")
         assert(
-            !(options.contains("XX") && options.contains("NX")),
-            "XX and NX options are mutually exclusive."
+            !(options.contains(.onlyAddNewElements) && options.contains(.onlyUpdateExistingElements)),
+            ".onlyAddNewElements and .onlyUpdateExistingElements options are mutually exclusive."
         )
 
         var args: [RESPValue] = [.init(bulk: key)]
-        args.append(convertingContentsOf: options)
-
-        for (element, score) in elements {
-            args.append(.init(bulk: score.description))
-            args.append(element.convertedToRESPValue())
+        args.add(contentsOf: options) { (array, option) in
+            array.append(.init(bulk: option.rawValue))
+        }
+        args.add(contentsOf: elements, overestimatedCountBeingAdded: elements.count * 2) { (array, next) in
+            array.append(.init(bulk: next.score.description))
+            array.append(next.element.convertedToRESPValue())
         }
 
         return send(command: "ZADD", with: args)
@@ -99,7 +108,7 @@ extension RedisClient {
     public func zadd<Value: RESPValueConvertible>(
         _ element: (element: Value, score: Double),
         to key: String,
-        options: Set<String> = []
+        options: Set<RedisSortedSetAddOption> = []
     ) -> EventLoopFuture<Bool> {
         return zadd([element], to: key, options: options)
             .map { return $0 == 1 }
@@ -490,6 +499,20 @@ extension RedisClient {
 
 // MARK: Intersect and Union
 
+/// The supported methods for aggregating results from the `zunionstore` or `zinterstore` commands in Redis.
+///
+/// For more information on these values, see
+/// [https://redis.io/commands/zunionstore](https://redis.io/commands/zunionstore)
+/// [https://redis.io/commands/zinterstore](https://redis.io/commands/zinterstore)
+public enum RedisSortedSetAggregateMethod: String {
+    /// Add the score of all matching elements in the source SortedSets.
+    case sum = "SUM"
+    /// Use the minimum score of the matching elements in the source SortedSets.
+    case min = "MIN"
+    /// Use the maximum score of the matching elements in the source SortedSets.
+    case max = "MAX"
+}
+
 extension RedisClient {
     /// Calculates the union of two or more sorted sets and stores the result.
     /// - Note: This operation overwrites any value stored at the destination key.
@@ -499,14 +522,14 @@ extension RedisClient {
     ///     - destination: The key of the new sorted set from the result.
     ///     - sources: The list of sorted set keys to treat as the source of the union.
     ///     - weights: The multiplying factor to apply to the corresponding `sources` key based on index of the two parameters.
-    ///     - aggregateMethod: The method of aggregating the values of the union. Supported values are "SUM", "MIN", and "MAX".
+    ///     - aggregateMethod: The method of aggregating the values of the union. If one isn't specified, Redis will default to `.sum`.
     /// - Returns: The number of elements in the new sorted set.
     @inlinable
     public func zunionstore(
         as destination: String,
         sources: [String],
         weights: [Int]? = nil,
-        aggregateMethod aggregate: String? = nil
+        aggregateMethod aggregate: RedisSortedSetAggregateMethod? = nil
     ) -> EventLoopFuture<Int> {
         return _zopstore(command: "ZUNIONSTORE", sources, destination, weights, aggregate)
     }
@@ -519,14 +542,14 @@ extension RedisClient {
     ///     - destination: The key of the new sorted set from the result.
     ///     - sources: The list of sorted set keys to treat as the source of the intersection.
     ///     - weights: The multiplying factor to apply to the corresponding `sources` key based on index of the two parameters.
-    ///     - aggregateMethod: The method of aggregating the values of the intersection. Supported values are "SUM", "MIN", and "MAX".
+    ///     - aggregateMethod: The method of aggregating the values of the intersection. If one isn't specified, Redis will default to `.sum`.
     /// - Returns: The number of elements in the new sorted set.
     @inlinable
     public func zinterstore(
         as destination: String,
         sources: [String],
         weights: [Int]? = nil,
-        aggregateMethod aggregate: String? = nil
+        aggregateMethod aggregate: RedisSortedSetAggregateMethod? = nil
     ) -> EventLoopFuture<Int> {
         return _zopstore(command: "ZINTERSTORE", sources, destination, weights, aggregate)
     }
@@ -537,8 +560,8 @@ extension RedisClient {
         _ sources: [String],
         _ destination: String,
         _ weights: [Int]?,
-        _ aggregate: String?) -> EventLoopFuture<Int>
-    {
+        _ aggregate: RedisSortedSetAggregateMethod?
+    ) -> EventLoopFuture<Int> {
         assert(sources.count > 0, "At least 1 source key should be provided.")
 
         var args: [RESPValue] = [
@@ -556,10 +579,8 @@ extension RedisClient {
         }
 
         if let a = aggregate {
-            assert(a == "SUM" || a == "MIN" || a == "MAX", "Aggregate method provided is unsupported.")
-
             args.append(.init(bulk: "AGGREGATE"))
-            args.append(.init(bulk: a))
+            args.append(.init(bulk: a.rawValue))
         }
 
         return send(command: command, with: args)

Sources/RedisNIO/Extensions/StandardLibrary.swift

@@ -30,7 +30,7 @@ extension Array where Element == RESPValue {
     /// Adds the elements of a collection to this array, delegating the details of how they are added to the given closure.
     ///
     /// When your closure will be doing more than a simple transform of the element value, such as when you're adding both the key _and_ value from a `KeyValuePair`,
-    /// you should set the `overestimatedCountBeingAdded` to a value you do not expect to reach in order to prevent multiple allocations from the increasing
+    /// you should set the `overestimatedCountBeingAdded` to a value you do not expect to exceed in order to prevent multiple allocations from the increasing
     /// element count.
     ///
     /// For example:

Tests/RedisNIOTests/Commands/SortedSetCommandsTests.swift

@@ -47,20 +47,22 @@ final class SortedSetCommandsTests: XCTestCase {
     func test_zadd() throws {
         _ = try connection.send(command: "FLUSHALL").wait()
 
-        XCTAssertThrowsError(try connection.zadd([(30, 2)], to: #function, options: ["INCR"]).wait())
-
         var count = try connection.zadd([(30, 2)], to: #function).wait()
         XCTAssertEqual(count, 1)
         count = try connection.zadd([(30, 5)], to: #function).wait()
         XCTAssertEqual(count, 0)
-        count = try connection.zadd([(30, 6), (31, 0), (32, 1)], to: #function, options: ["NX"]).wait()
+        count = try connection.zadd([(30, 6), (31, 0), (32, 1)], to: #function, options: [.onlyAddNewElements]).wait()
         XCTAssertEqual(count, 2)
-        count = try connection.zadd([(32, 2), (33, 3)], to: #function, options: ["XX", "CH"]).wait()
+        count = try connection.zadd(
+            [(32, 2), (33, 3)],
+            to: #function,
+            options: [.onlyUpdateExistingElements, .returnChangedCount]
+        ).wait()
         XCTAssertEqual(count, 1)
 
-        var success = try connection.zadd((30, 7), to: #function, options: ["CH"]).wait()
+        var success = try connection.zadd((30, 7), to: #function, options: [.returnChangedCount]).wait()
         XCTAssertTrue(success)
-        success = try connection.zadd((30, 8), to: #function, options: ["NX"]).wait()
+        success = try connection.zadd((30, 8), to: #function, options: [.onlyAddNewElements]).wait()
         XCTAssertFalse(success)
     }
 
@@ -227,7 +229,7 @@ final class SortedSetCommandsTests: XCTestCase {
             as: #function+#file,
             sources: [key, #function, #file],
             weights: [3, 2, 1],
-            aggregateMethod: "MAX"
+            aggregateMethod: .max
         ).wait()
         XCTAssertEqual(unionCount, 10)
         let rank = try connection.zrank(of: 10, in: #function+#file).wait()
@@ -243,7 +245,7 @@ final class SortedSetCommandsTests: XCTestCase {
             as: #file,
             sources: [key, #function],
             weights: [3, 2],
-            aggregateMethod: "MIN"
+            aggregateMethod: .min
         ).wait()
         XCTAssertEqual(unionCount, 2)
         let rank = try connection.zrank(of: 10, in: #file).wait()