Merge pull request #246 from jogehl/main

Add HSETEX

Author: bogdanp05

tests/test_asyncio/test_commands.py

@@ -2151,6 +2151,32 @@ async def test_hstrlen(self, r: valkey.Valkey):
         assert await r.hstrlen("a", "1") == 2
         assert await r.hstrlen("a", "2") == 3
 
+    @skip_if_server_version_lt("9.0.0")
+    async def test_hsetex(self, r):
+        assert await r.hsetex("a", "field1", "value1", ex=5) == 1
+        assert await r.hget("a", "field1") == b"value1"
+        assert await r.hsetex("a", "field1", "value2", ex=5) == 1
+        assert await r.hget("a", "field1") == b"value2"
+
+    @skip_if_server_version_lt("9.0.0")
+    async def test_hsetex_invalid_params(self, r):
+        with pytest.raises(exceptions.DataError):
+            await r.hsetex("a", "field1", "value1", ex=5, px=5000)
+
+    @skip_if_server_version_lt("9.0.0")
+    async def test_hsetex_px(self, r):
+        assert await r.hsetex("a", "field1", "value1", px=5000) == 1
+        assert await r.hget("a", "field1") == b"value1"
+        assert await r.hsetex("a", "field1", "value2", px=5000) == 1
+        assert await r.hget("a", "field1") == b"value2"
+
+    @skip_if_server_version_lt("9.0.0")
+    async def test_hsetex_mapping(self, r):
+        mapping = {"field1": "value1", "field2": "value2"}
+        assert await r.hsetex("a", mapping=mapping, ex=5) == 1
+        assert await r.hget("a", "field1") == b"value1"
+        assert await r.hget("a", "field2") == b"value2"
+
     # SORT
     async def test_sort_basic(self, r: valkey.Valkey):
         await r.rpush("a", "3", "2", "1", "4")

tests/test_commands.py

@@ -3236,6 +3236,32 @@ def test_hstrlen(self, r):
         assert r.hstrlen("a", "1") == 2
         assert r.hstrlen("a", "2") == 3
 
+    @skip_if_server_version_lt("9.0.0")
+    def test_hsetex(self, r):
+        assert r.hsetex("a", "field1", "value1", ex=5) == 1
+        assert r.hget("a", "field1") == b"value1"
+        assert r.hsetex("a", "field1", "value2", ex=5) == 1
+        assert r.hget("a", "field1") == b"value2"
+
+    @skip_if_server_version_lt("9.0.0")
+    def test_hsetex_invalid_params(self, r):
+        with pytest.raises(exceptions.DataError):
+            r.hsetex("a", "field1", "value1", ex=5, px=5000)  # Both ex and px provided
+
+    @skip_if_server_version_lt("9.0.0")
+    def test_hsetex_px(self, r):
+        assert r.hsetex("a", "field1", "value1", px=5000) == 1
+        assert r.hget("a", "field1") == b"value1"
+        assert r.hsetex("a", "field1", "value2", px=5000) == 1
+        assert r.hget("a", "field1") == b"value2"
+
+    @skip_if_server_version_lt("9.0.0")
+    def test_hsetex_mapping(self, r):
+        mapping = {"field1": "value1", "field2": "value2"}
+        assert r.hsetex("a", mapping=mapping, ex=5) == 1
+        assert r.hget("a", "field1") == b"value1"
+        assert r.hget("a", "field2") == b"value2"
+
     # SORT
     def test_sort_basic(self, r):
         r.rpush("a", "3", "2", "1", "4")

valkey/commands/core.py

@@ -5055,6 +5055,82 @@ def hsetnx(self, name: str, key: str, value: str) -> Union[Awaitable[bool], bool
         """
         return self.execute_command("HSETNX", name, key, value)
 
+    def hsetex(
+        self,
+        name: str,
+        key: Optional[str] = None,
+        value: Optional[str] = None,
+        mapping: Optional[dict] = None,
+        items: Optional[list] = None,
+        ex: Union[ExpiryT, None] = None,
+        px: Union[ExpiryT, None] = None,
+        exat: Union[AbsExpiryT, None] = None,
+        pxat: Union[AbsExpiryT, None] = None,
+        keepttl: bool = False,
+        nx: bool = False,
+        xx: bool = False,
+        fnx: bool = False,
+        fxx: bool = False,
+    ) -> Union[Awaitable[bool], bool]:
+        """
+        Set key to value within hash ``name``,
+        ``mapping`` accepts a dict of key/value pairs to be added to hash ``name``.
+        ``items`` accepts a list of key/value pairs to be added to hash ``name``.
+        Set expiration options for the hash fields.
+
+        For more information see https://valkey.io/commands/hsetex
+        """
+
+        if key is None and not mapping and not items:
+            raise DataError("'hsetex' with no key value pairs")
+        if int(key is not None) + int(mapping is not None) + int(items is not None) > 1:
+            raise DataError(
+                "Only one of 'key/value', 'mapping', or 'items' can be specified."
+            )
+
+        if int(keepttl) + sum(arg is not None for arg in [ex, px, exat, pxat]) > 1:
+            raise DataError(
+                "Only one of 'ex', 'px', 'exat', 'pxat', or 'keepttl' can be specified."
+            )
+        if nx and xx:
+            raise DataError("Only one of 'nx' or 'xx' can be specified.")
+        if fnx and fxx:
+            raise DataError("Only one of 'fnx' or 'fxx' can be specified.")
+        pieces = []
+        if ex is not None:
+            pieces.extend(["EX", ex])
+        if px is not None:
+            pieces.extend(["PX", px])
+        if exat is not None:
+            pieces.extend(["EXAT", exat])
+        if pxat is not None:
+            pieces.extend(["PXAT", pxat])
+        if nx:
+            pieces.append("NX")
+        if xx:
+            pieces.append("XX")
+        if fnx:
+            pieces.append("FNX")
+        if fxx:
+            pieces.append("FXX")
+        pieces.append("FIELDS")
+        if key is not None and value is not None:
+            pieces.append(1)  # for one field
+            pieces.append(key)
+            pieces.append(value)
+        if mapping:
+            pieces.append(len(mapping))
+            for key, value in mapping.items():
+                if key is None or value is None:
+                    raise DataError("'hsetex' mapping contains None key or value")
+                pieces.append(key)
+                pieces.append(value)
+        if items:
+            pieces.append(len(items) // 2)
+            pieces.extend(items)
+
+        return self.execute_command("HSETEX", name, *pieces)
+
     def hmset(self, name: str, mapping: dict) -> Union[Awaitable[str], str]:
         """
         Set key to value within hash ``name`` for each corresponding