implement python json mget command

Signed-off-by: BoazBD <boazbd@amazon.com>

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,
@@ -187,6 +231,7 @@ async def arrlen(
     args = ["JSON.ARRLEN", key]
     if path:
         args.append(path)
+
     return cast(
         Optional[TJsonResponse[int]],
         await client.custom_command(args),
@@ -241,7 +286,6 @@ async def clear(
     args = ["JSON.CLEAR", key]
     if path:
         args.append(path)
-
     return cast(int, await client.custom_command(args))
 
 

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):
@@ -277,6 +393,12 @@ async def test_json_type(self, glide_client: TGlideClient):
         result = await json.type(glide_client, key, "[*]")
         assert result == b"string"  # Expecting only the first type (string for key1)
 
+        print("START")
+        json_str = await json.get(glide_client, "doc", "$")
+        json.loads(str(json_str))
+        await json.get(glide_client, "doc", "$")
+        print("END")
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_json_arrlen(self, glide_client: TGlideClient):