Python: Add FUNCTION STATS command (valkey-io#1794)

* Add cluster and standalone versions of FUNCTION STATS commands * Add transaction implementation for FUNCTION STATS * Start adding tests for FUNCTION STATS * Try finish implementing tests * Fix some issues in the tests * Fix tests * Document check_function_stats_response * Run linters and formatters * Import pytest * Fix import ordering * Fix mypy errors * Fix transaction tests * Address PR comments * Fix black lint * Address minor PR comments --------- Co-authored-by: aaron-congo <[email protected]>
Bit-Quill · Jul 4, 2024 · b2d36d7 · b2d36d7
1 parent c0e9e74
commit b2d36d7
Show file tree

Hide file tree

Showing 8 changed files with 318 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -208,6 +208,7 @@
 * Node: Added LINDEX command ([#999](https://github.com/aws/glide-for-redis/pull/999))
 * Python, Node: Added ZPOPMAX command ([#996](https://github.com/aws/glide-for-redis/pull/996), [#1009](https://github.com/aws/glide-for-redis/pull/1009))
 * Python: Added DBSIZE command ([#1040](https://github.com/aws/glide-for-redis/pull/1040))
+* Python: Added FUNCTION STATS command ([#1794](https://github.com/aws/glide-for-redis/pull/1794))
 
 #### Features
 * Python, Node: Added support in Lua Scripts ([#775](https://github.com/aws/glide-for-redis/pull/775), [#860](https://github.com/aws/glide-for-redis/pull/860))

diff --git a/python/python/glide/async_commands/cluster_commands.py b/python/python/glide/async_commands/cluster_commands.py
@@ -18,6 +18,7 @@
     TClusterResponse,
     TEncodable,
     TFunctionListResponse,
+    TFunctionStatsResponse,
     TResult,
     TSingleNodeRoute,
 )
@@ -555,6 +556,48 @@ async def fcall_ro_route(
             await self._execute_command(RequestType.FCallReadOnly, args, route),
         )
 
+    async def function_stats(
+        self, route: Optional[Route] = None
+    ) -> TClusterResponse[TFunctionStatsResponse]:
+        """
+        Returns information about the function that's currently running and information about the
+        available execution engines.
+
+        See https://valkey.io/commands/function-stats/ for more details
+
+        Args:
+            route (Optional[Route]): Specifies the routing configuration for the command. The client
+                will route the command to the nodes defined by `route`.
+
+        Returns:
+            TClusterResponse[TFunctionStatsResponse]: A `Mapping` with two keys:
+                - `running_script` with information about the running script.
+                - `engines` with information about available engines and their stats.
+                See example for more details.
+
+        Examples:
+            >>> await client.function_stats(RandomNode())
+                {
+                    'running_script': {
+                        'name': 'foo',
+                        'command': ['FCALL', 'foo', '0', 'hello'],
+                        'duration_ms': 7758
+                    },
+                    'engines': {
+                        'LUA': {
+                            'libraries_count': 1,
+                            'functions_count': 1,
+                        }
+                    }
+                }
+
+        Since: Redis version 7.0.0.
+        """
+        return cast(
+            TClusterResponse[TFunctionStatsResponse],
+            await self._execute_command(RequestType.FunctionStats, [], route),
+        )
+
     async def function_dump(
         self, route: Optional[Route] = None
     ) -> TClusterResponse[bytes]:

diff --git a/python/python/glide/async_commands/standalone_commands.py b/python/python/glide/async_commands/standalone_commands.py
@@ -13,7 +13,14 @@
     _build_sort_args,
 )
 from glide.async_commands.transaction import BaseTransaction, Transaction
-from glide.constants import OK, TOK, TEncodable, TFunctionListResponse, TResult
+from glide.constants import (
+    OK,
+    TOK,
+    TEncodable,
+    TFunctionListResponse,
+    TFunctionStatsResponse,
+    TResult,
+)
 from glide.protobuf.redis_request_pb2 import RequestType
 
 
@@ -362,6 +369,42 @@ async def function_delete(self, library_name: TEncodable) -> TOK:
             ),
         )
 
+    async def function_stats(self) -> TFunctionStatsResponse:
+        """
+        Returns information about the function that's currently running and information about the
+        available execution engines.
+
+        See https://valkey.io/commands/function-stats/ for more details
+
+        Returns:
+            TFunctionStatsResponse: A `Mapping` with two keys:
+                - `running_script` with information about the running script.
+                - `engines` with information about available engines and their stats.
+                See example for more details.
+
+        Examples:
+            >>> await client.function_stats()
+                {
+                    'running_script': {
+                        'name': 'foo',
+                        'command': ['FCALL', 'foo', '0', 'hello'],
+                        'duration_ms': 7758
+                    },
+                    'engines': {
+                        'LUA': {
+                            'libraries_count': 1,
+                            'functions_count': 1,
+                        }
+                    }
+                }
+
+        Since: Redis version 7.0.0.
+        """
+        return cast(
+            TFunctionStatsResponse,
+            await self._execute_command(RequestType.FunctionStats, []),
+        )
+
     async def function_dump(self) -> bytes:
         """
         Returns the serialized payload of all loaded libraries.

diff --git a/python/python/glide/async_commands/transaction.py b/python/python/glide/async_commands/transaction.py
@@ -2019,6 +2019,23 @@ def fcall_ro(
             args.extend(arguments)
         return self.append_command(RequestType.FCallReadOnly, args)
 
+    def function_stats(self: TTransaction) -> TTransaction:
+        """
+        Returns information about the function that's currently running and information about the
+        available execution engines.
+
+        See https://valkey.io/commands/function-stats/ for more details
+
+        Command Response:
+            TFunctionStatsResponse: A `Mapping` with two keys:
+                - `running_script` with information about the running script.
+                - `engines` with information about available engines and their stats.
+                See example for more details.
+
+        Since: Redis version 7.0.0.
+        """
+        return self.append_command(RequestType.FunctionStats, [])
+
     def xadd(
         self: TTransaction,
         key: TEncodable,

diff --git a/python/python/glide/constants.py b/python/python/glide/constants.py
@@ -42,3 +42,12 @@
         Union[bytes, List[Mapping[bytes, Union[bytes, Set[bytes]]]]],
     ]
 ]
+TFunctionStatsResponse = Mapping[
+    bytes,
+    Union[
+        None,
+        Mapping[
+            bytes, Union[Mapping[bytes, Mapping[bytes, int]], bytes, int, List[bytes]]
+        ],
+    ],
+]
diff --git a/python/python/tests/test_async_client.py b/python/python/tests/test_async_client.py
@@ -76,7 +76,7 @@
     ProtocolVersion,
     ServerCredentials,
 )
-from glide.constants import OK, TEncodable, TResult
+from glide.constants import OK, TEncodable, TFunctionStatsResponse, TResult
 from glide.glide_client import GlideClient, GlideClusterClient, TGlideClient
 from glide.routes import (
     AllNodes,
@@ -91,6 +91,7 @@
 from tests.conftest import create_client
 from tests.utils.utils import (
     check_function_list_response,
+    check_function_stats_response,
     check_if_server_version_lt,
     compare_maps,
     convert_bytes_to_string_object,
@@ -8087,6 +8088,153 @@ async def test_function_delete_with_routing(
             await glide_client.function_delete(lib_name)
         assert "Library not found" in str(e)
 
+    @pytest.mark.parametrize("cluster_mode", [False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_function_stats(self, glide_client: GlideClient):
+        min_version = "7.0.0"
+        if await check_if_server_version_lt(glide_client, min_version):
+            return pytest.mark.skip(reason=f"Redis version required >= {min_version}")
+
+        lib_name = "functionStats"
+        func_name = lib_name
+        assert await glide_client.function_flush(FlushMode.SYNC) == OK
+
+        # function $funcName returns first argument
+        code = generate_lua_lib_code(lib_name, {func_name: "return args[1]"}, False)
+        assert await glide_client.function_load(code, True) == lib_name.encode()
+
+        response = await glide_client.function_stats()
+        check_function_stats_response(response, [], 1, 1)
+
+        code = generate_lua_lib_code(
+            lib_name + "_2",
+            {func_name + "_2": "return 'OK'", func_name + "_3": "return 42"},
+            False,
+        )
+        assert (
+            await glide_client.function_load(code, True) == (lib_name + "_2").encode()
+        )
+
+        response = await glide_client.function_stats()
+        check_function_stats_response(response, [], 2, 3)
+
+        assert await glide_client.function_flush(FlushMode.SYNC) == OK
+
+        response = await glide_client.function_stats()
+        check_function_stats_response(response, [], 0, 0)
+
+    @pytest.mark.parametrize("cluster_mode", [True])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_function_stats_cluster(self, glide_client: GlideClusterClient):
+        min_version = "7.0.0"
+        if await check_if_server_version_lt(glide_client, min_version):
+            return pytest.mark.skip(reason=f"Redis version required >= {min_version}")
+
+        lib_name = "functionStats_without_route"
+        func_name = lib_name
+        assert await glide_client.function_flush(FlushMode.SYNC) == OK
+
+        # function $funcName returns first argument
+        code = generate_lua_lib_code(lib_name, {func_name: "return args[1]"}, False)
+        assert await glide_client.function_load(code, True) == lib_name.encode()
+
+        response = await glide_client.function_stats()
+        for node_response in response.values():
+            check_function_stats_response(
+                cast(TFunctionStatsResponse, node_response), [], 1, 1
+            )
+
+        code = generate_lua_lib_code(
+            lib_name + "_2",
+            {func_name + "_2": "return 'OK'", func_name + "_3": "return 42"},
+            False,
+        )
+        assert (
+            await glide_client.function_load(code, True) == (lib_name + "_2").encode()
+        )
+
+        response = await glide_client.function_stats()
+        for node_response in response.values():
+            check_function_stats_response(
+                cast(TFunctionStatsResponse, node_response), [], 2, 3
+            )
+
+        assert await glide_client.function_flush(FlushMode.SYNC) == OK
+
+        response = await glide_client.function_stats()
+        for node_response in response.values():
+            check_function_stats_response(
+                cast(TFunctionStatsResponse, node_response), [], 0, 0
+            )
+
+    @pytest.mark.parametrize("cluster_mode", [True])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    @pytest.mark.parametrize("single_route", [True, False])
+    async def test_function_stats_with_routing(
+        self, glide_client: GlideClusterClient, single_route: bool
+    ):
+        min_version = "7.0.0"
+        if await check_if_server_version_lt(glide_client, min_version):
+            return pytest.mark.skip(reason=f"Redis version required >= {min_version}")
+
+        route = (
+            SlotKeyRoute(SlotType.PRIMARY, get_random_string(10))
+            if single_route
+            else AllPrimaries()
+        )
+        lib_name = "functionStats_with_route_" + str(single_route)
+        func_name = lib_name
+        assert await glide_client.function_flush(FlushMode.SYNC, route) == OK
+
+        # function $funcName returns first argument
+        code = generate_lua_lib_code(lib_name, {func_name: "return args[1]"}, False)
+        assert await glide_client.function_load(code, True, route) == lib_name.encode()
+
+        response = await glide_client.function_stats(route)
+        if single_route:
+            check_function_stats_response(
+                cast(TFunctionStatsResponse, response), [], 1, 1
+            )
+        else:
+            for node_response in response.values():
+                check_function_stats_response(
+                    cast(TFunctionStatsResponse, node_response), [], 1, 1
+                )
+
+        code = generate_lua_lib_code(
+            lib_name + "_2",
+            {func_name + "_2": "return 'OK'", func_name + "_3": "return 42"},
+            False,
+        )
+        assert (
+            await glide_client.function_load(code, True, route)
+            == (lib_name + "_2").encode()
+        )
+
+        response = await glide_client.function_stats(route)
+        if single_route:
+            check_function_stats_response(
+                cast(TFunctionStatsResponse, response), [], 2, 3
+            )
+        else:
+            for node_response in response.values():
+                check_function_stats_response(
+                    cast(TFunctionStatsResponse, node_response), [], 2, 3
+                )
+
+        assert await glide_client.function_flush(FlushMode.SYNC, route) == OK
+
+        response = await glide_client.function_stats(route)
+        if single_route:
+            check_function_stats_response(
+                cast(TFunctionStatsResponse, response), [], 0, 0
+            )
+        else:
+            for node_response in response.values():
+                check_function_stats_response(
+                    cast(TFunctionStatsResponse, node_response), [], 0, 0
+                )
+
     @pytest.mark.parametrize("cluster_mode", [True])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_fcall_with_key(self, glide_client: GlideClusterClient):

diff --git a/python/python/tests/test_transaction.py b/python/python/tests/test_transaction.py
@@ -159,6 +159,18 @@ async def transaction_test(
         args.append(OK)
         transaction.function_flush(FlushMode.SYNC)
         args.append(OK)
+        transaction.function_stats()
+        args.append(
+            {
+                b"running_script": None,
+                b"engines": {
+                    b"LUA": {
+                        b"libraries_count": 0,
+                        b"functions_count": 0,
+                    }
+                },
+            }
+        )
 
     transaction.dbsize()
     args.append(0)