feat: implement query profiling (#542)

Author: daniel-sanche

datastore/samples/snippets/snippets.py

@@ -278,15 +278,16 @@ def sum_query_property_filter(client):
     # Execute sum aggregation query with filters
     completed_tasks = client.query(kind="Task").add_filter("done", "=", True)
     completed_tasks_query = client.aggregation_query(query=completed_tasks).sum(
-        property_ref="hours",
-        alias="total_completed_sum_hours"
+        property_ref="hours", alias="total_completed_sum_hours"
     )
 
     completed_query_result = completed_tasks_query.fetch()
     for aggregation_results in completed_query_result:
         for aggregation_result in aggregation_results:
             if aggregation_result.alias == "total_completed_sum_hours":
-                print(f"Total sum of hours in completed tasks is {aggregation_result.value}")
+                print(
+                    f"Total sum of hours in completed tasks is {aggregation_result.value}"
+                )
     # [END datastore_sum_aggregation_query_with_filters]
     return tasks
 
@@ -339,15 +340,16 @@ def avg_query_property_filter(client):
     # Execute average aggregation query with filters
     completed_tasks = client.query(kind="Task").add_filter("done", "=", True)
     completed_tasks_query = client.aggregation_query(query=completed_tasks).avg(
-        property_ref="hours",
-        alias="total_completed_avg_hours"
+        property_ref="hours", alias="total_completed_avg_hours"
     )
 
     completed_query_result = completed_tasks_query.fetch()
     for aggregation_results in completed_query_result:
         for aggregation_result in aggregation_results:
             if aggregation_result.alias == "total_completed_avg_hours":
-                print(f"Total average of hours in completed tasks is {aggregation_result.value}")
+                print(
+                    f"Total average of hours in completed tasks is {aggregation_result.value}"
+                )
     # [END datastore_avg_aggregation_query_with_filters]
     return tasks
 
@@ -375,9 +377,11 @@ def multiple_aggregations_query(client):
         [
             datastore.aggregation.CountAggregation(alias="count_aggregation"),
             datastore.aggregation.SumAggregation(
-                property_ref="hours", alias="sum_aggregation"),
+                property_ref="hours", alias="sum_aggregation"
+            ),
             datastore.aggregation.AvgAggregation(
-                property_ref="hours", alias="avg_aggregation")
+                property_ref="hours", alias="avg_aggregation"
+            ),
         ]
     )
 
@@ -389,6 +393,96 @@ def multiple_aggregations_query(client):
     return tasks
 
 
+def explain_analyze_entity(client):
+    # [START datastore_query_explain_analyze_entity]
+    # Build the query with explain_options
+    # analzye = true to get back the query stats, plan info, and query results
+    query = client.query(
+        kind="Task", explain_options=datastore.ExplainOptions(analyze=True)
+    )
+
+    # initiate the query
+    iterator = query.fetch()
+
+    # explain_metrics is only available after query is completed
+    for task_result in iterator:
+        print(task_result)
+
+    # get the plan summary
+    plan_summary = iterator.explain_metrics.plan_summary
+    print(f"Indexes used: {plan_summary.indexes_used}")
+
+    # get the execution stats
+    execution_stats = iterator.explain_metrics.execution_stats
+    print(f"Results returned: {execution_stats.results_returned}")
+    print(f"Execution duration: {execution_stats.execution_duration}")
+    print(f"Read operations: {execution_stats.read_operations}")
+    print(f"Debug stats: {execution_stats.debug_stats}")
+    # [END datastore_query_explain_analyze_entity]
+
+
+def explain_entity(client):
+    # [START datastore_query_explain_entity]
+    # Build the query with explain_options
+    # by default (analyze = false), only plan_summary property is available
+    query = client.query(kind="Task", explain_options=datastore.ExplainOptions())
+
+    # initiate the query
+    iterator = query.fetch()
+
+    # get the plan summary
+    plan_summary = iterator.explain_metrics.plan_summary
+    print(f"Indexes used: {plan_summary.indexes_used}")
+    # [END datastore_query_explain_entity]
+
+
+def explain_analyze_aggregation(client):
+    # [START datastore_query_explain_analyze_aggregation]
+    # Build the aggregation query with explain_options
+    # analzye = true to get back the query stats, plan info, and query results
+    all_tasks_query = client.query(kind="Task")
+    count_query = client.aggregation_query(
+        all_tasks_query, explain_options=datastore.ExplainOptions(analyze=True)
+    ).count()
+
+    # initiate the query
+    iterator = count_query.fetch()
+
+    # explain_metrics is only available after query is completed
+    for task_result in iterator:
+        print(task_result)
+
+    # get the plan summary
+    plan_summary = iterator.explain_metrics.plan_summary
+    print(f"Indexes used: {plan_summary.indexes_used}")
+
+    # get the execution stats
+    execution_stats = iterator.explain_metrics.execution_stats
+    print(f"Results returned: {execution_stats.results_returned}")
+    print(f"Execution duration: {execution_stats.execution_duration}")
+    print(f"Read operations: {execution_stats.read_operations}")
+    print(f"Debug stats: {execution_stats.debug_stats}")
+    # [END datastore_query_explain_analyze_aggregation]
+
+
+def explain_aggregation(client):
+    # [START datastore_query_explain_aggregation]
+    # Build the aggregation query with explain_options
+    # by default (analyze = false), only plan_summary property is available
+    all_tasks_query = client.query(kind="Task")
+    count_query = client.aggregation_query(
+        all_tasks_query, explain_options=datastore.ExplainOptions()
+    ).count()
+
+    # initiate the query
+    iterator = count_query.fetch()
+
+    # get the plan summary
+    plan_summary = iterator.explain_metrics.plan_summary
+    print(f"Indexes used: {plan_summary.indexes_used}")
+    # [END datastore_query_explain_aggregation]
+
+
 def main(project_id):
     client = datastore.Client(project_id)
 

datastore/samples/snippets/snippets_test.py

@@ -52,17 +52,15 @@ def setup_indexes(request):
 
     indexes = []
     done_property_index = datastore_admin_v1.Index.IndexedProperty(
-        name='done',
-        direction=datastore_admin_v1.Index.Direction.ASCENDING
+        name="done", direction=datastore_admin_v1.Index.Direction.ASCENDING
     )
     hour_property_index = datastore_admin_v1.Index.IndexedProperty(
-        name='hours',
-        direction=datastore_admin_v1.Index.Direction.ASCENDING
+        name="hours", direction=datastore_admin_v1.Index.Direction.ASCENDING
     )
     done_hour_index = datastore_admin_v1.Index(
-        kind='Task',
+        kind="Task",
         ancestor=datastore_admin_v1.Index.AncestorMode.NONE,
-        properties=[done_property_index, hour_property_index]
+        properties=[done_property_index, hour_property_index],
     )
     indexes.append(done_hour_index)
 
@@ -157,9 +155,7 @@ def test_count_query_with_stale_read(self, capsys, client):
     def test_sum_query_on_kind(self, capsys, client):
         tasks = snippets.sum_query_on_kind(client)
         captured = capsys.readouterr()
-        assert (
-            captured.out.strip() == "Total sum of hours in tasks is 9"
-        )
+        assert captured.out.strip() == "Total sum of hours in tasks is 9"
         assert captured.err == ""
 
         client.entities_to_delete.extend(tasks)
@@ -168,9 +164,7 @@ def test_sum_query_on_kind(self, capsys, client):
     def test_sum_query_property_filter(self, capsys, client):
         tasks = snippets.sum_query_property_filter(client)
         captured = capsys.readouterr()
-        assert (
-            captured.out.strip() == "Total sum of hours in completed tasks is 8"
-        )
+        assert captured.out.strip() == "Total sum of hours in completed tasks is 8"
         assert captured.err == ""
 
         client.entities_to_delete.extend(tasks)
@@ -179,9 +173,7 @@ def test_sum_query_property_filter(self, capsys, client):
     def test_avg_query_on_kind(self, capsys, client):
         tasks = snippets.avg_query_on_kind(client)
         captured = capsys.readouterr()
-        assert (
-            captured.out.strip() == "Total average of hours in tasks is 3.0"
-        )
+        assert captured.out.strip() == "Total average of hours in tasks is 3.0"
         assert captured.err == ""
 
         client.entities_to_delete.extend(tasks)
@@ -201,15 +193,57 @@ def test_avg_query_property_filter(self, capsys, client):
     def test_multiple_aggregations_query(self, capsys, client):
         tasks = snippets.multiple_aggregations_query(client)
         captured = capsys.readouterr()
+        assert "avg_aggregation value is 3.0" in captured.out
+        assert "count_aggregation value is 3" in captured.out
+        assert "sum_aggregation value is 9" in captured.out
+        assert captured.err == ""
+
+        client.entities_to_delete.extend(tasks)
+
+    @backoff.on_exception(backoff.expo, AssertionError, max_time=240)
+    def test_explain_analyze_entity(self, capsys, client):
+        snippets.explain_analyze_entity(client)
+        captured = capsys.readouterr()
         assert (
-            'avg_aggregation value is 3.0' in captured.out
+            "Indexes used: [{'properties': '(__name__ ASC)', 'query_scope': 'Collection group'}]"
+            in captured.out
         )
+        assert "Results returned: 0" in captured.out
+        assert "Execution duration: 0:00" in captured.out
+        assert "Read operations: 0" in captured.out
+        assert "Debug stats: {" in captured.out
+        assert captured.err == ""
+
+    @backoff.on_exception(backoff.expo, AssertionError, max_time=240)
+    def test_explain_entity(self, capsys, client):
+        snippets.explain_entity(client)
+        captured = capsys.readouterr()
         assert (
-            'count_aggregation value is 3' in captured.out
+            "Indexes used: [{'properties': '(__name__ ASC)', 'query_scope': 'Collection group'}]"
+            in captured.out
         )
+        assert captured.err == ""
+
+    @backoff.on_exception(backoff.expo, AssertionError, max_time=240)
+    def test_explain_analyze_aggregation(self, capsys, client):
+        snippets.explain_analyze_aggregation(client)
+        captured = capsys.readouterr()
         assert (
-            'sum_aggregation value is 9' in captured.out
+            "Indexes used: [{'properties': '(__name__ ASC)', 'query_scope': 'Collection group'}]"
+            in captured.out
         )
+        assert "Results returned: 1" in captured.out
+        assert "Execution duration: 0:00" in captured.out
+        assert "Read operations: 1" in captured.out
+        assert "Debug stats: {" in captured.out
         assert captured.err == ""
 
-        client.entities_to_delete.extend(tasks)
+    @backoff.on_exception(backoff.expo, AssertionError, max_time=240)
+    def test_explain_aggregation(self, capsys, client):
+        snippets.explain_aggregation(client)
+        captured = capsys.readouterr()
+        assert (
+            "Indexes used: [{'properties': '(__name__ ASC)', 'query_scope': 'Collection group'}]"
+            in captured.out
+        )
+        assert captured.err == ""