ES-2276 | support level 2 query profiling (#355)

aMahanna · web-flow · commit 0d56bcfea46c · 2024-12-10T09:03:45.000-05:00
* ES-2276 | initial commit

* fix: `profile` typing
diff --git a/arango/aql.py b/arango/aql.py
@@ -264,7 +264,7 @@ def execute(
         cache: Optional[bool] = None,
         memory_limit: int = 0,
         fail_on_warning: Optional[bool] = None,
-        profile: Optional[bool] = None,
+        profile: Optional[Union[bool, int]] = None,
         max_transaction_size: Optional[int] = None,
         max_warning_count: Optional[int] = None,
         intermediate_commit_count: Optional[int] = None,
@@ -317,8 +317,12 @@ def execute(
             this behaviour, so it does not need to be set per-query.
         :type fail_on_warning: bool
         :param profile: Return additional profiling details in the cursor,
-            unless the query cache is used.
-        :type profile: bool
+            unless the query cache is used. If set to True or 1, then query profiling
+            information can be fetched with `cursor.profile()`. If set to 2, additional
+            execution stats per query plan node are included via "nodes" in
+            `cursor.statistics()`, as well as a the query plan which can be fetched
+            with `cursor.plan()`.
+        :type profile: bool | int
         :param max_transaction_size: Transaction size limit in bytes.
         :type max_transaction_size: int
         :param max_warning_count: Max number of warnings returned.
diff --git a/arango/cursor.py b/arango/cursor.py
@@ -40,6 +40,7 @@ class Cursor:
         "_count",
         "_cached",
         "_stats",
+        "_plan",
         "_profile",
         "_warnings",
         "_has_more",
@@ -63,6 +64,7 @@ def __init__(
         self._count: Optional[int] = None
         self._cached = None
         self._stats = None
+        self._plan = None
         self._profile = None
         self._warnings = None
         self._next_batch_id: Optional[str] = None
@@ -132,12 +134,18 @@ def _update(self, data: Json) -> Json:
                 self._warnings = extra["warnings"]
                 result["warnings"] = extra["warnings"]
 
+            if "plan" in extra:
+                self._plan = extra["plan"]
+                result["plan"] = extra["plan"]
+
             if "stats" in extra:
                 stats = extra["stats"]
                 if "writesExecuted" in stats:
                     stats["modified"] = stats.pop("writesExecuted")
                 if "writesIgnored" in stats:
                     stats["ignored"] = stats.pop("writesIgnored")
+                if "documentLookups" in stats:
+                    stats["lookups"] = stats.pop("documentLookups")
                 if "scannedFull" in stats:
                     stats["scanned_full"] = stats.pop("scannedFull")
                 if "scannedIndex" in stats:
@@ -159,6 +167,9 @@ def _update(self, data: Json) -> Json:
                 if "peakMemoryUsage" in stats:
                     stats["peak_memory_usage"] = stats.pop("peakMemoryUsage")
 
+                if "intermediateCommits" in stats:
+                    stats["intermediate_commits"] = stats.pop("intermediateCommits")
+
                 self._stats = stats
                 result["statistics"] = stats
 
@@ -239,6 +250,14 @@ def warnings(self) -> Optional[Sequence[Json]]:
         """
         return self._warnings
 
+    def plan(self) -> Optional[Json]:
+        """Return query execution plan.
+
+        :return: Query execution plan.
+        :rtype: dict
+        """
+        return self._plan
+
     def empty(self) -> bool:
         """Check if the current batch is empty.
 
diff --git a/arango/request.py b/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 = "8.1.1"
+    driver_version = "8.2.0"
     driver_header = "python-arango/" + driver_version + " (" + flags + ")"
     normalized_headers: Headers = {
         "charset": "utf-8",
diff --git a/tests/test_cursor.py b/tests/test_cursor.py
@@ -22,7 +22,7 @@ def test_cursor_from_execute_query(db, col, docs):
         batch_size=2,
         ttl=1000,
         optimizer_rules=["+all"],
-        profile=True,
+        profile=2,
     )
     cursor_id = cursor.id
     assert "Cursor" in repr(cursor)
@@ -41,12 +41,25 @@ def test_cursor_from_execute_query(db, col, docs):
     assert "http_requests" in statistics
     assert "scanned_full" in statistics
     assert "scanned_index" in statistics
+    assert "nodes" in statistics
+
     assert cursor.warnings() == []
 
     profile = cursor.profile()
     assert profile["initializing"] > 0
     assert profile["parsing"] > 0
 
+    plan = cursor.plan()
+    assert set(plan.keys()) == {
+        "nodes",
+        "rules",
+        "collections",
+        "variables",
+        "estimatedCost",
+        "estimatedNrItems",
+        "isModificationQuery",
+    }
+
     assert clean_doc(cursor.next()) == docs[0]
     assert cursor.id == cursor_id
     assert cursor.has_more() is True
@@ -106,7 +119,7 @@ def test_cursor_write_query(db, col, docs):
         batch_size=1,
         ttl=1000,
         optimizer_rules=["+all"],
-        profile=True,
+        profile=1,
         max_runtime=0.0,
     )
     cursor_id = cursor.id
