diff --git a/python/python/glide/__init__.py b/python/python/glide/__init__.py index cf817c128a..27a43f0a7f 100644 --- a/python/python/glide/__init__.py +++ b/python/python/glide/__init__.py @@ -33,6 +33,7 @@ UpdateOptions, ) from glide.async_commands.server_modules import json +from glide.async_commands.server_modules.json import JsonArrPopOptions, JsonGetOptions from glide.async_commands.sorted_set import ( AggregationType, GeoSearchByBox, @@ -184,7 +185,6 @@ "InfBound", "InfoSection", "InsertPosition", - "json", "LexBoundary", "Limit", "ListDirection", @@ -212,6 +212,10 @@ "ClusterScanCursor" # PubSub "PubSubMsg", + # Json + "json", + "JsonGetOptions", + "JsonArrPopOptions", # Logger "Logger", "LogLevel", diff --git a/python/python/glide/async_commands/server_modules/json.py b/python/python/glide/async_commands/server_modules/json.py index d75c2c7a4b..755c90a58e 100644 --- a/python/python/glide/async_commands/server_modules/json.py +++ b/python/python/glide/async_commands/server_modules/json.py @@ -54,6 +54,36 @@ def get_options(self) -> List[str]: return args +from typing import List, Optional, Union + + +class JsonArrPopOptions: + """ + Options for the JSON.ARRPOP command. + + Args: + path (TEncodable): The JSON path within the document from which to pop an element. + index (Optional[int]): The index of the element to pop. If not specified, will pop the last element. + Out of boundary indexes are rounded to their respective array boundaries. Defaults to None. + """ + + def __init__(self, path: TEncodable, index: Optional[int] = None): + self.path = path + self.index = index + + def get_options(self) -> List[TEncodable]: + """ + Get the options as a list of arguments for the JSON.ARRPOP command. + + Returns: + List[TEncodable]: A list containing the path and, if specified, the index. + """ + args = [self.path] + if self.index is not None: + args.append(str(self.index)) + return args + + async def set( client: TGlideClient, key: TEncodable, @@ -139,6 +169,66 @@ async def get( return cast(bytes, await client.custom_command(args)) +async def arrpop( + client: TGlideClient, + key: TEncodable, + options: Optional[JsonArrPopOptions] = None, +) -> Optional[TJsonResponse[bytes]]: + """ + Pops the last element from the array located at the specified path within the JSON document stored at `key`. + If `options.index` is provided, it pops the element at that index instead of the last element. + + See https://valkey.io/commands/json.arrpop/ for more details. + + Args: + client (TGlideClient): The client to execute the command. + key (TEncodable): The key of the JSON document. + options (Optional[JsonArrPopOptions]): Options including the path and optional index. See `JsonArrPopOptions`. + + Returns: + Optional[TJsonResponse[bytes]]: + For JSONPath (`path` starts with `$`): + Returns a list of bytes string replies for every possible path, representing the popped JSON values, + or None for JSON values matching the path that are not an array or are an empty array. + If a value is not an array, its corresponding return value is null. + For legacy path (`path` starts with `.`): + Returns a bytes string representing the popped JSON value, or None if the array at `path` is empty. + If multiple paths match, the value from the first matching array that is not empty is returned. + If the JSON value at `path` is not a array or if `path` doesn't exist, an error is raised. + If `key` doesn't exist, an error is raised. + + Examples: + >>> from glide import json + >>> await json.set(client, "doc", "$", '{"a": [1, 2, true], "b": {"a": [3, 4, ["value", 3, false], 5], "c": {"a": 42}}}') + b'OK' # JSON is successfully set + >>> await json.arrpop(client, "doc", JsonArrPopOptions(path="$.a", index=1)) + [b'2'] # Pop second element from array at path $.a + >>> await json.arrpop(client, "doc", JsonArrPopOptions(path="$..a")) + [b'true', b'5', None] # Pop last elements from all arrays matching path `..a` + + #### Using a legacy path (..) to pop the first matching array + >>> await json.arrpop(client, "doc", JsonArrPopOptions(path="..a")) + b"1" # First match popped (from array at path $.a) + + #### Even though only one value is returned from `..a`, subsequent arrays are also affected + >>> await json.get(client, "doc", "$..a") + b"[[], [3, 4], 42]" # Remaining elements after pop show the changes + + >>> await json.set(client, "doc", "$", '[[], ["a"], ["a", "b", "c"]]') + b'OK' # JSON is successfully set + >>> await json.arrpop(client, "doc", JsonArrPopOptions(path=".", index=-1)) + b'["a","b","c"]' # Pop last elements at path `.` + """ + args = ["JSON.ARRPOP", key] + if options: + args.extend(options.get_options()) + + return cast( + Optional[TJsonResponse[bytes]], + await client.custom_command(args), + ) + + async def delete( client: TGlideClient, key: TEncodable, diff --git a/python/python/tests/tests_server_modules/test_json.py b/python/python/tests/tests_server_modules/test_json.py index 2c9ddabce4..38ad53e5c6 100644 --- a/python/python/tests/tests_server_modules/test_json.py +++ b/python/python/tests/tests_server_modules/test_json.py @@ -5,7 +5,7 @@ import pytest from glide.async_commands.core import ConditionalChange, InfoSection from glide.async_commands.server_modules import json -from glide.async_commands.server_modules.json import JsonGetOptions +from glide.async_commands.server_modules.json import JsonArrPopOptions, JsonGetOptions from glide.config import ProtocolVersion from glide.constants import OK from glide.exceptions import RequestError @@ -196,3 +196,76 @@ async def test_json_toggle(self, glide_client: TGlideClient): with pytest.raises(RequestError): assert await json.toggle(glide_client, "non_exiting_key", "$") + + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_json_arrpop(self, glide_client: TGlideClient): + key = get_random_string(5) + key2 = get_random_string(5) + + json_value = '{"a": [1, 2, true], "b": {"a": [3, 4, ["value", 3, false] ,5], "c": {"a": 42}}}' + assert await json.set(glide_client, key, "$", json_value) == OK + + assert await json.arrpop( + glide_client, key, JsonArrPopOptions(path="$.a", index=1) + ) == [b"2"] + assert ( + await json.arrpop(glide_client, key, JsonArrPopOptions(path="$..a")) + ) == [b"true", b"5", None] + + assert ( + await json.arrpop(glide_client, key, JsonArrPopOptions(path="..a")) == b"1" + ) + # Even if only one array element was returned, ensure second array at `..a` was popped + assert await json.get(glide_client, key, "$..a") == b"[[],[3,4],42]" + + # Out of index + assert await json.arrpop( + glide_client, key, JsonArrPopOptions(path="$..a", index=10) + ) == [None, b"4", None] + + assert ( + await json.arrpop( + glide_client, key, JsonArrPopOptions(path="..a", index=-10) + ) + == b"3" + ) + + # Path is not an array + assert await json.arrpop(glide_client, key, JsonArrPopOptions(path="$")) == [ + None + ] + with pytest.raises(RequestError): + assert await json.arrpop(glide_client, key, JsonArrPopOptions(path=".")) + with pytest.raises(RequestError): + assert await json.arrpop(glide_client, key) + + # Non existing path + assert ( + await json.arrpop( + glide_client, key, JsonArrPopOptions(path="$.non_existing_path") + ) + == [] + ) + with pytest.raises(RequestError): + assert await json.arrpop( + glide_client, key, JsonArrPopOptions(path="non_existing_path") + ) + + with pytest.raises(RequestError): + await json.arrpop( + glide_client, "non_existing_key", JsonArrPopOptions(path="$.a") + ) + with pytest.raises(RequestError): + await json.arrpop( + glide_client, "non_existing_key", JsonArrPopOptions(path=".a") + ) + + assert ( + await json.set(glide_client, key2, "$", '[[], ["a"], ["a", "b", "c"]]') + == OK + ) + assert ( + await json.arrpop(glide_client, key2, JsonArrPopOptions(path=".", index=-1)) + == b'["a","b","c"]' + )