implements db context manager

RocketPy-Team · May 24, 2024 · 77efcc2 · 77efcc2
1 parent 7879742
commit 77efcc2
Show file tree

Hide file tree

Showing 9 changed files with 143 additions and 105 deletions.
diff --git a/lib/controllers/environment.py b/lib/controllers/environment.py
@@ -70,7 +70,8 @@ async def create_env(self) -> Union[EnvCreated, HTTPException]:
             views.EnvCreated
         """
         try:
-            await EnvRepository.fetch_env(self.env).create_env()
+            with EnvRepository.fetch_env(self.env) as env_repo:
+                await env_repo.create_env()
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.environment.create_env: {exc_str}")
@@ -100,7 +101,9 @@ async def get_env_by_id(env_id: str) -> Union[Env, HTTPException]:
             HTTP 404 Not Found: If the env is not found in the database.
         """
         try:
-            read_env = await EnvRepository().get_env_by_id(env_id)
+            with EnvRepository() as env_repo:
+                await env_repo.get_env_by_id(env_id)
+                read_env = env_repo.env
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.environment.get_env_by_id: {exc_str}")
@@ -176,8 +179,9 @@ async def update_env_by_id(
             HTTP 404 Not Found: If the env is not found in the database.
         """
         try:
-            env_repo = await EnvRepository.fetch_env(self.env).create_env()
-            await env_repo.delete_env_by_id(env_id)
+            with EnvRepository.fetch_env(self.env) as env_repo:
+                await env_repo.create_env()
+                await env_repo.delete_env_by_id(env_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.environment.update_env: {exc_str}")
@@ -209,7 +213,8 @@ async def delete_env_by_id(
             HTTP 404 Not Found: If the env is not found in the database.
         """
         try:
-            await EnvRepository().delete_env_by_id(env_id)
+            with EnvRepository() as env_repo:
+                await env_repo.delete_env_by_id(env_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.environment.delete_env: {exc_str}")
@@ -243,6 +248,7 @@ async def simulate_env(
         try:
             read_env = await cls.get_env_by_id(env_id)
             rocketpy_env = cls.get_rocketpy_env(read_env)
+
             env_simulation_numbers = EnvData.parse_obj(
                 rocketpy_env.all_info_returned()
             )

diff --git a/lib/controllers/flight.py b/lib/controllers/flight.py
@@ -107,9 +107,11 @@ async def create_flight(self) -> Union[FlightCreated, HTTPException]:
             views.FlightCreated
         """
         try:
-            await FlightRepository.fetch_flight(self.flight).create_flight(
-                motor_kind=self.motor_kind, rocket_option=self.rocket_option
-            )
+            with FlightRepository.fetch_flight(self.flight) as flight_repo:
+                await flight_repo.create_flight(
+                    motor_kind=self.motor_kind,
+                    rocket_option=self.rocket_option,
+                )
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.flight.create_flight: {exc_str}")
@@ -139,7 +141,9 @@ async def get_flight_by_id(flight_id: str) -> Union[Flight, HTTPException]:
             HTTP 404 Not Found: If the flight is not found in the database.
         """
         try:
-            read_flight = await FlightRepository().get_flight_by_id(flight_id)
+            with FlightRepository() as flight_repo:
+                await flight_repo.get_flight_by_id(flight_id)
+                read_flight = flight_repo.flight
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.flight.get_flight_by_id: {exc_str}")
@@ -215,12 +219,12 @@ async def update_flight_by_id(
             HTTP 404 Not Found: If the flight is not found in the database.
         """
         try:
-            flight_repo = await FlightRepository.fetch_flight(
-                self.flight
-            ).create_flight(
-                motor_kind=self.motor_kind, rocket_option=self.rocket_option
-            )
-            await flight_repo.delete_flight_by_id(flight_id)
+            with FlightRepository.fetch_flight(self.flight) as flight_repo:
+                await flight_repo.create_flight(
+                    motor_kind=self.motor_kind,
+                    rocket_option=self.rocket_option,
+                )
+                await flight_repo.delete_flight_by_id(flight_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.flight.update_flight: {exc_str}")
@@ -257,13 +261,12 @@ async def update_env_by_flight_id(
             new_flight = read_flight.dict()
             new_flight["environment"] = env
             new_flight = Flight(**new_flight)
-            flight_repo = await FlightRepository.fetch_flight(
-                new_flight
-            ).create_flight(
-                motor_kind=read_flight.rocket.motor.motor_kind,
-                rocket_option=read_flight.rocket.rocket_option,
-            )
-            await flight_repo.delete_flight_by_id(flight_id)
+            with FlightRepository.fetch_flight(new_flight) as flight_repo:
+                await flight_repo.create_flight(
+                    motor_kind=read_flight.rocket.motor.motor_kind,
+                    rocket_option=read_flight.rocket.rocket_option,
+                )
+                await flight_repo.delete_flight_by_id(flight_id)
         except HTTPException as e:
             raise e from e
         except Exception as e:
@@ -305,10 +308,11 @@ async def update_rocket_by_flight_id(
             new_flight = read_flight.dict()
             new_flight["rocket"] = updated_rocket
             new_flight = Flight(**new_flight)
-            flight_repo = await FlightRepository.fetch_flight(
-                new_flight
-            ).create_flight(motor_kind=motor_kind, rocket_option=rocket_option)
-            await flight_repo.delete_flight_by_id(flight_id)
+            with FlightRepository.fetch_flight(new_flight) as flight_repo:
+                await flight_repo.create_flight(
+                    motor_kind=motor_kind, rocket_option=rocket_option
+                )
+                await flight_repo.delete_flight_by_id(flight_id)
         except HTTPException as e:
             raise e from e
         except Exception as e:
@@ -342,7 +346,8 @@ async def delete_flight_by_id(
             HTTP 404 Not Found: If the flight is not found in the database.
         """
         try:
-            await FlightRepository().delete_flight_by_id(flight_id)
+            with FlightRepository() as flight_repo:
+                await flight_repo.delete_flight_by_id(flight_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.flight.delete_flight: {exc_str}")
@@ -376,8 +381,7 @@ async def simulate_flight(
         """
         try:
             read_flight = await cls.get_flight_by_id(flight_id=flight_id)
-            rocketpy_flight = cls.get_rocketpy_flight(flight=read_flight)
-            flight = rocketpy_flight
+            flight = cls.get_rocketpy_flight(read_flight)
 
             _initial_conditions = InitialConditions(
                 initial_altitude="Attitude - e0: {:.3f} | e1: {:.3f} | e2: {:.3f} | e3: {:.3f}".format(

diff --git a/lib/controllers/motor.py b/lib/controllers/motor.py
@@ -118,9 +118,8 @@ async def create_motor(self) -> Union[MotorCreated, HTTPException]:
             views.MotorCreated
         """
         try:
-            await MotorRepository.fetch_motor(self.motor).create_motor(
-                motor_kind=self.motor_kind
-            )
+            with MotorRepository.fetch_motor(self.motor) as motor_repo:
+                await motor_repo.create_motor(motor_kind=self.motor_kind)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.motor.create_motor: {exc_str}")
@@ -150,7 +149,9 @@ async def get_motor_by_id(motor_id: str) -> Union[Motor, HTTPException]:
             HTTP 404 Not Found: If the motor is not found in the database.
         """
         try:
-            read_motor = await MotorRepository().get_motor_by_id(motor_id)
+            with MotorRepository() as motor_repo:
+                await motor_repo.get_motor_by_id(motor_id)
+                read_motor = motor_repo.motor
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.motor.get_motor_by_id: {exc_str}")
@@ -226,10 +227,9 @@ async def update_motor_by_id(
             HTTP 404 Not Found: If the motor is not found in the database.
         """
         try:
-            motor_repo = await MotorRepository.fetch_motor(
-                self.motor
-            ).create_motor(motor_kind=self.motor_kind)
-            await motor_repo.delete_motor_by_id(motor_id)
+            with MotorRepository.fetch_motor(self.motor) as motor_repo:
+                await motor_repo.create_motor(motor_kind=self.motor_kind)
+                await motor_repo.delete_motor_by_id(motor_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.motor.update_motor: {exc_str}")
@@ -261,7 +261,8 @@ async def delete_motor_by_id(
             HTTP 404 Not Found: If the motor is not found in the database.
         """
         try:
-            await MotorRepository().delete_motor_by_id(motor_id)
+            with MotorRepository() as motor_repo:
+                await motor_repo.delete_motor_by_id(motor_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.motor.delete_motor: {exc_str}")
@@ -293,7 +294,7 @@ async def simulate_motor(
             HTTP 404 Not Found: If the motor does not exist in the database.
         """
         try:
-            read_motor = await MotorRepository().get_motor_by_id(motor_id)
+            read_motor = await cls.get_motor_by_id(motor_id)
             motor = cls.get_rocketpy_motor(read_motor)
 
             motor_simulation_numbers = MotorData(

diff --git a/lib/controllers/rocket.py b/lib/controllers/rocket.py
@@ -155,9 +155,11 @@ async def create_rocket(self) -> Union[RocketCreated, HTTPException]:
             views.RocketCreated
         """
         try:
-            await RocketRepository.fetch_rocket(self.rocket).create_rocket(
-                rocket_option=self.rocket_option, motor_kind=self.motor_kind
-            )
+            with RocketRepository.fetch_rocket(self.rocket) as rocket_repo:
+                await rocket_repo.create_rocket(
+                    rocket_option=self.rocket_option,
+                    motor_kind=self.motor_kind,
+                )
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.rocket.create_rocket: {exc_str}")
@@ -189,7 +191,9 @@ async def get_rocket_by_id(
             HTTP 404 Not Found: If the rocket is not found in the database.
         """
         try:
-            read_rocket = await RocketRepository().get_rocket_by_id(rocket_id)
+            with RocketRepository() as rocket_repo:
+                await rocket_repo.get_rocket_by_id(rocket_id)
+                read_rocket = rocket_repo.rocket
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.rocket.get_rocket_by_id: {exc_str}")
@@ -264,12 +268,12 @@ async def update_rocket_by_id(
             HTTP 404 Not Found: If the rocket is not found in the database.
         """
         try:
-            rocket_repo = await RocketRepository.fetch_rocket(
-                self.rocket
-            ).create_rocket(
-                rocket_option=self.rocket_option, motor_kind=self.motor_kind
-            )
-            await rocket_repo.delete_rocket_by_id(rocket_id)
+            with RocketRepository.fetch_rocket(self.rocket) as rocket_repo:
+                await rocket_repo.create_rocket(
+                    rocket_option=self.rocket_option,
+                    motor_kind=self.motor_kind,
+                )
+                await rocket_repo.delete_rocket_by_id(rocket_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.rocket.update_rocket: {exc_str}")
@@ -301,7 +305,8 @@ async def delete_rocket_by_id(
             HTTP 404 Not Found: If the rocket is not found in the database.
         """
         try:
-            await RocketRepository().delete_rocket_by_id(rocket_id)
+            with RocketRepository() as rocket_repo:
+                await rocket_repo.delete_rocket_by_id(rocket_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"controllers.rocket.delete_rocket: {exc_str}")

diff --git a/lib/repositories/environment.py b/lib/repositories/environment.py
@@ -18,6 +18,7 @@ class EnvRepository(Repository):
 
     def __init__(self):
         super().__init__("environments")
+        self._env = None
 
     @classmethod
     def fetch_env(cls, environment: Env):
@@ -42,6 +43,17 @@ def env_id(self) -> str:
     def env_id(self, env_id: "str"):
         self._env_id = env_id
 
+    async def insert_env(self, env_data: dict):
+        await self.collection.insert_one(env_data)
+        return self
+
+    async def find_env(self, env_id: str):
+        return await self.collection.find_one({"env_id": env_id})
+
+    async def delete_env(self, env_id: str):
+        await self.collection.delete_one({"env_id": env_id})
+        return self
+
     async def create_env(self):
         """
         Creates a non-existing models.Env in the database
@@ -69,17 +81,18 @@ async def get_env_by_id(self, env_id: str) -> Union[Env, None]:
         Gets a models.Env from the database
 
         Returns:
-            models.Env
+            self
         """
         try:
             read_env = await self.find_env(env_id)
             parsed_env = Env.parse_obj(read_env) if read_env else None
+            self.env = parsed_env
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"repositories.environment.get_env: {exc_str}")
             raise Exception(f"Error getting environment: {exc_str}") from e
         else:
-            return parsed_env
+            return self
         finally:
             logger.info(
                 f"Call to repositories.environment.get_env completed for Env {env_id}"
@@ -90,26 +103,17 @@ async def delete_env_by_id(self, env_id: str):
         Deletes a models.Env from the database
 
         Returns:
-            None
+            self
         """
         try:
             await self.delete_env(env_id)
         except Exception as e:
             exc_str = parse_error(e)
             logger.error(f"repositories.environment.delete_env: {exc_str}")
             raise Exception(f"Error deleting environment: {exc_str}") from e
+        else:
+            return self
         finally:
             logger.info(
                 f"Call to repositories.environment.delete_env completed for Env {env_id}"
             )
-
-    async def insert_env(self, env_data: dict):
-        await self.collection.insert_one(env_data)
-        return self
-
-    async def find_env(self, env_id: str):
-        return await self.collection.find_one({"env_id": env_id})
-
-    async def delete_env(self, env_id: str):
-        await self.collection.delete_one({"env_id": env_id})
-        return self