Merge pull request #4504 from opsmill/lgu-perf-benchmarks

Instantiate DB queries benchmarking framework

Author: LucasG0

backend/infrahub/core/query/__init__.py

@@ -14,7 +14,6 @@
 from infrahub import config
 from infrahub.core.constants import PermissionLevel
 from infrahub.core.timestamp import Timestamp
-from infrahub.database.constants import DatabaseType, Neo4jRuntime
 from infrahub.exceptions import QueryError
 
 if TYPE_CHECKING:
@@ -24,6 +23,7 @@
     from infrahub.core.branch import Branch
     from infrahub.database import InfrahubDatabase
 
+
 RETURN_TYPE = TypeVar("RETURN_TYPE")
 
 
@@ -518,9 +518,7 @@ def _get_params_for_neo4j_shell(self) -> str:
 
         return ":params { " + ", ".join(params) + " }"
 
-    async def execute(
-        self, db: InfrahubDatabase, profile: bool = False, runtime: Neo4jRuntime = Neo4jRuntime.DEFAULT
-    ) -> Self:
+    async def execute(self, db: InfrahubDatabase) -> Self:
         # Ensure all mandatory params have been provided
         # Ensure at least 1 return obj has been defined
 
@@ -529,12 +527,6 @@ async def execute(
 
         query_str = self.get_query()
 
-        if profile:
-            query_str = "PROFILE\n" + query_str
-
-        if runtime != Neo4jRuntime.DEFAULT and db.db_type == DatabaseType.NEO4J:
-            query_str = f"CYPHER runtime={runtime.value}\n" + query_str
-
         if self.type == QueryType.READ:
             if self.limit or self.offset:
                 results = await db.execute_query(query=query_str, params=self.params, name=self.name)

backend/infrahub/database/__init__.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import random
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypeVar, Union
 
 from neo4j import (
@@ -30,7 +31,7 @@
 from infrahub.log import get_logger
 from infrahub.utils import InfrahubStringEnum
 
-from .constants import DatabaseType
+from .constants import DatabaseType, Neo4jRuntime
 from .memgraph import DatabaseManagerMemgraph
 from .metrics import QUERY_EXECUTION_METRICS, TRANSACTION_RETRIES
 from .neo4j import DatabaseManagerNeo4j
@@ -50,6 +51,12 @@
 log = get_logger()
 
 
+@dataclass
+class QueryConfig:
+    neo4j_runtime: Neo4jRuntime = Neo4jRuntime.DEFAULT
+    profile_memory: bool = False
+
+
 class InfrahubDatabaseMode(InfrahubStringEnum):
     DRIVER = "driver"
     SESSION = "session"
@@ -134,13 +141,15 @@ def __init__(
         session: Optional[AsyncSession] = None,
         session_mode: InfrahubDatabaseSessionMode = InfrahubDatabaseSessionMode.WRITE,
         transaction: Optional[AsyncTransaction] = None,
-    ) -> None:
+        queries_names_to_config: Optional[dict[str, QueryConfig]] = None,
+    ):
         self._mode: InfrahubDatabaseMode = mode
         self._driver: AsyncDriver = driver
         self._session: Optional[AsyncSession] = session
         self._session_mode: InfrahubDatabaseSessionMode = session_mode
         self._is_session_local: bool = False
         self._transaction: Optional[AsyncTransaction] = transaction
+        self.queries_names_to_config = queries_names_to_config if queries_names_to_config is not None else {}
 
         if schemas:
             self._schemas: dict[str, SchemaBranch] = {schema.name: schema for schema in schemas}
@@ -189,6 +198,7 @@ def start_session(self, read_only: bool = False, schemas: Optional[list[SchemaBr
             db_manager=self.manager,
             driver=self._driver,
             session_mode=session_mode,
+            queries_names_to_config=self.queries_names_to_config,
         )
 
     def start_transaction(self, schemas: Optional[list[SchemaBranch]] = None) -> InfrahubDatabase:
@@ -200,6 +210,7 @@ def start_transaction(self, schemas: Optional[list[SchemaBranch]] = None) -> Inf
             driver=self._driver,
             session=self._session,
             session_mode=self._session_mode,
+            queries_names_to_config=self.queries_names_to_config,
         )
 
     async def session(self) -> AsyncSession:
@@ -274,14 +285,8 @@ async def close(self) -> None:
     async def execute_query(
         self, query: str, params: Optional[dict[str, Any]] = None, name: Optional[str] = "undefined"
     ) -> list[Record]:
-        with trace.get_tracer(__name__).start_as_current_span("execute_db_query") as span:
-            span.set_attribute("query", query)
-            if name:
-                span.set_attribute("query_name", name)
-
-            with QUERY_EXECUTION_METRICS.labels(self._session_mode.value, name).time():
-                response = await self.run_query(query=query, params=params)
-                return [item async for item in response]
+        results, _ = await self.execute_query_with_metadata(query=query, params=params, name=name)
+        return results
 
     async def execute_query_with_metadata(
         self, query: str, params: Optional[dict[str, Any]] = None, name: Optional[str] = "undefined"
@@ -291,6 +296,17 @@ async def execute_query_with_metadata(
             if name:
                 span.set_attribute("query_name", name)
 
+            try:
+                query_config = self.queries_names_to_config[name]
+                if self.db_type == DatabaseType.NEO4J:
+                    runtime = self.queries_names_to_config[name].neo4j_runtime
+                    if runtime != Neo4jRuntime.DEFAULT:
+                        query = f"CYPHER runtime = {runtime.value}\n" + query
+                if query_config.profile_memory:
+                    query = "PROFILE\n" + query
+            except KeyError:
+                pass  # No specific config for this query
+
             with QUERY_EXECUTION_METRICS.labels(self._session_mode.value, name).time():
                 response = await self.run_query(query=query, params=params, name=name)
                 results = [item async for item in response]

backend/infrahub/database/analyzer.py

@@ -11,6 +11,7 @@ class Neo4jRuntime(str, Enum):
     INTERPRETED = "interpreted"
     SLOTTED = "slotted"
     PIPELINED = "pipelined"
+    PARALLEL = "parallel"
 
 
 class IndexType(str, Enum):