Python: Add command GetEX (valkey-io#1612)

* Add Python cmd GetEX * Python: Add command GetEx * addressing comments & updating changelog * changelog & request type --------- Co-authored-by: TJ Zhang <[email protected]>
Bit-Quill · Jun 19, 2024 · 50e5ff3 · 50e5ff3
1 parent 53bbd2f
commit 50e5ff3
Show file tree

Hide file tree

Showing 8 changed files with 237 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -44,6 +44,7 @@
 * Python: Added TOUCH command ([#1582](https://github.com/aws/glide-for-redis/pull/1582))
 * Python: Added BITOP command ([#1596](https://github.com/aws/glide-for-redis/pull/1596))
 * Python: Added BITPOS command ([#1604](https://github.com/aws/glide-for-redis/pull/1604))
+* Python: Added GETEX command ([#1612](https://github.com/aws/glide-for-redis/pull/1612))
 
 ### Breaking Changes
 * Node: Update XREAD to return a Map of Map ([#1494](https://github.com/aws/glide-for-redis/pull/1494))

diff --git a/glide-core/src/protobuf/redis_request.proto b/glide-core/src/protobuf/redis_request.proto
@@ -230,6 +230,7 @@ enum RequestType {
     XGroupCreateConsumer = 189;
     XGroupDelConsumer = 190;
     RandomKey = 191;
+    GetEx = 192;
 }
 
 message Command {

diff --git a/glide-core/src/request_type.rs b/glide-core/src/request_type.rs
@@ -200,6 +200,7 @@ pub enum RequestType {
     XGroupCreateConsumer = 189,
     XGroupDelConsumer = 190,
     RandomKey = 191,
+    GetEx = 192,
 }
 
 fn get_two_word_command(first: &str, second: &str) -> Cmd {
@@ -403,6 +404,7 @@ impl From<::protobuf::EnumOrUnknown<ProtobufRequestType>> for RequestType {
             ProtobufRequestType::XGroupCreateConsumer => RequestType::XGroupCreateConsumer,
             ProtobufRequestType::XGroupDelConsumer => RequestType::XGroupDelConsumer,
             ProtobufRequestType::RandomKey => RequestType::RandomKey,
+            ProtobufRequestType::GetEx => RequestType::GetEx,
         }
     }
 }
@@ -604,6 +606,7 @@ impl RequestType {
             }
             RequestType::XGroupDelConsumer => Some(get_two_word_command("XGROUP", "DELCONSUMER")),
             RequestType::RandomKey => Some(cmd("RANDOMKEY")),
+            RequestType::GetEx => Some(cmd("GETEX")),
         }
     }
 }
diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py
@@ -5,8 +5,10 @@
 from glide.async_commands.core import (
     ConditionalChange,
     ExpireOptions,
+    ExpiryGetEx,
     ExpirySet,
     ExpiryType,
+    ExpiryTypeGetEx,
     FlushMode,
     InfoSection,
     InsertPosition,

diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py
@@ -71,6 +71,23 @@ class ExpiryType(Enum):
     KEEP_TTL = 4, Type[None]  # Equivalent to `KEEPTTL` in the Redis API
 
 
+class ExpiryTypeGetEx(Enum):
+    """GetEx option: The type of the expiry.
+    - EX - Set the specified expire time, in seconds. Equivalent to `EX` in the Redis API.
+    - PX - Set the specified expire time, in milliseconds. Equivalent to `PX` in the Redis API.
+    - UNIX_SEC - Set the specified Unix time at which the key will expire, in seconds. Equivalent to `EXAT` in the Redis API.
+    - UNIX_MILLSEC - Set the specified Unix time at which the key will expire, in milliseconds. Equivalent to `PXAT` in the
+        Redis API.
+    - PERSIST - Remove the time to live associated with the key. Equivalent to `PERSIST` in the Redis API.
+    """
+
+    SEC = 0, Union[int, timedelta]  # Equivalent to `EX` in the Redis API
+    MILLSEC = 1, Union[int, timedelta]  # Equivalent to `PX` in the Redis API
+    UNIX_SEC = 2, Union[int, datetime]  # Equivalent to `EXAT` in the Redis API
+    UNIX_MILLSEC = 3, Union[int, datetime]  # Equivalent to `PXAT` in the Redis API
+    PERSIST = 4, Type[None]  # Equivalent to `PERSIST` in the Redis API
+
+
 class InfoSection(Enum):
     """
     INFO option: a specific section of information:
@@ -324,6 +341,61 @@ def get_cmd_args(self) -> List[str]:
         return [self.cmd_arg] if self.value is None else [self.cmd_arg, self.value]
 
 
+class ExpiryGetEx:
+    """GetEx option: Represents the expiry type and value to be executed with "GetEx" command."""
+
+    def __init__(
+        self,
+        expiry_type: ExpiryTypeGetEx,
+        value: Optional[Union[int, datetime, timedelta]],
+    ) -> None:
+        """
+        Args:
+            - expiry_type (ExpiryType): The expiry type.
+            - value (Optional[Union[int, datetime, timedelta]]): The value of the expiration type. The type of expiration
+                determines the type of expiration value:
+                - SEC: Union[int, timedelta]
+                - MILLSEC: Union[int, timedelta]
+                - UNIX_SEC: Union[int, datetime]
+                - UNIX_MILLSEC: Union[int, datetime]
+                - PERSIST: Type[None]
+        """
+        self.set_expiry_type_and_value(expiry_type, value)
+
+    def set_expiry_type_and_value(
+        self,
+        expiry_type: ExpiryTypeGetEx,
+        value: Optional[Union[int, datetime, timedelta]],
+    ):
+        if not isinstance(value, get_args(expiry_type.value[1])):
+            raise ValueError(
+                f"The value of {expiry_type} should be of type {expiry_type.value[1]}"
+            )
+        self.expiry_type = expiry_type
+        if self.expiry_type == ExpiryTypeGetEx.SEC:
+            self.cmd_arg = "EX"
+            if isinstance(value, timedelta):
+                value = int(value.total_seconds())
+        elif self.expiry_type == ExpiryTypeGetEx.MILLSEC:
+            self.cmd_arg = "PX"
+            if isinstance(value, timedelta):
+                value = int(value.total_seconds() * 1000)
+        elif self.expiry_type == ExpiryTypeGetEx.UNIX_SEC:
+            self.cmd_arg = "EXAT"
+            if isinstance(value, datetime):
+                value = int(value.timestamp())
+        elif self.expiry_type == ExpiryTypeGetEx.UNIX_MILLSEC:
+            self.cmd_arg = "PXAT"
+            if isinstance(value, datetime):
+                value = int(value.timestamp() * 1000)
+        elif self.expiry_type == ExpiryTypeGetEx.PERSIST:
+            self.cmd_arg = "PERSIST"
+        self.value = str(value) if value else None
+
+    def get_cmd_args(self) -> List[str]:
+        return [self.cmd_arg] if self.value is None else [self.cmd_arg, self.value]
+
+
 class InsertPosition(Enum):
     BEFORE = "BEFORE"
     AFTER = "AFTER"
@@ -4761,3 +4833,43 @@ async def srandmember_count(self, key: str, count: int) -> List[str]:
             List[str],
             await self._execute_command(RequestType.SRandMember, [key, str(count)]),
         )
+
+    async def getex(
+        self,
+        key: str,
+        expiry: Optional[ExpiryGetEx] = None,
+    ) -> Optional[str]:
+        """
+        Get the value of `key` and optionally set its expiration. `GETEX` is similar to `GET`.
+        See https://valkey.io/commands/getex for more details.
+
+        Args:
+            key (str): The key to get.
+            expiry (Optional[ExpirySet], optional): set expiriation to the given key.
+                Equivalent to [`EX` | `PX` | `EXAT` | `PXAT` | `PERSIST`] in the Redis API.
+
+        Returns:
+            Optional[str]:
+                If `key` exists, return the value stored at `key`
+                If `key` does not exist, return `None`
+
+        Examples:
+            >>> await client.set("key", "value")
+                'OK'
+            >>> await client.getex("key")
+                'value'
+            >>> await client.getex("key", ExpiryGetEx(ExpiryTypeGetEx.SEC, 1))
+                'value'
+            >>> time.sleep(1)
+            >>> await client.getex("key")
+                None
+
+        Since: Redis version 6.2.0.
+        """
+        args = [key]
+        if expiry is not None:
+            args.extend(expiry.get_cmd_args())
+        return cast(
+            Optional[str],
+            await self._execute_command(RequestType.GetEx, args),
+        )
diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py
@@ -8,6 +8,7 @@
 from glide.async_commands.core import (
     ConditionalChange,
     ExpireOptions,
+    ExpiryGetEx,
     ExpirySet,
     FlushMode,
     GeospatialData,
@@ -3304,6 +3305,30 @@ def flushall(
             args.append(flush_mode.value)
         return self.append_command(RequestType.FlushAll, args)
 
+    def getex(
+        self: TTransaction, key: str, expiry: Optional[ExpiryGetEx] = None
+    ) -> TTransaction:
+        """
+        Get the value of `key` and optionally set its expiration. GETEX is similar to GET.
+        See https://valkey.io/commands/getex for more details.
+
+        Args:
+            key (str): The key to get.
+            expiry (Optional[ExpirySet], optional): set expiriation to the given key.
+                Equivalent to [`EX` | `PX` | `EXAT` | `PXAT` | `PERSIST`] in the Redis API.
+
+        Command Response:
+            Optional[str]:
+                If `key` exists, return the value stored at `key`
+                If 'key` does not exist, return 'None'
+
+        Since: Redis version 6.2.0.
+        """
+        args = [key]
+        if expiry is not None:
+            args.extend(expiry.get_cmd_args())
+        return self.append_command(RequestType.GetEx, args)
+
 
 class Transaction(BaseTransaction):
     """

diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py
@@ -16,8 +16,10 @@
 from glide.async_commands.core import (
     ConditionalChange,
     ExpireOptions,
+    ExpiryGetEx,
     ExpirySet,
     ExpiryType,
+    ExpiryTypeGetEx,
     FlushMode,
     InfBound,
     InfoSection,
@@ -5234,6 +5236,41 @@ async def test_flushall(self, redis_client: TRedisClient):
                 assert await redis_client.flushall(FlushMode.SYNC, AllPrimaries()) is OK
             assert await redis_client.dbsize() == 0
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_getex(self, redis_client: TRedisClient):
+        min_version = "6.2.0"
+        if await check_if_server_version_lt(redis_client, min_version):
+            return pytest.mark.skip(reason=f"Redis version required >= {min_version}")
+
+        key1 = get_random_string(10)
+        non_existing_key = get_random_string(10)
+        value = get_random_string(10)
+
+        assert await redis_client.set(key1, value) == OK
+        assert await redis_client.getex(non_existing_key) is None
+        assert await redis_client.getex(key1) == value
+        assert await redis_client.ttl(key1) == -1
+
+        # setting expiration timer
+        assert (
+            await redis_client.getex(key1, ExpiryGetEx(ExpiryTypeGetEx.MILLSEC, 50))
+            == value
+        )
+        assert await redis_client.ttl(key1) != -1
+
+        # setting and clearing expiration timer
+        assert await redis_client.set(key1, value) == OK
+        assert (
+            await redis_client.getex(key1, ExpiryGetEx(ExpiryTypeGetEx.SEC, 10))
+            == value
+        )
+        assert (
+            await redis_client.getex(key1, ExpiryGetEx(ExpiryTypeGetEx.PERSIST, None))
+            == value
+        )
+        assert await redis_client.ttl(key1) == -1
+
 
 class TestMultiKeyCommandCrossSlot:
     @pytest.mark.parametrize("cluster_mode", [True])
@@ -5356,6 +5393,50 @@ def test_expiry_cmd_args(self):
         )
         assert exp_unix_millisec_datetime.get_cmd_args() == ["PXAT", "1682639759342"]
 
+    def test_get_expiry_cmd_args(self):
+        exp_sec = ExpiryGetEx(ExpiryTypeGetEx.SEC, 5)
+        assert exp_sec.get_cmd_args() == ["EX", "5"]
+
+        exp_sec_timedelta = ExpiryGetEx(ExpiryTypeGetEx.SEC, timedelta(seconds=5))
+        assert exp_sec_timedelta.get_cmd_args() == ["EX", "5"]
+
+        exp_millsec = ExpiryGetEx(ExpiryTypeGetEx.MILLSEC, 5)
+        assert exp_millsec.get_cmd_args() == ["PX", "5"]
+
+        exp_millsec_timedelta = ExpiryGetEx(
+            ExpiryTypeGetEx.MILLSEC, timedelta(seconds=5)
+        )
+        assert exp_millsec_timedelta.get_cmd_args() == ["PX", "5000"]
+
+        exp_millsec_timedelta = ExpiryGetEx(
+            ExpiryTypeGetEx.MILLSEC, timedelta(seconds=5)
+        )
+        assert exp_millsec_timedelta.get_cmd_args() == ["PX", "5000"]
+
+        exp_unix_sec = ExpiryGetEx(ExpiryTypeGetEx.UNIX_SEC, 1682575739)
+        assert exp_unix_sec.get_cmd_args() == ["EXAT", "1682575739"]
+
+        exp_unix_sec_datetime = ExpiryGetEx(
+            ExpiryTypeGetEx.UNIX_SEC,
+            datetime(2023, 4, 27, 23, 55, 59, 342380, timezone.utc),
+        )
+        assert exp_unix_sec_datetime.get_cmd_args() == ["EXAT", "1682639759"]
+
+        exp_unix_millisec = ExpiryGetEx(ExpiryTypeGetEx.UNIX_MILLSEC, 1682586559964)
+        assert exp_unix_millisec.get_cmd_args() == ["PXAT", "1682586559964"]
+
+        exp_unix_millisec_datetime = ExpiryGetEx(
+            ExpiryTypeGetEx.UNIX_MILLSEC,
+            datetime(2023, 4, 27, 23, 55, 59, 342380, timezone.utc),
+        )
+        assert exp_unix_millisec_datetime.get_cmd_args() == ["PXAT", "1682639759342"]
+
+        exp_persist = ExpiryGetEx(
+            ExpiryTypeGetEx.PERSIST,
+            None,
+        )
+        assert exp_persist.get_cmd_args() == ["PERSIST"]
+
     def test_expiry_raises_on_value_error(self):
         with pytest.raises(ValueError):
             ExpirySet(ExpiryType.SEC, 5.5)

diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py
@@ -9,6 +9,8 @@
 from glide.async_commands.bitmap import BitmapIndexType, BitwiseOperation, OffsetOptions
 from glide.async_commands.command_args import Limit, ListDirection, OrderBy
 from glide.async_commands.core import (
+    ExpiryGetEx,
+    ExpiryTypeGetEx,
     FlushMode,
     InsertPosition,
     StreamAddOptions,
@@ -64,6 +66,7 @@ async def transaction_test(
     key18 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # sort
     key19 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # bitmap
     key20 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # bitmap
+    key22 = "{{{}}}:{}".format(keyslot, get_random_string(3))  # getex
 
     value = datetime.now(timezone.utc).strftime("%m/%d/%Y, %H:%M:%S")
     value2 = get_random_string(5)
@@ -480,6 +483,15 @@ async def transaction_test(
         transaction.flushall(FlushMode.SYNC)
         args.append(OK)
 
+    min_version = "6.2.0"
+    if not await check_if_server_version_lt(redis_client, min_version):
+        transaction.set(key22, "value")
+        args.append(OK)
+        transaction.getex(key22)
+        args.append("value")
+        transaction.getex(key22, ExpiryGetEx(ExpiryTypeGetEx.SEC, 1))
+        args.append("value")
+
     min_version = "7.0.0"
     if not await check_if_server_version_lt(redis_client, min_version):
         transaction.zadd(key16, {"a": 1, "b": 2, "c": 3, "d": 4})