Adding usePlanCache option for AQL queries.

Author: apetenchea

arango/aql.py

@@ -144,6 +144,36 @@ def response_handler(resp: Response) -> bool:
 
         return self._execute(request, response_handler)
 
+    def plan_entries(self) -> Result[Jsons]:
+        """Return a list of all AQL query plan cache entries.
+
+        :return: List of AQL query plan cache entries.
+        :rtype: list
+        :raise arango.exceptions.AQLCacheEntriesError: If retrieval fails.
+        """
+        request = Request(method="get", endpoint="/_api/query-plan-cache")
+
+        def response_handler(resp: Response) -> Jsons:
+            if not resp.is_success:
+                raise AQLCacheEntriesError(resp, request)
+            result: Jsons = resp.body
+            return result
+
+        return self._execute(request, response_handler)
+
+    def clear_plan(self) -> Result[None]:
+        """Clear the AQL query plan cache.
+
+        :raises arango.exceptions.AQLCacheClearError: If clearing the cache fails.
+        """
+        request = Request(method="delete", endpoint="/_api/query-plan-cache")
+
+        def response_handler(resp: Response) -> None:
+            if not resp.is_success:
+                raise AQLCacheClearError(resp, request)
+
+        return self._execute(request, response_handler)
+
 
 class AQL(ApiGroup):
     """AQL (ArangoDB Query Language) API wrapper.
@@ -277,6 +307,7 @@ def execute(
         allow_dirty_read: bool = False,
         allow_retry: bool = False,
         force_one_shard_attribute_value: Optional[str] = None,
+        use_plan_cache: Optional[bool] = None,
     ) -> Result[Cursor]:
         """Execute the query and return the result cursor.
 
@@ -388,6 +419,8 @@ def execute(
             shipped to a wrong DB server and may not return results
             (i.e. empty result set). Use at your own risk.
         :param force_one_shard_attribute_value: str | None
+        :param use_plan_cache: If set to True, the query plan cache is used.
+        :param use_plan_cache: bool | None
         :return: Result cursor.
         :rtype: arango.cursor.Cursor
         :raise arango.exceptions.AQLQueryExecuteError: If execute fails.
@@ -399,8 +432,6 @@ def execute(
             data["ttl"] = ttl
         if bind_vars is not None:
             data["bindVars"] = bind_vars
-        if cache is not None:
-            data["cache"] = cache
         if memory_limit is not None:
             data["memoryLimit"] = memory_limit
 
@@ -437,6 +468,10 @@ def execute(
             options["allowRetry"] = allow_retry
         if force_one_shard_attribute_value is not None:
             options["forceOneShardAttributeValue"] = force_one_shard_attribute_value
+        if cache is not None:
+            options["cache"] = cache
+        if use_plan_cache is not None:
+            options["usePlanCache"] = use_plan_cache
 
         if options:
             data["options"] = options

tests/test_aql.py

@@ -1,5 +1,7 @@
+import pytest
 from packaging import version
 
+from arango.errno import FORBIDDEN
 from arango.exceptions import (
     AQLCacheClearError,
     AQLCacheConfigureError,
@@ -346,6 +348,86 @@ def test_aql_function_management(db, bad_db):
     assert db.aql.functions() == []
 
 
+def test_cache_results_management(db, bad_db, col, docs, cluster):
+    if cluster:
+        pytest.skip("Cluster mode does not support query result cache management")
+
+    aql = db.aql
+    cache = aql.cache
+
+    # Sanity check, just see if the response is OK.
+    _ = cache.properties()
+    with pytest.raises(AQLCachePropertiesError) as err:
+        _ = bad_db.aql.cache.properties()
+    assert err.value.error_code == FORBIDDEN
+
+    # Turn on caching
+    result = cache.configure(mode="on")
+    assert result["mode"] == "on"
+    result = cache.properties()
+    assert result["mode"] == "on"
+    with pytest.raises(AQLCacheConfigureError) as err:
+        _ = bad_db.aql.cache.configure(mode="on")
+    assert err.value.error_code == FORBIDDEN
+
+    # Run a simple query to use the cache
+    col.insert(docs[0])
+    _ = aql.execute(
+        query="FOR doc IN @@collection RETURN doc",
+        bind_vars={"@collection": col.name},
+        cache=True,
+    )
+
+    # Check the entries
+    entries = cache.entries()
+    assert isinstance(entries, list)
+    assert len(entries) > 0
+
+    with pytest.raises(AQLCacheEntriesError) as err:
+        _ = bad_db.aql.cache.entries()
+    assert err.value.error_code == FORBIDDEN
+
+    # Clear the cache
+    cache.clear()
+    entries = cache.entries()
+    assert len(entries) == 0
+    with pytest.raises(AQLCacheClearError) as err:
+        bad_db.aql.cache.clear()
+    assert err.value.error_code == FORBIDDEN
+
+
+def test_cache_plan_management(db, bad_db, col, docs, db_version):
+    if db_version < version.parse("3.12.4"):
+        pytest.skip("Query plan cache is supported in ArangoDB 3.12.4+")
+
+    aql = db.aql
+    cache = aql.cache
+
+    # Run a simple query to use the cache
+    col.insert(docs[0])
+    _ = aql.execute(
+        query="FOR doc IN @@collection RETURN doc",
+        bind_vars={"@collection": col.name},
+        use_plan_cache=True,
+    )
+
+    # Check the entries
+    entries = cache.plan_entries()
+    assert isinstance(entries, list)
+    assert len(entries) > 0
+    with pytest.raises(AQLCacheEntriesError) as err:
+        _ = bad_db.aql.cache.plan_entries()
+    assert err.value.error_code == FORBIDDEN
+
+    # Clear the cache
+    cache.clear_plan()
+    entries = cache.plan_entries()
+    assert len(entries) == 0
+    with pytest.raises(AQLCacheClearError) as err:
+        bad_db.aql.cache.clear_plan()
+    assert err.value.error_code == FORBIDDEN
+
+
 def test_aql_cache_management(db, bad_db):
     # Test get AQL cache properties
     properties = db.aql.cache.properties()