From a07bb7f7cbf4ca14ea1ca70af4e203d1cefee05b Mon Sep 17 00:00:00 2001 From: Fatih Acar Date: Tue, 7 Jan 2025 19:40:45 +0100 Subject: [PATCH] fix(backend): do not cleanup labels for each db result row There is a few microseconds overhead for this function... Also add a span for Query.execute in order to know the overhead of fetching the results. Signed-off-by: Fatih Acar --- backend/infrahub/core/query/__init__.py | 7 +++++-- backend/tests/unit/core/test_query.py | 20 +++++++++++--------- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/backend/infrahub/core/query/__init__.py b/backend/infrahub/core/query/__init__.py index 095d28fcc1..37ededf594 100644 --- a/backend/infrahub/core/query/__init__.py +++ b/backend/infrahub/core/query/__init__.py @@ -10,6 +10,7 @@ from neo4j.graph import Node as Neo4jNode from neo4j.graph import Path as Neo4jPath from neo4j.graph import Relationship as Neo4jRelationship +from opentelemetry import trace from infrahub import config from infrahub.core.constants import PermissionLevel @@ -161,7 +162,7 @@ def cleanup_return_labels(labels: list[str]) -> list[str]: class QueryResult: def __init__(self, data: list[Union[Neo4jNode, Neo4jRelationship, list[Neo4jNode]]], labels: list[str]): self.data = data - self.labels = cleanup_return_labels(labels) + self.labels = labels self.branch_score: int = 0 self.time_score: int = 0 self.permission_score = PermissionLevel.DEFAULT @@ -523,6 +524,7 @@ def _get_params_for_neo4j_shell(self) -> str: return ":params { " + ", ".join(params) + " }" + @trace.get_tracer(__name__).start_as_current_span("Query.execute") async def execute(self, db: InfrahubDatabase) -> Self: # Ensure all mandatory params have been provided # Ensure at least 1 return obj has been defined @@ -552,7 +554,8 @@ async def execute(self, db: InfrahubDatabase) -> Self: if not results and self.raise_error_if_empty: raise QueryError(query=query_str, params=self.params) - self.results = [QueryResult(data=result, labels=self.return_labels) for result in results] + clean_labels = cleanup_return_labels(self.return_labels) + self.results = [QueryResult(data=result, labels=clean_labels) for result in results] self.has_been_executed = True return self diff --git a/backend/tests/unit/core/test_query.py b/backend/tests/unit/core/test_query.py index 5405af6337..80138f9e96 100644 --- a/backend/tests/unit/core/test_query.py +++ b/backend/tests/unit/core/test_query.py @@ -174,12 +174,14 @@ async def test_query_result_getters(neo4j_factory): qr = QueryResult( data=[n1, r1, r2, n2], - labels=[ - "n1", - "r1", - "r2", - "n2", - ], + labels=cleanup_return_labels( + [ + "n1", + "r1", + "r2", + "n2", + ] + ), ) assert list(qr.get_rels()) == [r1, r2] assert list(qr.get_nodes()) == [n1, n2] @@ -228,9 +230,9 @@ async def test_sort_results_by_time(neo4j_factory): }, ) - qr1 = QueryResult(data=[n1, n2, r1], labels=["n1", "n2", "r"]) - qr2 = QueryResult(data=[n1, n2, r2], labels=["n1", "n2", "r"]) - qr3 = QueryResult(data=[n1, n2, r3], labels=["n1", "n2", "r"]) + qr1 = QueryResult(data=[n1, n2, r1], labels=cleanup_return_labels(["n1", "n2", "r"])) + qr2 = QueryResult(data=[n1, n2, r2], labels=cleanup_return_labels(["n1", "n2", "r"])) + qr3 = QueryResult(data=[n1, n2, r3], labels=cleanup_return_labels(["n1", "n2", "r"])) results = sort_results_by_time(results=[qr1, qr2, qr3], rel_label="r") assert list(results) == [qr3, qr1, qr2]