Adding query cache management (#36)

Author: apetenchea

arangoasync/aql.py

@@ -1,11 +1,15 @@
-__all__ = ["AQL"]
+__all__ = ["AQL", "AQLQueryCache"]
 
 
 from typing import Optional
 
 from arangoasync.cursor import Cursor
 from arangoasync.errno import HTTP_NOT_FOUND
 from arangoasync.exceptions import (
+    AQLCacheClearError,
+    AQLCacheConfigureError,
+    AQLCacheEntriesError,
+    AQLCachePropertiesError,
     AQLQueryClearError,
     AQLQueryExecuteError,
     AQLQueryExplainError,
@@ -23,13 +27,196 @@
 from arangoasync.typings import (
     Json,
     Jsons,
+    QueryCacheProperties,
     QueryExplainOptions,
     QueryProperties,
     QueryTrackingConfiguration,
     Result,
 )
 
 
+class AQLQueryCache:
+    """AQL Query Cache API wrapper.
+
+    Args:
+        executor: API executor. Required to execute the API requests.
+    """
+
+    def __init__(self, executor: ApiExecutor) -> None:
+        self._executor = executor
+
+    @property
+    def name(self) -> str:
+        """Return the name of the current database."""
+        return self._executor.db_name
+
+    @property
+    def serializer(self) -> Serializer[Json]:
+        """Return the serializer."""
+        return self._executor.serializer
+
+    @property
+    def deserializer(self) -> Deserializer[Json, Jsons]:
+        """Return the deserializer."""
+        return self._executor.deserializer
+
+    def __repr__(self) -> str:
+        return f"<AQLQueryCache in {self.name}>"
+
+    async def entries(self) -> Result[Jsons]:
+        """Return a list of all AQL query results cache entries.
+
+
+        Returns:
+            list: List of AQL query results cache entries.
+
+        Raises:
+            AQLCacheEntriesError: If retrieval fails.
+
+        References:
+            - `list-the-entries-of-the-aql-query-results-cache <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-results-cache/#list-the-entries-of-the-aql-query-results-cache>`__
+        """  # noqa: E501
+        request = Request(method=Method.GET, endpoint="/_api/query-cache/entries")
+
+        def response_handler(resp: Response) -> Jsons:
+            if not resp.is_success:
+                raise AQLCacheEntriesError(resp, request)
+            return self.deserializer.loads_many(resp.raw_body)
+
+        return await self._executor.execute(request, response_handler)
+
+    async def plan_entries(self) -> Result[Jsons]:
+        """Return a list of all AQL query plan cache entries.
+
+        Returns:
+            list: List of AQL query plan cache entries.
+
+        Raises:
+            AQLCacheEntriesError: If retrieval fails.
+
+        References:
+            - `list-the-entries-of-the-aql-query-plan-cache <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-plan-cache/#list-the-entries-of-the-aql-query-plan-cache>`__
+        """  # noqa: E501
+        request = Request(method=Method.GET, endpoint="/_api/query-plan-cache")
+
+        def response_handler(resp: Response) -> Jsons:
+            if not resp.is_success:
+                raise AQLCacheEntriesError(resp, request)
+            return self.deserializer.loads_many(resp.raw_body)
+
+        return await self._executor.execute(request, response_handler)
+
+    async def clear(self) -> Result[None]:
+        """Clear the AQL query results cache.
+
+        Raises:
+            AQLCacheClearError: If clearing the cache fails.
+
+        References:
+            - `clear-the-aql-query-results-cache <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-results-cache/#clear-the-aql-query-results-cache>`__
+        """  # noqa: E501
+        request = Request(method=Method.DELETE, endpoint="/_api/query-cache")
+
+        def response_handler(resp: Response) -> None:
+            if not resp.is_success:
+                raise AQLCacheClearError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
+    async def clear_plan(self) -> Result[None]:
+        """Clear the AQL query plan cache.
+
+        Raises:
+            AQLCacheClearError: If clearing the cache fails.
+
+        References:
+            - `clear-the-aql-query-plan-cache <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-plan-cache/#clear-the-aql-query-plan-cache>`__
+        """  # noqa: E501
+        request = Request(method=Method.DELETE, endpoint="/_api/query-plan-cache")
+
+        def response_handler(resp: Response) -> None:
+            if not resp.is_success:
+                raise AQLCacheClearError(resp, request)
+
+        return await self._executor.execute(request, response_handler)
+
+    async def properties(self) -> Result[QueryCacheProperties]:
+        """Return the current AQL query results cache configuration.
+
+        Returns:
+            QueryCacheProperties: Current AQL query cache properties.
+
+        Raises:
+            AQLCachePropertiesError: If retrieval fails.
+
+        References:
+            - `get-the-aql-query-results-cache-configuration <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-results-cache/#get-the-aql-query-results-cache-configuration>`__
+        """  # noqa: E501
+        request = Request(method=Method.GET, endpoint="/_api/query-cache/properties")
+
+        def response_handler(resp: Response) -> QueryCacheProperties:
+            if not resp.is_success:
+                raise AQLCachePropertiesError(resp, request)
+            return QueryCacheProperties(self.deserializer.loads(resp.raw_body))
+
+        return await self._executor.execute(request, response_handler)
+
+    async def configure(
+        self,
+        mode: Optional[str] = None,
+        max_results: Optional[int] = None,
+        max_results_size: Optional[int] = None,
+        max_entry_size: Optional[int] = None,
+        include_system: Optional[bool] = None,
+    ) -> Result[QueryCacheProperties]:
+        """Configure the AQL query results cache.
+
+        Args:
+            mode (str | None): Cache mode. Allowed values are `"off"`, `"on"`,
+                and `"demand"`.
+            max_results (int | None): Max number of query results stored per
+                database-specific cache.
+            max_results_size (int | None): Max cumulative size of query results stored
+                per database-specific cache.
+            max_entry_size (int | None): Max entry size of each query result stored per
+                database-specific cache.
+            include_system (bool | None): Store results of queries in system collections.
+
+        Returns:
+            QueryCacheProperties: Updated AQL query cache properties.
+
+        Raises:
+            AQLCacheConfigureError: If setting the configuration fails.
+
+        References:
+            - `set-the-aql-query-results-cache-configuration <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-results-cache/#set-the-aql-query-results-cache-configuration>`__
+        """  # noqa: E501
+        data: Json = dict()
+        if mode is not None:
+            data["mode"] = mode
+        if max_results is not None:
+            data["maxResults"] = max_results
+        if max_results_size is not None:
+            data["maxResultsSize"] = max_results_size
+        if max_entry_size is not None:
+            data["maxEntrySize"] = max_entry_size
+        if include_system is not None:
+            data["includeSystem"] = include_system
+
+        request = Request(
+            method=Method.PUT,
+            endpoint="/_api/query-cache/properties",
+            data=self.serializer.dumps(data),
+        )
+
+        def response_handler(resp: Response) -> QueryCacheProperties:
+            if not resp.is_success:
+                raise AQLCacheConfigureError(resp, request)
+            return QueryCacheProperties(self.deserializer.loads(resp.raw_body))
+
+        return await self._executor.execute(request, response_handler)
+
+
 class AQL:
     """AQL (ArangoDB Query Language) API wrapper.
 
@@ -58,6 +245,11 @@ def deserializer(self) -> Deserializer[Json, Jsons]:
         """Return the deserializer."""
         return self._executor.deserializer
 
+    @property
+    def cache(self) -> AQLQueryCache:
+        """Return the AQL Query Cache API wrapper."""
+        return AQLQueryCache(self._executor)
+
     def __repr__(self) -> str:
         return f"<AQL in {self.name}>"
 

arangoasync/database.py

@@ -27,6 +27,7 @@
     PermissionResetError,
     PermissionUpdateError,
     ServerStatusError,
+    ServerVersionError,
     TransactionAbortError,
     TransactionCommitError,
     TransactionExecuteError,
@@ -1189,6 +1190,32 @@ def response_handler(resp: Response) -> Any:
 
         return await self._executor.execute(request, response_handler)
 
+    async def version(self, details: bool = False) -> Result[Json]:
+        """Return the server version information.
+
+        Args:
+            details (bool): If `True`, return detailed version information.
+
+        Returns:
+            dict: Server version information.
+
+        Raises:
+            ServerVersionError: If the operation fails on the server side.
+
+        References:
+            - `get-the-server-version <https://docs.arangodb.com/stable/develop/http-api/administration/#get-the-server-version>`__
+        """  # noqa: E501
+        request = Request(
+            method=Method.GET, endpoint="/_api/version", params={"details": details}
+        )
+
+        def response_handler(resp: Response) -> Json:
+            if not resp.is_success:
+                raise ServerVersionError(resp, request)
+            return self.deserializer.loads(resp.raw_body)
+
+        return await self._executor.execute(request, response_handler)
+
 
 class StandardDatabase(Database):
     """Standard database API wrapper.

arangoasync/exceptions.py

@@ -71,6 +71,22 @@ def __init__(
         self.http_headers = resp.headers
 
 
+class AQLCacheClearError(ArangoServerError):
+    """Failed to clear the query cache."""
+
+
+class AQLCacheConfigureError(ArangoServerError):
+    """Failed to configure query cache properties."""
+
+
+class AQLCacheEntriesError(ArangoServerError):
+    """Failed to retrieve AQL cache entries."""
+
+
+class AQLCachePropertiesError(ArangoServerError):
+    """Failed to retrieve query cache properties."""
+
+
 class AQLQueryClearError(ArangoServerError):
     """Failed to clear slow AQL queries."""
 
@@ -251,6 +267,10 @@ class ServerStatusError(ArangoServerError):
     """Failed to retrieve server status."""
 
 
+class ServerVersionError(ArangoServerError):
+    """Failed to retrieve server version."""
+
+
 class TransactionAbortError(ArangoServerError):
     """Failed to abort transaction."""
 

arangoasync/typings.py

@@ -1096,6 +1096,9 @@ class QueryProperties(JsonWrapper):
             store intermediate and final results temporarily on disk if the number
             of rows produced by the query exceeds the specified value.
         stream (bool | None): Can be enabled to execute the query lazily.
+        use_plan_cache (bool | None): Set this option to `True` to utilize
+            a cached query plan or add the execution plan of this query to the
+            cache if it’s not in the cache yet.
 
     Example:
         .. code-block:: json
@@ -1136,6 +1139,7 @@ def __init__(
         spill_over_threshold_memory_usage: Optional[int] = None,
         spill_over_threshold_num_rows: Optional[int] = None,
         stream: Optional[bool] = None,
+        use_plan_cache: Optional[bool] = None,
     ) -> None:
         data: Json = dict()
         if allow_dirty_reads is not None:
@@ -1178,6 +1182,8 @@ def __init__(
             data["spillOverThresholdNumRows"] = spill_over_threshold_num_rows
         if stream is not None:
             data["stream"] = stream
+        if use_plan_cache is not None:
+            data["usePlanCache"] = use_plan_cache
         super().__init__(data)
 
     @property
@@ -1260,6 +1266,10 @@ def spill_over_threshold_num_rows(self) -> Optional[int]:
     def stream(self) -> Optional[bool]:
         return self._data.get("stream")
 
+    @property
+    def use_plan_cache(self) -> Optional[bool]:
+        return self._data.get("usePlanCache")
+
 
 class QueryExecutionPlan(JsonWrapper):
     """The execution plan of an AQL query.
@@ -1598,3 +1608,46 @@ def max_plans(self) -> Optional[int]:
     @property
     def optimizer(self) -> Optional[Json]:
         return self._data.get("optimizer")
+
+
+class QueryCacheProperties(JsonWrapper):
+    """AQL Cache Configuration.
+
+    Example:
+        .. code-block:: json
+
+           {
+             "mode" : "demand",
+             "maxResults" : 128,
+             "maxResultsSize" : 268435456,
+             "maxEntrySize" : 16777216,
+             "includeSystem" : false
+           }
+
+    References:
+        - `get-the-aql-query-results-cache-configuration <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-results-cache/#get-the-aql-query-results-cache-configuration>`__
+        - `set-the-aql-query-results-cache-configuration <https://docs.arangodb.com/stable/develop/http-api/queries/aql-query-results-cache/#set-the-aql-query-results-cache-configuration>`__
+    """  # noqa: E501
+
+    def __init__(self, data: Json) -> None:
+        super().__init__(data)
+
+    @property
+    def mode(self) -> str:
+        return cast(str, self._data.get("mode", ""))
+
+    @property
+    def max_results(self) -> int:
+        return cast(int, self._data.get("maxResults", 0))
+
+    @property
+    def max_results_size(self) -> int:
+        return cast(int, self._data.get("maxResultsSize", 0))
+
+    @property
+    def max_entry_size(self) -> int:
+        return cast(int, self._data.get("maxEntrySize", 0))
+
+    @property
+    def include_system(self) -> bool:
+        return cast(bool, self._data.get("includeSystem", False))