From c4364c5e708647c68a880385ebc5f0b4a0f4d60e Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 17 Jun 2024 17:53:19 -0700 Subject: [PATCH] Python: add BITOP command --- CHANGELOG.md | 1 + python/python/glide/__init__.py | 2 +- python/python/glide/async_commands/bitmap.py | 12 +++ python/python/glide/async_commands/core.py | 37 +++++++- .../glide/async_commands/transaction.py | 26 +++++- python/python/tests/test_async_client.py | 88 ++++++++++++++++++- python/python/tests/test_transaction.py | 9 +- 7 files changed, 170 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a6d92db95d..945035a090 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ * Python: Added SETBIT command ([#1571](https://github.com/aws/glide-for-redis/pull/1571)) * Python: Added GETBIT command ([#1575](https://github.com/aws/glide-for-redis/pull/1575)) * Python: Added BITCOUNT command ([#1592](https://github.com/aws/glide-for-redis/pull/1592)) +* Python: Added BITOP command ([#1596](https://github.com/aws/glide-for-redis/pull/1596)) ### Breaking Changes * Node: Update XREAD to return a Map of Map ([#1494](https://github.com/aws/glide-for-redis/pull/1494)) diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py index 65bb5fac80..18d52b8d27 100644 --- a/python/python/glide/__init__.py +++ b/python/python/glide/__init__.py @@ -1,6 +1,6 @@ # Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0 -from glide.async_commands.bitmap import BitmapIndexType, OffsetOptions +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 ( ConditionalChange, diff --git a/python/python/glide/async_commands/bitmap.py b/python/python/glide/async_commands/bitmap.py index 03c5966251..8b073bd2a1 100644 --- a/python/python/glide/async_commands/bitmap.py +++ b/python/python/glide/async_commands/bitmap.py @@ -48,3 +48,15 @@ def to_args(self) -> List[str]: args.append(self.index_type.value) return args + + +class BitwiseOperation(Enum): + """ + Enumeration defining the bitwise operation to use in the `BITOP` command. Specifies the bitwise operation to + perform between the passed in keys. + """ + + AND = "AND" + OR = "OR" + XOR = "XOR" + NOT = "NOT" diff --git a/python/python/glide/async_commands/core.py b/python/python/glide/async_commands/core.py index 5625cbf3d4..5d8fbc7149 100644 --- a/python/python/glide/async_commands/core.py +++ b/python/python/glide/async_commands/core.py @@ -16,7 +16,7 @@ get_args, ) -from glide.async_commands.bitmap import OffsetOptions +from glide.async_commands.bitmap import BitwiseOperation, OffsetOptions from glide.async_commands.command_args import Limit, ListDirection, OrderBy from glide.async_commands.sorted_set import ( AggregationType, @@ -4202,6 +4202,41 @@ async def getbit(self, key: str, offset: int) -> int: await self._execute_command(RequestType.GetBit, [key, str(offset)]), ) + async def bitop( + self, operation: BitwiseOperation, destination: str, keys: List[str] + ) -> int: + """ + Perform a bitwise operation between multiple keys (containing string values) and store the result in the + `destination`. + + See https://valkey.io/commands/bitop for more details. + + Note: + When in cluster mode, `destination` and all `keys` must map to the same hash slot. + + Args: + operation (BitwiseOperation): The bitwise operation to perform. + destination (str): The key that will store the resulting string. + keys (List[str]): The list of keys to perform the bitwise operation on. + + Returns: + int: The size of the string stored in `destination`. + + Examples: + >>> await client.set("key1", "A") # "A" has binary value 01000001 + >>> await client.set("key1", "B") # "B" has binary value 01000010 + >>> await client.bitop(BitwiseOperation.AND, "destination", ["key1", "key2"]) + 1 # The size of the resulting string stored in "destination" is 1 + >>> await client.get("destination") + "@" # "@" has binary value 01000000 + """ + return cast( + int, + await self._execute_command( + RequestType.BitOp, [operation.value, destination] + keys + ), + ) + async def object_encoding(self, key: str) -> Optional[str]: """ Returns the internal encoding for the Redis object stored at `key`. diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py index d8e451f505..4bc0662be9 100644 --- a/python/python/glide/async_commands/transaction.py +++ b/python/python/glide/async_commands/transaction.py @@ -3,7 +3,7 @@ import threading from typing import List, Mapping, Optional, Tuple, TypeVar, Union -from glide.async_commands.bitmap import OffsetOptions +from glide.async_commands.bitmap import BitwiseOperation, OffsetOptions from glide.async_commands.command_args import Limit, ListDirection, OrderBy from glide.async_commands.core import ( ConditionalChange, @@ -2915,6 +2915,30 @@ def getbit(self: TTransaction, key: str, offset: int) -> TTransaction: """ return self.append_command(RequestType.GetBit, [key, str(offset)]) + def bitop( + self: TTransaction, + operation: BitwiseOperation, + destination: str, + keys: List[str], + ) -> TTransaction: + """ + Perform a bitwise operation between multiple keys (containing string values) and store the result in the + `destination`. + + See https://valkey.io/commands/bitop for more details. + + Args: + operation (BitwiseOperation): The bitwise operation to perform. + destination (str): The key that will store the resulting string. + keys (List[str]): The list of keys to perform the bitwise operation on. + + Command response: + int: The size of the string stored in `destination`. + """ + return self.append_command( + RequestType.BitOp, [operation.value, destination] + keys + ) + def object_encoding(self: TTransaction, key: str) -> TTransaction: """ Returns the internal encoding for the Redis object stored at `key`. diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py index 3c0be92527..2b8f14f55c 100644 --- a/python/python/tests/test_async_client.py +++ b/python/python/tests/test_async_client.py @@ -11,7 +11,7 @@ import pytest from glide import ClosingError, RequestError, Script -from glide.async_commands.bitmap import BitmapIndexType, OffsetOptions +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 ( ConditionalChange, @@ -4426,6 +4426,91 @@ async def test_getbit(self, redis_client: TRedisClient): with pytest.raises(RequestError): await redis_client.getbit(set_key, 0) + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_bitop(self, redis_client: TRedisClient): + key1 = f"{{testKey}}:1-{get_random_string(10)}" + key2 = f"{{testKey}}:2-{get_random_string(10)}" + keys = [key1, key2] + destination = f"{{testKey}}:3-{get_random_string(10)}" + non_existing_key1 = f"{{testKey}}:4-{get_random_string(10)}" + non_existing_key2 = f"{{testKey}}:5-{get_random_string(10)}" + non_existing_keys = [non_existing_key1, non_existing_key2] + set_key = f"{{testKey}}:6-{get_random_string(10)}" + value1 = "foobar" + value2 = "abcdef" + + assert await redis_client.set(key1, value1) == OK + assert await redis_client.set(key2, value2) == OK + assert await redis_client.bitop(BitwiseOperation.AND, destination, keys) == 6 + assert await redis_client.get(destination) == "`bc`ab" + assert await redis_client.bitop(BitwiseOperation.OR, destination, keys) == 6 + assert await redis_client.get(destination) == "goofev" + + # reset values for simplicity of results in XOR + assert await redis_client.set(key1, "a") == OK + assert await redis_client.set(key2, "b") == OK + assert await redis_client.bitop(BitwiseOperation.XOR, destination, keys) == 1 + assert await redis_client.get(destination) == "\u0003" + + # test single source key + assert await redis_client.bitop(BitwiseOperation.AND, destination, [key1]) == 1 + assert await redis_client.get(destination) == "a" + assert await redis_client.bitop(BitwiseOperation.OR, destination, [key1]) == 1 + assert await redis_client.get(destination) == "a" + assert await redis_client.bitop(BitwiseOperation.XOR, destination, [key1]) == 1 + assert await redis_client.get(destination) == "a" + assert await redis_client.bitop(BitwiseOperation.NOT, destination, [key1]) == 1 + # currently, attempting to get the value from destination after the above NOT incorrectly raises an error + # TODO: update with a GET call once fix is implemented for https://github.com/aws/glide-for-redis/issues/1447 + + assert await redis_client.setbit(key1, 0, 1) == 0 + assert await redis_client.bitop(BitwiseOperation.NOT, destination, [key1]) == 1 + assert await redis_client.get(destination) == "\u001e" + + # stores None when all keys hold empty strings + assert ( + await redis_client.bitop( + BitwiseOperation.AND, destination, non_existing_keys + ) + == 0 + ) + assert await redis_client.get(destination) is None + assert ( + await redis_client.bitop( + BitwiseOperation.OR, destination, non_existing_keys + ) + == 0 + ) + assert await redis_client.get(destination) is None + assert ( + await redis_client.bitop( + BitwiseOperation.XOR, destination, non_existing_keys + ) + == 0 + ) + assert await redis_client.get(destination) is None + assert ( + await redis_client.bitop( + BitwiseOperation.NOT, destination, [non_existing_key1] + ) + == 0 + ) + assert await redis_client.get(destination) is None + + # invalid argument - source key list cannot be empty + with pytest.raises(RequestError): + await redis_client.bitop(BitwiseOperation.OR, destination, []) + + # invalid arguments - NOT cannot be passed more than 1 key + with pytest.raises(RequestError): + await redis_client.bitop(BitwiseOperation.NOT, destination, [key1, key2]) + + assert await redis_client.sadd(set_key, [value1]) == 1 + # invalid argument - source key has the wrong type + with pytest.raises(RequestError): + await redis_client.bitop(BitwiseOperation.AND, destination, [set_key]) + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_object_encoding(self, redis_client: TRedisClient): @@ -4581,6 +4666,7 @@ async def test_multi_key_command_returns_cross_slot_error( "abc", "zxy", ListDirection.LEFT, ListDirection.LEFT, 1 ), redis_client.msetnx({"abc": "abc", "zxy": "zyx"}), + redis_client.bitop(BitwiseOperation.OR, "abc", ["zxy", "lkn"]), ] if not await check_if_server_version_lt(redis_client, "7.0.0"): diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py index 3b5e34023d..236dba235b 100644 --- a/python/python/tests/test_transaction.py +++ b/python/python/tests/test_transaction.py @@ -6,7 +6,7 @@ import pytest from glide import RequestError -from glide.async_commands.bitmap import BitmapIndexType, OffsetOptions +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 InsertPosition, StreamAddOptions, TrimByMinId from glide.async_commands.sorted_set import ( @@ -362,6 +362,13 @@ async def transaction_test( transaction.bitcount(key20, OffsetOptions(1, 1)) args.append(6) + transaction.set(key19, "abcdef") + args.append(OK) + transaction.bitop(BitwiseOperation.AND, key19, [key19, key20]) + args.append(6) + transaction.get(key19) + args.append("`bc`ab") + if not await check_if_server_version_lt(redis_client, "7.0.0"): transaction.bitcount(key20, OffsetOptions(5, 30, BitmapIndexType.BIT)) args.append(17)