Python - Implement JSON.MGET command

CHANGELOG.md

@@ -12,6 +12,7 @@
 * Python: Add commands FT.ALIASADD, FT.ALIASDEL, FT.ALIASUPDATE([#2471](https://github.com/valkey-io/valkey-glide/pull/2471))
 * Python: Python FT.DROPINDEX command ([#2437](https://github.com/valkey-io/valkey-glide/pull/2437))
 * Python: Python: Added FT.CREATE command([#2413](https://github.com/valkey-io/valkey-glide/pull/2413))
+* Python: Add JSON.MGET command ([#2507](https://github.com/valkey-io/valkey-glide/pull/2507))
 * Python: Add JSON.ARRLEN command ([#2403](https://github.com/valkey-io/valkey-glide/pull/2403))
 * Python: Add JSON.CLEAR command ([#2418](https://github.com/valkey-io/valkey-glide/pull/2418))
 * Python: Add JSON.TYPE command ([#2409](https://github.com/valkey-io/valkey-glide/pull/2409))

python/python/glide/async_commands/server_modules/json.py

@@ -204,6 +204,57 @@ async def get(
     return cast(TJsonResponse[Optional[bytes]], await client.custom_command(args))
 
 
+async def mget(
+    client: TGlideClient,
+    keys: List[TEncodable],
+    path: Optional[TEncodable] = None,
+) -> List[Optional[bytes]]:
+    """
+    Retrieves the JSON values at the specified `path` stored at multiple `keys`.
+
+    Note:
+        In cluster mode, if keys in `keys` map to different hash slots, the command
+        will be split across these slots and executed separately for each. This means the command
+        is atomic only at the slot level. If one or more slot-specific requests fail, the entire
+        call will return the first encountered error, even though some requests may have succeeded
+        while others did not. If this behavior impacts your application logic, consider splitting
+        the request into sub-requests per slot to ensure atomicity.
+
+    Args:
+        client (TGlideClient): The client to execute the command.
+        keys (List[TEncodable]): A list of keys for the JSON documents.
+        path (Optional[TEncodable]): The path within the JSON documents. Default is root `$`.
+
+    Returns:
+        List[Optional[bytes]]:
+            List[Optional[bytes]]:
+            For JSONPath (`path` starts with `$`):
+                Returns a list of byte representations of the values found at the given path for each key.
+                If `path` does not exist within the key, the entry will be an empty array.
+            For legacy path (`path` doesn't starts with `$`):
+                Returns a list of byte representations of the values found at the given path for each key.
+                If `path` does not exist within the key, the entry will be None.
+            If a key doesn't exist, the corresponding list element will be None.
+
+
+    Examples:
+        >>> from glide import json as glideJson
+        >>> import json
+        >>> json_strs = await glideJson.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 glideJson.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 glideJson.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 path:
+        args.append(path)
+
+    return cast(TJsonResponse[Optional[bytes]], await client.custom_command(args))
+
+
 async def arrappend(
     client: TGlideClient,
     key: TEncodable,

python/python/tests/tests_server_modules/test_json.py

@@ -162,6 +162,139 @@ 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):
+        key1 = get_random_string(5)
+        key2 = get_random_string(5)
+
+        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
+        )
+
+        # Test with root JSONPath
+        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
+
+        # Retrieves the full JSON objects from multiple keys.
+        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
+
+        # Retrieves the value of the 'b' field for multiple keys.
+        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
+
+        # Retrieves all values of 'b' fields using recursive path for multiple keys
+        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
+
+        # retrieves the value of the nested 'b.b' field for multiple keys
+        result = await json.mget(
+            glide_client,
+            [key1, key2],
+            ".b.b",
+        )
+        expected_result = [b"2.5", b"4"]
+        assert result == expected_result
+
+        # Legacy path that exists in only one of the keys
+        result = await json.mget(
+            glide_client,
+            [key1, key2],
+            ".b.c",
+        )
+        expected_result = [b"true", None]
+        assert result == expected_result
+
+        # JSONPath doesn't exist
+        result = await json.mget(
+            glide_client,
+            [key1, key2],
+            "$non_existing_path",
+        )
+        expected_result = [b"[]", b"[]"]
+        assert result == expected_result
+
+        # Legacy path doesn't exist
+        result = await json.mget(
+            glide_client,
+            [key1, key2],
+            ".non_existing_path",
+        )
+        assert result == [None, None]
+
+        # One key doesn't exist
+        result = await json.mget(
+            glide_client,
+            [key1, "{non_existing_key}"],
+            "$.a",
+        )
+        assert result == [b"[1.0]", None]
+
+        # Both keys don't exist
+        result = await json.mget(
+            glide_client,
+            ["{non_existing_key}1", "{non_existing_key}2"],
+            "$a",
+        )
+        assert result == [None, None]
+
+        # Test with only one key
+        result = await json.mget(
+            glide_client,
+            [key1],
+            "$.a",
+        )
+        expected_result = [b"[1.0]"]
+        assert result == expected_result
+
+        # No path given
+        result = await json.mget(
+            glide_client,
+            [key1, key2],
+        )
+        assert result == [None]
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_json_del(self, glide_client: TGlideClient):