Merge pull request #251 from ArangoDB-Community/feature/de-546-stats-field-in-explain-api

[DE-546] Adding stats fields to the explain API

Author: CryptoNinjaGeek

CHANGELOG.md

@@ -0,0 +1,7 @@
+main
+----
+
+* Adding peak_memory_usage as a new property of AQL queries, available since ArangoDB 3.11.
+
+* The explain method of AQL queries includes the "stats" field in the returned object. Note that the REST API returns
+  it separately from the "plan" field, but for now we have to merge them together to ensure backward compatibility.

arango/aql.py

@@ -213,11 +213,19 @@ def response_handler(resp: Response) -> Union[Json, Jsons]:
             if not resp.is_success:
                 raise AQLQueryExplainError(resp, request)
             if "plan" in resp.body:
-                plan: Json = resp.body["plan"]
-                return plan
+                result: Json = resp.body["plan"]
+                if "stats" in resp.body:
+                    result["stats"] = resp.body["stats"]
+                return result
             else:
-                plans: Jsons = resp.body["plans"]
-                return plans
+                results: Jsons = resp.body["plans"]
+                if "stats" in resp.body:
+                    # Although "plans" contains an array, "stats" is a single object.
+                    # We need to duplicate "stats" for each plan in order to preserve
+                    # the original structure.
+                    for plan in results:
+                        plan["stats"] = resp.body["stats"]
+                return results
 
         return self._execute(request, response_handler)
 

arango/formatter.py

@@ -322,6 +322,10 @@ def format_aql_query(body: Json) -> Json:
         result["stream"] = body["stream"]
     if "user" in body:
         result["user"] = body["user"]
+
+    # New in 3.11
+    if "peakMemoryUsage" in body:
+        result["peak_memory_usage"] = body["peakMemoryUsage"]
     return verify_format(body, result)
 
 

arango/request.py

@@ -12,7 +12,7 @@ def normalize_headers(
     if driver_flags is not None:
         for flag in driver_flags:
             flags = flags + flag + ";"
-    driver_version = "7.5.3"
+    driver_version = "7.5.8"
     driver_header = "python-arango/" + driver_version + " (" + flags + ")"
     normalized_headers: Headers = {
         "charset": "utf-8",

setup.py

@@ -24,6 +24,7 @@
         "PyJWT",
         "setuptools>=42",
         "importlib_metadata>=4.7.1",
+        "packaging>=23.1",
     ],
     extras_require={
         "dev": [

tester.sh

@@ -21,7 +21,10 @@ if [[ "$tests" != "all" && "$tests" != "community" && "$tests" != "enterprise" ]
     exit 1
 fi
 
-version="${3:-3.10.6}"
+# 3.11.0
+# 3.10.6
+# 3.9.9
+version="${3:-3.11.0}"
 
 if [[ -n "$4" && "$4" != "notest" ]]; then
     echo "Invalid argument. Use 'notest' to only start the docker container, without running the tests."

tests/conftest.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 
 import pytest
+from packaging import version
 
 from arango import ArangoClient, formatter
 from arango.database import StandardDatabase
@@ -40,6 +41,7 @@ class GlobalData:
     enterprise: bool = None
     secret: str = None
     root_password: str = None
+    db_version: version = version.parse("0.0.0")
 
 
 global_data = GlobalData()
@@ -66,7 +68,7 @@ def pytest_configure(config):
         password=config.getoption("passwd"),
         superuser_token=generate_jwt(secret),
     )
-    sys_db.version()
+    db_version = sys_db.version()
 
     # Create a user and non-system database for testing.
     username = generate_username()
@@ -118,6 +120,7 @@ def pytest_configure(config):
     global_data.username = username
     global_data.password = password
     global_data.db_name = tst_db_name
+    global_data.db_version = version.parse(db_version)
     global_data.sys_db = sys_db
     global_data.tst_db = tst_db
     global_data.bad_db = bad_db
@@ -247,6 +250,11 @@ def mock_verify_format(body, result):
     monkeypatch.setattr(formatter, "verify_format", mock_verify_format)
 
 
+@pytest.fixture(autouse=False)
+def db_version():
+    return global_data.db_version
+
+
 @pytest.fixture(autouse=False)
 def url():
     return global_data.url

tests/test_aql.py

@@ -1,3 +1,5 @@
+from packaging import version
+
 from arango.exceptions import (
     AQLCacheClearError,
     AQLCacheConfigureError,
@@ -26,37 +28,60 @@ def test_aql_attributes(db, username):
     assert repr(db.aql.cache) == f"<AQLQueryCache in {db.name}>"
 
 
-def test_aql_query_management(db, bad_db, col, docs):
-    plan_fields = [
+def test_aql_query_management(db_version, db, bad_db, col, docs):
+    explain_fields = [
         "estimatedNrItems",
         "estimatedCost",
         "rules",
         "variables",
         "collections",
+        "stats",
     ]
+    stats_fields = {
+        "0.0.0": [
+            "rulesExecuted",
+            "rulesSkipped",
+            "plansCreated",
+        ],
+        "3.10.4": [
+            "peakMemoryUsage",
+            "executionTime",
+        ],
+    }
+
     # Test explain invalid query
     with assert_raises(AQLQueryExplainError) as err:
         db.aql.explain("INVALID QUERY")
     assert err.value.error_code == 1501
 
     # Test explain valid query with all_plans set to False
-    plan = db.aql.explain(
+    explain = db.aql.explain(
         f"FOR d IN {col.name} RETURN d",
         all_plans=False,
         opt_rules=["-all", "+use-index-range"],
     )
-    assert all(field in plan for field in plan_fields)
+    assert all(field in explain for field in explain_fields)
+    for v, fields in stats_fields.items():
+        if db_version >= version.parse(v):
+            assert all(field in explain["stats"] for field in fields)
+        else:
+            assert all(field not in explain["stats"] for field in fields)
 
     # Test explain valid query with all_plans set to True
-    plans = db.aql.explain(
+    explanations = db.aql.explain(
         f"FOR d IN {col.name} RETURN d",
         all_plans=True,
         opt_rules=["-all", "+use-index-range"],
         max_plans=10,
     )
-    for plan in plans:
-        assert all(field in plan for field in plan_fields)
-    assert len(plans) < 10
+    for explain in explanations:
+        assert all(field in explain for field in explain_fields)
+        for v, fields in stats_fields.items():
+            if db_version >= version.parse(v):
+                assert all(field in explain["stats"] for field in fields)
+            else:
+                assert all(field not in explain["stats"] for field in fields)
+    assert len(explanations) < 10
 
     # Test validate invalid query
     with assert_raises(AQLQueryValidateError) as err:
@@ -161,7 +186,7 @@ def test_aql_query_management(db, bad_db, col, docs):
     assert new_tracking["track_bind_vars"] is True
     assert new_tracking["track_slow_queries"] is True
 
-    # Kick off some long lasting queries in the background
+    # Kick off some long-lasting queries in the background
     db.begin_async_execution().aql.execute("RETURN SLEEP(100)")
     db.begin_async_execution().aql.execute("RETURN SLEEP(50)")
 
@@ -174,6 +199,8 @@ def test_aql_query_management(db, bad_db, col, docs):
         assert "state" in query
         assert "bind_vars" in query
         assert "runtime" in query
+        if db_version >= version.parse("3.11"):
+            assert "peak_memory_usage" in query
     assert len(queries) == 2
 
     # Test list queries with bad database