Minor changes

Author: LucasG0

backend/infrahub/database/__init__.py

@@ -78,7 +78,7 @@ def get_branch_name(branch: Optional[Union[Branch, str]] = None) -> str:
 
 
 class DatabaseSchemaManager:
-    def __init__(self, db: InfrahubDatabase):
+    def __init__(self, db: InfrahubDatabase) -> None:
         self._db = db
 
     def get(self, name: str, branch: Optional[Union[Branch, str]] = None, duplicate: bool = True) -> MainSchemaTypes:

backend/tests/helpers/query_benchmark/car_person_generators.py

@@ -31,7 +31,10 @@ async def load_cars(self, nb_cars: int, persons: Optional[dict[str, Node]] = Non
                 await car_node.new(db=self.db, name=car_name, nbr_seats=4, owner=random_person)
             else:
                 await car_node.new(db=self.db, name=car_name, nbr_seats=4)
-            await self._save_obj(car_node)
+
+            async with self.db.start_session():
+                await car_node.save(db=self.db)
+
             cars[car_name] = car_node
 
         return cars
@@ -60,7 +63,10 @@ async def load_persons(self, nb_persons: int, cars: Optional[dict[str, Node]] =
                 await person_node.new(db=self.db, name=person_name, cars=random_cars)
             else:
                 await person_node.new(db=self.db, name=person_name)
-            await self._save_obj(person_node)
+
+            async with self.db.start_session():
+                await person_node.save(db=self.db)
+
             persons_names_to_nodes[person_name] = person_node
 
         return persons_names_to_nodes

backend/tests/helpers/query_benchmark/data_generator.py

@@ -5,7 +5,6 @@
 from rich.console import Console
 from rich.progress import Progress
 
-from infrahub.core.node import Node
 from tests.helpers.query_benchmark.db_query_profiler import InfrahubDatabaseProfiler, ProfilerEnabler, query_analyzer
 
 
@@ -26,11 +25,6 @@ async def init(self) -> None:
     async def load_data(self, nb_elements: int) -> None:
         raise NotImplementedError("Abstract method")
 
-    async def _save_obj(self, obj: Node) -> None:
-        async with self.db.start_session() as dbs:
-            async with dbs.start_transaction() as dbt:
-                await obj.save(db=dbt)
-
 
 async def load_data_and_profile(
     data_generator: DataGenerator,