diff --git a/python/python/glide/async_commands/server_modules/json.py b/python/python/glide/async_commands/server_modules/json.py index 7da1c6c4aa..4149ce0127 100644 --- a/python/python/glide/async_commands/server_modules/json.py +++ b/python/python/glide/async_commands/server_modules/json.py @@ -139,6 +139,50 @@ async def get( return cast(bytes, await client.custom_command(args)) +async def mget( + client: TGlideClient, + keys: List[TEncodable], + paths: Optional[Union[TEncodable, List[TEncodable]]] = None, + options: Optional[JsonGetOptions] = None, +) -> Optional[List[bytes]]: + """ + Retrieves the JSON values at the specified `paths` stored at multiple `keys`. + + See https://valkey.io/commands/json.mget/ for more details. + + Args: + client (TGlideClient): The Redis client to execute the command. + keys (List[TEncodable]): A list of keys for the JSON documents. + paths (Optional[Union[TEncodable, List[TEncodable]]]): The path or list of paths within the JSON documents. Default is root `$`. + options (Optional[JsonGetOptions]): Options for formatting the byte representation of the JSON data. See `JsonGetOptions`. + + Returns: + Optional[List[bytes]]: A list of bytes representations of the returned values. + If a key doesn't exist, its corresponding entry will be `None`. + + Examples: + >>> from glide import json as redisJson + >>> import json + >>> json_strs = await redisJson.mget(client, ["doc1", "doc2"], ["$"]) + >>> [json.loads(js) for js in json_strs] # Parse JSON strings to Python data + [[{"a": 1.0, "b": 2}], [{"a": 2.0, "b": {"a": 3.0, "b" : 4.0}}]] # JSON objects retrieved from keys `doc1` and `doc2` + >>> await redisJson.mget(client, ["doc1", "doc2"], ["$.a"]) + [b"[1.0]", b"[2.0]"] # Returns values at path '$.a' for the JSON documents stored at `doc1` and `doc2`. + >>> await redisJson.mget(client, ["doc1"], ["$.non_existing_path"]) + [None] # Returns an empty array since the path '$.non_existing_path' does not exist in the JSON document stored at `doc1`. + """ + args = ["JSON.MGET"] + keys + if options: + args.extend(options.get_options()) + if paths: + if isinstance(paths, (str, bytes)): + paths = [paths] + args.extend(paths) + + results = await client.custom_command(args) + return [result if result is not None else None for result in results] + + async def arrlen( 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 794e885cfe..dca2890f6c 100644 --- a/python/python/tests/tests_server_modules/test_json.py +++ b/python/python/tests/tests_server_modules/test_json.py @@ -142,6 +142,122 @@ async def test_json_get_formatting(self, glide_client: TGlideClient): expected_result = b'[\n~{\n~~"a":*1.0,\n~~"b":*2,\n~~"c":*{\n~~~"d":*3,\n~~~"e":*4\n~~}\n~}\n]' assert result == expected_result + @pytest.mark.parametrize("cluster_mode", [True, False]) + @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) + async def test_json_mget(self, glide_client: TGlideClient): + key = get_random_string(5) + key1 = f"{{{key}}}1" + key2 = f"{{{key}}}2" + # The prefix ensures that both keys hash to the same slot + + json1_value = {"a": 1.0, "b": {"a": 1, "b": 2.5, "c": True}} + json2_value = {"a": 3.0, "b": {"a": 1, "b": 4}} + + assert ( + await json.set(glide_client, key1, "$", OuterJson.dumps(json1_value)) == OK + ) + assert ( + await json.set(glide_client, key2, "$", OuterJson.dumps(json2_value)) == OK + ) + + result = await json.mget( + glide_client, + [key1, key2], + ["$"], + ) + expected_result = [ + b'[{"a":1.0,"b":{"a":1,"b":2.5,"c":true}}]', + b'[{"a":3.0,"b":{"a":1,"b":4}}]', + ] + assert result == expected_result + + result = await json.mget( + glide_client, + [key1, key2], + ["."], + ) + expected_result = [ + b'{"a":1.0,"b":{"a":1,"b":2.5,"c":true}}', + b'{"a":3.0,"b":{"a":1,"b":4}}', + ] + assert result == expected_result + + result = await json.mget( + glide_client, + [key1, key2], + ["$.a"], + ) + expected_result = [b"[1.0]", b"[3.0]"] + assert result == expected_result + + result = await json.mget( + glide_client, + [key1, key2], + ["$.b"], + ) + expected_result = [b'[{"a":1,"b":2.5,"c":true}]', b'[{"a":1,"b":4}]'] + assert result == expected_result + + result = await json.mget( + glide_client, + [key1, key2], + ["$..b"], + ) + expected_result = [b'[{"a":1,"b":2.5,"c":true},2.5]', b'[{"a":1,"b":4},4]'] + assert result == expected_result + + result = await json.mget( + glide_client, + [key1, key2], + [".b.b"], + ) + expected_result = [b"2.5", b"4"] + assert result == expected_result + + # Path doesn't exist + result = await json.mget( + glide_client, + [key1, key2], + ["$non_existing_path"], + ) + expected_result = [b"[]", b"[]"] + assert result == expected_result + + # Keys don't exist + result = await json.mget( + glide_client, + ["{non_existing_key}1", "{non_existing_key}2"], + ["$a"], + ) + expected_result = [None, None] + assert result == expected_result + + # Test with only one key + result = await json.mget( + glide_client, + [key1], + ["$.a"], + ) + expected_result = [b"[1.0]"] + assert result == expected_result + + # Value at path isnt an object + result = await json.mget( + glide_client, + [key1, key2], + ["$.e"], + ) + expected_result = [b"[]", b"[]"] + assert result == expected_result + + # No path given + result = await json.mget( + glide_client, + [key1, key2], + ) + expected_result = [None] + assert result == expected_result + @pytest.mark.parametrize("cluster_mode", [True, False]) @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3]) async def test_del(self, glide_client: TGlideClient):