Implemented IFEQ set command in python with tests (valkey-io#2962)

* Define new function, change paramter type to Union, add conditionals for IFEQ Signed-off-by: Angraybill <[email protected]> * Tests for IFEQ, positive and negative cases Signed-off-by: Angraybill <[email protected]> * Update Changelog Signed-off-by: Angraybill <[email protected]> * Import OnlyIfEqual in init file Signed-off-by: Angraybill <[email protected]> --------- Signed-off-by: Angraybill <[email protected]>
Maayanshani25 · Jan 22, 2025 · 70584d5 · 70584d5
1 parent c6e3b32
commit 70584d5
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,7 +19,7 @@
 * Go: Add `BZPopMin` ([#2849](https://github.com/valkey-io/valkey-glide/pull/2849))
 * Java: Shadow `protobuf` dependency ([#2931](https://github.com/valkey-io/valkey-glide/pull/2931))
 * Java: Add `RESP2` support ([#2383](https://github.com/valkey-io/valkey-glide/pull/2383))
-* Node: Add `IFEQ` option ([#2909](https://github.com/valkey-io/valkey-glide/pull/2909))
+* Node, Python: Add `IFEQ` option ([#2909](https://github.com/valkey-io/valkey-glide/pull/2909), [#2962](https://github.com/valkey-io/valkey-glide/pull/2962))
 
 #### Breaking Changes
 

diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py
@@ -30,6 +30,7 @@
     FunctionRestorePolicy,
     InfoSection,
     InsertPosition,
+    OnlyIfEqual,
     UpdateOptions,
 )
 from glide.async_commands.server_modules import ft, glide_json, json_transaction
@@ -225,6 +226,7 @@
     "Script",
     "ScoreBoundary",
     "ConditionalChange",
+    "OnlyIfEqual",
     "ExpireOptions",
     "ExpiryGetEx",
     "ExpirySet",

diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py
@@ -79,6 +79,19 @@ class ConditionalChange(Enum):
     ONLY_IF_DOES_NOT_EXIST = "NX"
 
 
+@dataclass
+class OnlyIfEqual:
+    """
+    Change condition to the `SET` command,
+    For additional conditonal options see ConditionalChange
+    - comparison_value - value to compare to the current value of a key.
+    If comparison_value is equal to the key, it will overwrite the value of key to the new provided value
+    Equivalent to the IFEQ comparison-value in the Valkey API
+    """
+
+    comparison_value: TEncodable
+
+
 class ExpiryType(Enum):
     """SET option: The type of the expiry.
     - SEC - Set the specified expire time, in seconds. Equivalent to `EX` in the Valkey API.
@@ -435,7 +448,7 @@ async def set(
         self,
         key: TEncodable,
         value: TEncodable,
-        conditional_set: Optional[ConditionalChange] = None,
+        conditional_set: Optional[Union[ConditionalChange, OnlyIfEqual]] = None,
         expiry: Optional[ExpirySet] = None,
         return_old_value: bool = False,
     ) -> Optional[bytes]:
@@ -447,7 +460,7 @@ async def set(
             key (TEncodable): the key to store.
             value (TEncodable): the value to store with the given key.
             conditional_set (Optional[ConditionalChange], optional): set the key only if the given condition is met.
-                Equivalent to [`XX` | `NX`] in the Valkey API. Defaults to None.
+                Equivalent to [`XX` | `NX` | `IFEQ` comparison-value] in the Valkey API. Defaults to None.
             expiry (Optional[ExpirySet], optional): set expiriation to the given key.
                 Equivalent to [`EX` | `PX` | `EXAT` | `PXAT` | `KEEPTTL`] in the Valkey API. Defaults to None.
             return_old_value (bool, optional): Return the old value stored at key, or None if key did not exist.
@@ -463,16 +476,38 @@ async def set(
         Example:
             >>> await client.set(b"key", b"value")
                 'OK'
-            >>> await client.set("key", "new_value",conditional_set=ConditionalChange.ONLY_IF_EXISTS, expiry=Expiry(ExpiryType.SEC, 5))
+
+                # ONLY_IF_EXISTS -> Only set the key if it already exists
+                # expiry -> Set the amount of time until key expires
+            >>> await client.set("key", "new_value",conditional_set=ConditionalChange.ONLY_IF_EXISTS, expiry=ExpirySet(ExpiryType.SEC, 5))
                 'OK' # Set "new_value" to "key" only if "key" already exists, and set the key expiration to 5 seconds.
+
+                # ONLY_IF_DOES_NOT_EXIST -> Only set key if it does not already exist
             >>> await client.set("key", "value", conditional_set=ConditionalChange.ONLY_IF_DOES_NOT_EXIST,return_old_value=True)
                 b'new_value' # Returns the old value of "key".
             >>> await client.get("key")
                 b'new_value' # Value wasn't modified back to being "value" because of "NX" flag.
+
+
+                # ONLY_IF_EQUAL -> Only set key if provided value is equal to current value of the key
+            >>> await client.set("key", "value")
+                'OK' # Reset "key" to "value"
+            >>> await client.set("key", "new_value", conditional_set=OnlyIfEqual("different_value")
+                'None' # Did not rewrite value of "key" because provided value was not equal to the previous value of "key"
+            >>> await client.get("key")
+                b'value' # Still the original value because nothing got rewritten in the last call
+            >>> await client.set("key", "new_value", conditional_set=OnlyIfEqual("value")
+                'OK'
+            >>> await client.get("key")
+                b'newest_value" # Set "key" to "new_value" because the provided value was equal to the previous value of "key"
         """
         args = [key, value]
-        if conditional_set:
+        if isinstance(conditional_set, ConditionalChange):
             args.append(conditional_set.value)
+
+        elif isinstance(conditional_set, OnlyIfEqual):
+            args.extend(["IFEQ", conditional_set.comparison_value])
+
         if return_old_value:
             args.append("GET")
         if expiry is not None:

diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py
@@ -39,6 +39,7 @@
     InfBound,
     InfoSection,
     InsertPosition,
+    OnlyIfEqual,
     UpdateOptions,
 )
 from glide.async_commands.sorted_set import (
@@ -452,6 +453,7 @@ async def test_inflight_request_limit(
     async def test_conditional_set(self, glide_client: TGlideClient):
         key = get_random_string(10)
         value = get_random_string(10)
+
         res = await glide_client.set(
             key, value, conditional_set=ConditionalChange.ONLY_IF_EXISTS
         )
@@ -466,6 +468,29 @@ async def test_conditional_set(self, glide_client: TGlideClient):
         )
         assert res is None
         assert await glide_client.get(key) == value.encode()
+        # Tests for ONLY_IF_EQUAL below in test_set_only_if_equal()
+
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    @pytest.mark.skip_if_version_below("8.1.0")
+    async def test_set_only_if_equal(self, glide_client: TGlideClient):
+        key = get_random_string(10)
+        value = get_random_string(10)
+        value2 = get_random_string(10)
+        wrong_comparison_value = get_random_string(10)
+        while wrong_comparison_value == value:
+            wrong_comparison_value = get_random_string(10)
+
+        await glide_client.set(key, value)
+
+        res = await glide_client.set(
+            key, "foobar", conditional_set=OnlyIfEqual(wrong_comparison_value)
+        )
+        assert res is None
+        assert await glide_client.get(key) == value.encode()
+        res = await glide_client.set(key, value2, conditional_set=OnlyIfEqual(value))
+        assert res == OK
+        assert await glide_client.get(key) == value2.encode()
 
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])