Merge pull request #4628 from opsmill/dga-20241015-queryanalyzer

Add support for context to InfrahubDatabase

Author: LucasG0

backend/infrahub/database/__init__.py

@@ -182,6 +182,14 @@ def is_transaction(self) -> bool:
             return True
         return False
 
+    def get_context(self) -> dict[str, Any]:
+        """
+        This method is meant to be overridden by subclasses in order to fill in subclass attributes
+        to methods returning a copy of this object using self.__class__ constructor.
+        """
+
+        return {}
+
     def add_schema(self, schema: SchemaBranch, name: Optional[str] = None) -> None:
         self._schemas[name or schema.name] = schema
 
@@ -191,6 +199,8 @@ def start_session(self, read_only: bool = False, schemas: Optional[list[SchemaBr
         if read_only:
             session_mode = InfrahubDatabaseSessionMode.READ
 
+        context = self.get_context()
+
         return self.__class__(
             mode=InfrahubDatabaseMode.SESSION,
             db_type=self.db_type,
@@ -199,9 +209,12 @@ def start_session(self, read_only: bool = False, schemas: Optional[list[SchemaBr
             driver=self._driver,
             session_mode=session_mode,
             queries_names_to_config=self.queries_names_to_config,
+            **context,
         )
 
     def start_transaction(self, schemas: Optional[list[SchemaBranch]] = None) -> InfrahubDatabase:
+        context = self.get_context()
+
         return self.__class__(
             mode=InfrahubDatabaseMode.TRANSACTION,
             db_type=self.db_type,
@@ -211,6 +224,7 @@ def start_transaction(self, schemas: Optional[list[SchemaBranch]] = None) -> Inf
             session=self._session,
             session_mode=self._session_mode,
             queries_names_to_config=self.queries_names_to_config,
+            **context,
         )
 
     async def session(self) -> AsyncSession:

backend/tests/helpers/query_benchmark/data_generator.py

@@ -5,7 +5,7 @@
 from rich.console import Console
 from rich.progress import Progress
 
-from tests.helpers.query_benchmark.db_query_profiler import InfrahubDatabaseProfiler, ProfilerEnabler, query_analyzer
+from tests.helpers.query_benchmark.db_query_profiler import InfrahubDatabaseProfiler, ProfilerEnabler, QueryAnalyzer
 
 
 class DataGenerator:
@@ -33,6 +33,7 @@ async def load_data_and_profile(
     profile_frequency: int,
     graphs_output_location: Path,
     test_label: str,
+    query_analyzer: QueryAnalyzer,
     memory_profiling_rate: int = 25,
 ) -> None:
     """
@@ -67,7 +68,7 @@ async def load_data_and_profile(
             await data_generator.load_data(nb_elements=nb_elem_to_load)
             query_analyzer.increase_nb_elements_loaded(profile_frequency)
             profile_memory = i % memory_profiling_rate == 0 if memory_profiling_rate is not None else False
-            with ProfilerEnabler(profile_memory=profile_memory):
+            with ProfilerEnabler(profile_memory=profile_memory, query_analyzer=query_analyzer):
                 await func_call()
             progress.advance(task)
 

backend/tests/helpers/query_benchmark/db_query_profiler.py

@@ -147,35 +147,46 @@ def create_memory_graph(self, query_name: str, label: str, output_dir: Path) ->
 class ProfilerEnabler:
     profile_memory: bool
 
-    def __init__(self, profile_memory: bool) -> None:
+    def __init__(self, profile_memory: bool, query_analyzer: QueryAnalyzer) -> None:
         self.profile_memory = profile_memory
+        self.query_analyzer = query_analyzer
 
     def __enter__(self) -> None:
-        query_analyzer.profile_duration = True
-        query_analyzer.profile_memory = self.profile_memory
+        self.query_analyzer.profile_duration = True
+        self.query_analyzer.profile_memory = self.profile_memory
 
     def __exit__(
         self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
     ) -> None:
-        query_analyzer.profile_duration = False
-        query_analyzer.profile_memory = False
+        self.query_analyzer.profile_duration = False
+        self.query_analyzer.profile_memory = False
 
 
 # Tricky to have it as an attribute of InfrahubDatabaseProfiler as some copies of InfrahubDatabase are made
 # during start_session calls.
-query_analyzer = QueryAnalyzer()
+# query_analyzer = QueryAnalyzer()
 
 
 class InfrahubDatabaseProfiler(InfrahubDatabase):
+    def __init__(self, query_analyzer: QueryAnalyzer, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.query_analyzer = query_analyzer
+        # Note that any attribute added here should be added to get_context method.
+
+    def get_context(self) -> dict[str, Any]:
+        ctx = super().get_context()
+        ctx["query_analyzer"] = self.query_analyzer
+        return ctx
+
     async def execute_query_with_metadata(
         self, query: str, params: dict[str, Any] | None = None, name: str | None = "undefined"
     ) -> tuple[list[Record], dict[str, Any]]:
-        if not query_analyzer.profile_duration:
+        if not self.query_analyzer.profile_duration:
             # Profiling might be disabled to avoid capturing queries while loading data
             return await super().execute_query_with_metadata(query, params, name)
 
         # We don't want to memory profile all queries
-        if query_analyzer.profile_memory and name in self.queries_names_to_config:
+        if self.query_analyzer.profile_memory and name in self.queries_names_to_config:
             # Following call to super().execute_query_with_metadata() will use this value to set PROFILE option
             self.queries_names_to_config[name].profile_memory = True
             profile_memory = True
@@ -193,6 +204,6 @@ async def execute_query_with_metadata(
             query_name=str(name),
             start_time=time_start,
         )
-        query_analyzer.add_measurement(measurement)
+        self.query_analyzer.add_measurement(measurement)
 
         return response, metadata

backend/tests/query_benchmark/conftest.py

@@ -1,6 +1,12 @@
 import pytest
 
 from infrahub.core.schema import SchemaRoot
+from tests.helpers.query_benchmark.db_query_profiler import QueryAnalyzer
+
+
+@pytest.fixture(scope="session")
+def query_analyzer() -> QueryAnalyzer:
+    return QueryAnalyzer()
 
 
 @pytest.fixture

backend/tests/query_benchmark/test_node_manager_query.py

@@ -23,8 +23,6 @@
 
 log = get_logger()
 
-pytestmark = pytest.mark.skip("Not relevant to test this currently.")
-
 
 @pytest.mark.parametrize(
     "neo4j_image, neo4j_runtime",
@@ -43,14 +41,14 @@
         ),
     ],
 )
-async def test_query_persons(neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root):
+async def test_query_persons(query_analyzer, neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root):
     queries_names_to_config = {
         NodeGetListQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
         NodeListGetAttributeQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
         NodeListGetInfoQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
     }
     db_profiling_queries, default_branch = await start_db_and_create_default_branch(
-        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config
+        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config, query_analyzer=query_analyzer
     )
 
     registry.schema.register_schema(schema=car_person_schema_root, branch=default_branch.name)
@@ -67,6 +65,7 @@ async def test_query_persons(neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_
         profile_frequency=50,
         nb_elements=1000,
         graphs_output_location=graph_output_location,
+        query_analyzer=query_analyzer,
         test_label=f" data: {neo4j_image}" + f" runtime: {neo4j_runtime}",
     )
 
@@ -88,14 +87,16 @@ async def test_query_persons(neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_
         ),
     ],
 )
-async def test_query_persons_with_isolated_cars(neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root):
+async def test_query_persons_with_isolated_cars(
+    query_analyzer, neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root
+):
     queries_names_to_config = {
         NodeGetListQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
         NodeListGetAttributeQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
         NodeListGetInfoQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
     }
     db_profiling_queries, default_branch = await start_db_and_create_default_branch(
-        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config
+        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config, query_analyzer=query_analyzer
     )
 
     registry.schema.register_schema(schema=car_person_schema_root, branch=default_branch.name)
@@ -117,6 +118,7 @@ async def test_query_persons_with_isolated_cars(neo4j_image: str, neo4j_runtime:
         profile_frequency=50,
         nb_elements=1000,
         graphs_output_location=graph_output_location,
+        query_analyzer=query_analyzer,
         test_label=f" data: {neo4j_image}" + f" runtime: {neo4j_runtime}",
     )
 
@@ -138,14 +140,16 @@ async def test_query_persons_with_isolated_cars(neo4j_image: str, neo4j_runtime:
         ),
     ],
 )
-async def test_query_persons_with_connected_cars(neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root):
+async def test_query_persons_with_connected_cars(
+    query_analyzer, neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root
+):
     queries_names_to_config = {
         NodeGetListQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
         NodeListGetAttributeQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
         NodeListGetInfoQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime),
     }
     db_profiling_queries, default_branch = await start_db_and_create_default_branch(
-        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config
+        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config, query_analyzer=query_analyzer
     )
 
     registry.schema.register_schema(schema=car_person_schema_root, branch=default_branch.name)
@@ -163,5 +167,6 @@ async def test_query_persons_with_connected_cars(neo4j_image: str, neo4j_runtime
         profile_frequency=50,
         nb_elements=1000,
         graphs_output_location=graph_output_location,
+        query_analyzer=query_analyzer,
         test_label=f" data: {neo4j_image}" + f" runtime: {neo4j_runtime}",
     )

backend/tests/query_benchmark/test_node_unique_attribute_constraint.py

@@ -15,14 +15,13 @@
     CarGenerator,
 )
 from tests.helpers.query_benchmark.data_generator import load_data_and_profile
-from tests.helpers.utils import start_db_and_create_default_branch
+
+from .utils import start_db_and_create_default_branch
 
 RESULTS_FOLDER = Path(__file__).resolve().parent / "query_performance_results"
 
 log = get_logger()
 
-pytestmark = pytest.mark.skip("Not relevant to test this currently.")
-
 
 @pytest.mark.parametrize(
     "neo4j_image, neo4j_runtime",
@@ -42,11 +41,11 @@
     ],
 )
 async def test_query_unique_cars_single_attribute(
-    neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root
+    query_analyzer, neo4j_image: str, neo4j_runtime: Neo4jRuntime, car_person_schema_root
 ):
     queries_names_to_config = {NodeUniqueAttributeConstraintQuery.name: QueryConfig(neo4j_runtime=neo4j_runtime)}
     db_profiling_queries, default_branch = await start_db_and_create_default_branch(
-        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config
+        neo4j_image=neo4j_image, queries_names_to_config=queries_names_to_config, query_analyzer=query_analyzer
     )
 
     # Register schema
@@ -70,5 +69,6 @@ async def test_query_unique_cars_single_attribute(
         profile_frequency=50,
         nb_elements=1000,
         graphs_output_location=graph_output_location,
+        query_analyzer=query_analyzer,
         test_label=f" data: {neo4j_image}" + f" runtime: {neo4j_runtime}",
     )

backend/tests/query_benchmark/utils.py

@@ -6,21 +6,22 @@
 from infrahub.core.initialization import create_default_branch, create_global_branch, create_root_node
 from infrahub.core.schema.manager import SchemaManager
 from infrahub.database import InfrahubDatabaseMode, QueryConfig, get_db
-from infrahub.lock import initialize_lock
 from tests.helpers.constants import PORT_BOLT_NEO4J
-from tests.helpers.query_benchmark.db_query_profiler import InfrahubDatabaseProfiler
+from tests.helpers.query_benchmark.db_query_profiler import InfrahubDatabaseProfiler, QueryAnalyzer
 from tests.helpers.utils import start_neo4j_container
 
 
 async def start_db_and_create_default_branch(
-    neo4j_image: str, queries_names_to_config: Optional[dict[str, QueryConfig]] = None
+    neo4j_image: str, query_analyzer: QueryAnalyzer, queries_names_to_config: Optional[dict[str, QueryConfig]] = None
 ) -> Tuple[InfrahubDatabaseProfiler, Branch]:
     neo4j_container = start_neo4j_container(neo4j_image)
     config.SETTINGS.database.port = int(neo4j_container.get_exposed_port(PORT_BOLT_NEO4J))
     db = InfrahubDatabaseProfiler(
-        mode=InfrahubDatabaseMode.DRIVER, driver=await get_db(), queries_names_to_config=queries_names_to_config
+        mode=InfrahubDatabaseMode.DRIVER,
+        query_analyzer=query_analyzer,
+        driver=await get_db(),
+        queries_names_to_config=queries_names_to_config,
     )
-    initialize_lock()
 
     # Create default branch
     await create_root_node(db=db)