Python: adds json module and JSON.SET JSON.GET commands (valkey-io#1056)

CHANGELOG.md

@@ -16,6 +16,7 @@
 * Node: Added HVAL command ([#1022](https://github.com/aws/glide-for-redis/pull/1022))
 * Node: Added PERSIST command ([#1023](https://github.com/aws/glide-for-redis/pull/1023))
 * Node: Added Xadd, Xtrim commands. ([#1057](https://github.com/aws/glide-for-redis/pull/1057))
+* Python: Added json module and JSON.SET JSON.GET commands  ([#1056](https://github.com/aws/glide-for-redis/pull/1056))
 
 #### Features
 

glide-core/src/protobuf/redis_request.proto

@@ -129,6 +129,8 @@ enum RequestType {
     PTTL = 85;
     ZRemRangeByRank = 86;
     Persist = 87;
+    JsonSet = 88;
+    JsonGet = 89;
 }
 
 message Command {

glide-core/src/socket_listener.rs

@@ -359,6 +359,8 @@ fn get_command(request: &Command) -> Option<Cmd> {
         RequestType::PTTL => Some(cmd("PTTL")),
         RequestType::ZRemRangeByRank => Some(cmd("ZREMRANGEBYRANK")),
         RequestType::Persist => Some(cmd("PERSIST")),
+        RequestType::JsonSet => Some(cmd("JSON.SET")),
+        RequestType::JsonGet => Some(cmd("JSON.GET")),
     }
 }
 

python/pytest.ini

@@ -1,4 +1,4 @@
 [pytest]
 markers =
     smoke_test: mark a test as a build verification testing.
-addopts = -k "not test_redis_modules.py"
+addopts = -k "not redis_modules"

python/python/glide/__init__.py

@@ -8,6 +8,7 @@
     InfoSection,
     UpdateOptions,
 )
+from glide.async_commands.redis_modules import json
 from glide.async_commands.sorted_set import (
     InfBound,
     LexBoundary,
@@ -60,6 +61,7 @@
     "ExpiryType",
     "InfBound",
     "InfoSection",
+    "json",
     "LexBoundary",
     "Limit",
     "RangeByIndex",

python/python/glide/async_commands/redis_modules/json.py

@@ -0,0 +1,139 @@
+# Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0
+"""module for `RedisJSON` commands.
+
+    Examples:
+
+        >>> from glide import json as redisJson
+        >>> import json
+        >>> value = {'a': 1.0, 'b': 2}
+        >>> json_str = json.dumps(value) # Convert Python dictionary to JSON string using json.dumps()
+        >>> await redisJson.set(client, "doc", "$", json_str)
+            'OK'  # Indicates successful setting of the value at path '$' in the key stored at `doc`.
+        >>> json_get = await redisJson.get(client, "doc", "$") # Returns the value at path '$' in the JSON document stored at `doc` as JSON string.
+        >>> print(json_get)
+            "[{\"a\":1.0,\"b\":2}]" 
+        >>> json.loads(json_get)
+            [{"a": 1.0, "b" :2}] # JSON object retrieved from the key `doc` using json.loads()
+        """
+from typing import List, Optional, Union, cast
+
+from glide.async_commands.core import ConditionalChange
+from glide.constants import TOK
+from glide.protobuf.redis_request_pb2 import RequestType
+from glide.redis_client import TRedisClient
+
+
+class JsonGetOptions:
+    """
+    Represents options for formatting JSON data, to be used in  the [JSON.GET](https://redis.io/commands/json.get/) command.
+
+    Args:
+        indent (Optional[str]): Sets an indentation string for nested levels. Defaults to None.
+        newline (Optional[str]): Sets a string that's printed at the end of each line. Defaults to None.
+        space (Optional[str]): Sets a string that's put between a key and a value. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        indent: Optional[str] = None,
+        newline: Optional[str] = None,
+        space: Optional[str] = None,
+    ):
+        self.indent = indent
+        self.new_line = newline
+        self.space = space
+
+    def get_options(self) -> List[str]:
+        args = []
+        if self.indent:
+            args.extend(["INDENT", self.indent])
+        if self.new_line:
+            args.extend(["NEWLINE", self.new_line])
+        if self.space:
+            args.extend(["SPACE", self.space])
+        return args
+
+
+async def set(
+    client: TRedisClient,
+    key: str,
+    path: str,
+    value: str,
+    set_condition: Optional[ConditionalChange] = None,
+) -> Optional[TOK]:
+    """
+    Sets the JSON value at the specified `path` stored at `key`.
+
+    See https://redis.io/commands/json.set/ for more details.
+
+    Args:
+        client (TRedisClient): The Redis client to execute the command.
+        key (str): The key of the JSON document.
+        path (str): Represents the path within the JSON document where the value will be set.
+            The key will be modified only if `value` is added as the last child in the specified `path`, or if the specified `path` acts as the parent of a new child being added.
+        value (set): The value to set at the specific path, in JSON formatted str.
+        set_condition (Optional[ConditionalChange]): Set the value only if the given condition is met (within the key or path).
+            Equivalent to [`XX` | `NX`] in the Redis API. Defaults to None.
+
+    Returns:
+        Optional[TOK]: If the value is successfully set, returns OK.
+            If value isn't set because of `set_condition`, returns None.
+
+    Examples:
+        >>> from glide import json as redisJson
+        >>> import json
+        >>> value = {'a': 1.0, 'b': 2}
+        >>> json_str = json.dumps(value)
+        >>> await redisJson.set(client, "doc", "$", json_str)
+            'OK'  # Indicates successful setting of the value at path '$' in the key stored at `doc`.
+    """
+    args = [key, path, value]
+    if set_condition:
+        args.append(set_condition.value)
+
+    return cast(Optional[TOK], await client._execute_command(RequestType.JsonSet, args))
+
+
+async def get(
+    client: TRedisClient,
+    key: str,
+    paths: Optional[Union[str, List[str]]] = None,
+    options: Optional[JsonGetOptions] = None,
+) -> Optional[str]:
+    """
+    Retrieves the JSON value at the specified `paths` stored at `key`.
+
+    See https://redis.io/commands/json.get/ for more details.
+
+    Args:
+        client (TRedisClient): The Redis client to execute the command.
+        key (str): The key of the JSON document.
+        paths (Optional[Union[str, List[str]]]): The path or list of paths within the JSON document. Default is root `$`.
+        options (Optional[JsonGetOptions]): Options for formatting the string representation of the JSON data. See `JsonGetOptions`.
+
+    Returns:
+        str: A bulk string representation of the returned value.
+            If `key` doesn't exists, returns None.
+
+    Examples:
+        >>> from glide import json as redisJson
+        >>> import json
+        >>> json_str = await redisJson.get(client, "doc", "$")
+        >>> json.loads(json_str) # Parse JSON string to Python data
+            [{"a": 1.0, "b" :2}]  # JSON object retrieved from the key `doc` using json.loads()
+        >>> await redisJson.get(client, "doc", "$")
+            "[{\"a\":1.0,\"b\":2}]"  # Returns the value at path '$' in the JSON document stored at `doc`.
+        >>> await redisJson.get(client, "doc", ["$.a", "$.b"], json.JsonGetOptions(indent="  ", newline="\n", space=" "))
+            "{\n \"$.a\": [\n  1.0\n ],\n \"$.b\": [\n  2\n ]\n}"  # Returns the values at paths '$.a' and '$.b' in the JSON document stored at `doc`, with specified formatting options.
+        >>> await redisJson.get(client, "doc", "$.non_existing_path")
+            "[]"  # Returns an empty array since the path '$.non_existing_path' does not exist in the JSON document stored at `doc`.
+    """
+    args = [key]
+    if options:
+        args.extend(options.get_options())
+    if paths:
+        if isinstance(paths, str):
+            paths = [paths]
+        args.extend(paths)
+
+    return cast(str, await client._execute_command(RequestType.JsonGet, args))

python/python/tests/tests_redis_modules/test_json.py

@@ -0,0 +1,150 @@
+# Copyright GLIDE-for-Redis Project Contributors - SPDX Identifier: Apache-2.0
+
+import json as OuterJson
+
+import pytest
+from glide.async_commands.core import ConditionalChange, InfoSection
+from glide.async_commands.redis_modules import json
+from glide.async_commands.redis_modules.json import JsonGetOptions
+from glide.config import ProtocolVersion
+from glide.constants import OK
+from glide.redis_client import TRedisClient
+from tests.test_async_client import get_random_string, parse_info_response
+
+
+@pytest.mark.asyncio
+class TestJson:
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_json_module_is_loaded(self, redis_client: TRedisClient):
+        res = parse_info_response(await redis_client.info([InfoSection.MODULES]))
+        assert "ReJSON" in res["module"]
+
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_json_set_get(self, redis_client: TRedisClient):
+        key = get_random_string(5)
+
+        json_value = {"a": 1.0, "b": 2}
+        assert await json.set(redis_client, key, "$", OuterJson.dumps(json_value)) == OK
+
+        result = await json.get(redis_client, key, ".")
+        assert isinstance(result, str)
+        assert OuterJson.loads(result) == json_value
+
+        result = await json.get(redis_client, key, ["$.a", "$.b"])
+        assert isinstance(result, str)
+        assert OuterJson.loads(result) == {"$.a": [1.0], "$.b": [2]}
+
+        assert await json.get(redis_client, "non_existing_key", "$") is None
+        assert await json.get(redis_client, key, "$.d") == "[]"
+
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_json_set_get_multiple_values(self, redis_client: TRedisClient):
+        key = get_random_string(5)
+
+        assert (
+            await json.set(
+                redis_client,
+                key,
+                "$",
+                OuterJson.dumps({"a": {"c": 1, "d": 4}, "b": {"c": 2}, "c": True}),
+            )
+            == OK
+        )
+
+        result = await json.get(redis_client, key, "$..c")
+        assert isinstance(result, str)
+        assert OuterJson.loads(result) == [True, 1, 2]
+
+        result = await json.get(redis_client, key, ["$..c", "$.c"])
+        assert isinstance(result, str)
+        assert OuterJson.loads(result) == {"$..c": [True, 1, 2], "$.c": [True]}
+
+        assert await json.set(redis_client, key, "$..c", '"new_value"') == OK
+        result = await json.get(redis_client, key, "$..c")
+        assert isinstance(result, str)
+        assert OuterJson.loads(result) == ["new_value"] * 3
+
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_json_set_conditional_set(self, redis_client: TRedisClient):
+        key = get_random_string(5)
+        value = OuterJson.dumps({"a": 1.0, "b": 2})
+        assert (
+            await json.set(
+                redis_client,
+                key,
+                "$",
+                value,
+                ConditionalChange.ONLY_IF_EXISTS,
+            )
+            is None
+        )
+        assert (
+            await json.set(
+                redis_client,
+                key,
+                "$",
+                value,
+                ConditionalChange.ONLY_IF_DOES_NOT_EXIST,
+            )
+            == OK
+        )
+
+        assert (
+            await json.set(
+                redis_client,
+                key,
+                "$.a",
+                "4.5",
+                ConditionalChange.ONLY_IF_DOES_NOT_EXIST,
+            )
+            is None
+        )
+
+        assert await json.get(redis_client, key, ".a") == "1.0"
+
+        assert (
+            await json.set(
+                redis_client,
+                key,
+                "$.a",
+                "4.5",
+                ConditionalChange.ONLY_IF_EXISTS,
+            )
+            == OK
+        )
+
+        assert await json.get(redis_client, key, ".a") == "4.5"
+
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_json_get_formatting(self, redis_client: TRedisClient):
+        key = get_random_string(5)
+        assert (
+            await json.set(
+                redis_client,
+                key,
+                "$",
+                OuterJson.dumps({"a": 1.0, "b": 2, "c": {"d": 3, "e": 4}}),
+            )
+            == OK
+        )
+
+        result = await json.get(
+            redis_client, key, "$", JsonGetOptions(indent="  ", newline="\n", space=" ")
+        )
+
+        expected_result = '[\n  {\n    "a": 1.0,\n    "b": 2,\n    "c": {\n      "d": 3,\n      "e": 4\n    }\n  }\n]'
+        assert result == expected_result
+
+        result = await json.get(
+            redis_client, key, "$", JsonGetOptions(indent="~", newline="\n", space="*")
+        )
+
+        expected_result = (
+            '[\n~{\n~~"a":*1.0,\n~~"b":*2,\n~~"c":*{\n~~~"d":*3,\n~~~"e":*4\n~~}\n~}\n]'
+        )
+        assert result == expected_result