From fc730e6d10cf4f892a074dc084091c2ea600bdb1 Mon Sep 17 00:00:00 2001 From: GabrielBarberini Date: Wed, 29 May 2024 03:03:03 -0300 Subject: [PATCH] implements async context manager and threading lock --- lib/controllers/environment.py | 8 ++--- lib/controllers/flight.py | 12 +++---- lib/controllers/motor.py | 8 ++--- lib/controllers/rocket.py | 8 ++--- lib/repositories/repo.py | 57 +++++++++++++++++++--------------- 5 files changed, 50 insertions(+), 43 deletions(-) diff --git a/lib/controllers/environment.py b/lib/controllers/environment.py index 612e828..a087d1e 100644 --- a/lib/controllers/environment.py +++ b/lib/controllers/environment.py @@ -70,7 +70,7 @@ async def create_env(self) -> Union[EnvCreated, HTTPException]: views.EnvCreated """ try: - with EnvRepository() as env_repo: + async with EnvRepository() as env_repo: env_repo.fetch_env(self.env) await env_repo.create_env() except Exception as e: @@ -102,7 +102,7 @@ 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: - with EnvRepository() as env_repo: + async with EnvRepository() as env_repo: await env_repo.get_env_by_id(env_id) read_env = env_repo.env except Exception as e: @@ -180,7 +180,7 @@ async def update_env_by_id( HTTP 404 Not Found: If the env is not found in the database. """ try: - with EnvRepository() as env_repo: + async with EnvRepository() as env_repo: env_repo.fetch_env_by_id(env_id) await env_repo.create_env() await env_repo.delete_env_by_id(env_id) @@ -215,7 +215,7 @@ async def delete_env_by_id( HTTP 404 Not Found: If the env is not found in the database. """ try: - with EnvRepository() as env_repo: + async with EnvRepository() as env_repo: await env_repo.delete_env_by_id(env_id) except Exception as e: exc_str = parse_error(e) diff --git a/lib/controllers/flight.py b/lib/controllers/flight.py index 3f6641f..78c3ee1 100644 --- a/lib/controllers/flight.py +++ b/lib/controllers/flight.py @@ -107,7 +107,7 @@ async def create_flight(self) -> Union[FlightCreated, HTTPException]: views.FlightCreated """ try: - with FlightRepository() as flight_repo: + async with FlightRepository() as flight_repo: flight_repo.fetch_flight(self.flight) await flight_repo.create_flight( motor_kind=self.motor_kind, @@ -142,7 +142,7 @@ 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: - with FlightRepository() as flight_repo: + async with FlightRepository() as flight_repo: await flight_repo.get_flight_by_id(flight_id) read_flight = flight_repo.flight except Exception as e: @@ -220,7 +220,7 @@ async def update_flight_by_id( HTTP 404 Not Found: If the flight is not found in the database. """ try: - with FlightRepository() as flight_repo: + async with FlightRepository() as flight_repo: flight_repo.fetch_flight(self.flight) await flight_repo.create_flight( motor_kind=self.motor_kind, @@ -263,7 +263,7 @@ async def update_env_by_flight_id( new_flight = read_flight.dict() new_flight["environment"] = env new_flight = Flight(**new_flight) - with FlightRepository() as flight_repo: + async with FlightRepository() as flight_repo: flight_repo.fetch_flight(new_flight) await flight_repo.create_flight( motor_kind=read_flight.rocket.motor.motor_kind, @@ -311,7 +311,7 @@ async def update_rocket_by_flight_id( new_flight = read_flight.dict() new_flight["rocket"] = updated_rocket new_flight = Flight(**new_flight) - with FlightRepository() as flight_repo: + async with FlightRepository() as flight_repo: flight_repo.fetch_flight(new_flight) await flight_repo.create_flight( motor_kind=motor_kind, rocket_option=rocket_option @@ -350,7 +350,7 @@ async def delete_flight_by_id( HTTP 404 Not Found: If the flight is not found in the database. """ try: - with FlightRepository() as flight_repo: + async with FlightRepository() as flight_repo: await flight_repo.delete_flight_by_id(flight_id) except Exception as e: exc_str = parse_error(e) diff --git a/lib/controllers/motor.py b/lib/controllers/motor.py index b5b46e0..b702afe 100644 --- a/lib/controllers/motor.py +++ b/lib/controllers/motor.py @@ -118,7 +118,7 @@ async def create_motor(self) -> Union[MotorCreated, HTTPException]: views.MotorCreated """ try: - with MotorRepository() as motor_repo: + async with MotorRepository() as motor_repo: motor_repo.fetch_motor(self.motor) await motor_repo.create_motor(motor_kind=self.motor_kind) except Exception as e: @@ -150,7 +150,7 @@ 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: - with MotorRepository() as motor_repo: + async with MotorRepository() as motor_repo: await motor_repo.get_motor_by_id(motor_id) read_motor = motor_repo.motor except Exception as e: @@ -228,7 +228,7 @@ async def update_motor_by_id( HTTP 404 Not Found: If the motor is not found in the database. """ try: - with MotorRepository() as motor_repo: + async with MotorRepository() as motor_repo: motor_repo.fetch_motor(self.motor) await motor_repo.create_motor(motor_kind=self.motor_kind) await motor_repo.delete_motor_by_id(motor_id) @@ -263,7 +263,7 @@ async def delete_motor_by_id( HTTP 404 Not Found: If the motor is not found in the database. """ try: - with MotorRepository() as motor_repo: + async with MotorRepository() as motor_repo: await motor_repo.delete_motor_by_id(motor_id) except Exception as e: exc_str = parse_error(e) diff --git a/lib/controllers/rocket.py b/lib/controllers/rocket.py index b228b4f..8c0403f 100644 --- a/lib/controllers/rocket.py +++ b/lib/controllers/rocket.py @@ -155,7 +155,7 @@ async def create_rocket(self) -> Union[RocketCreated, HTTPException]: views.RocketCreated """ try: - with RocketRepository() as rocket_repo: + async with RocketRepository() as rocket_repo: rocket_repo.fetch_rocket(self.rocket) await rocket_repo.create_rocket( rocket_option=self.rocket_option, @@ -192,7 +192,7 @@ async def get_rocket_by_id( HTTP 404 Not Found: If the rocket is not found in the database. """ try: - with RocketRepository() as rocket_repo: + async with RocketRepository() as rocket_repo: await rocket_repo.get_rocket_by_id(rocket_id) read_rocket = rocket_repo.rocket except Exception as e: @@ -269,7 +269,7 @@ async def update_rocket_by_id( HTTP 404 Not Found: If the rocket is not found in the database. """ try: - with RocketRepository() as rocket_repo: + async with RocketRepository() as rocket_repo: rocket_repo.fetch_rocket(self.rocket) await rocket_repo.create_rocket( rocket_option=self.rocket_option, @@ -307,7 +307,7 @@ async def delete_rocket_by_id( HTTP 404 Not Found: If the rocket is not found in the database. """ try: - with RocketRepository() as rocket_repo: + async with RocketRepository() as rocket_repo: await rocket_repo.delete_rocket_by_id(rocket_id) except Exception as e: exc_str = parse_error(e) diff --git a/lib/repositories/repo.py b/lib/repositories/repo.py index 108c66d..6b9ed83 100644 --- a/lib/repositories/repo.py +++ b/lib/repositories/repo.py @@ -1,6 +1,6 @@ -import threading +import asyncio import logging -from lib.secrets import Secrets +from app.lib.secrets import Secrets from motor.motor_asyncio import AsyncIOMotorClient from pymongo.server_api import ServerApi @@ -11,10 +11,10 @@ class Repository: """ _instances = {} - _lock = threading.Lock() + _lock = asyncio.Lock() def __new__(cls): - with cls._lock: + async with cls._lock: if cls not in cls._instances: cls._instances[cls] = super(Repository, cls).__new__(cls) cls._instances[cls]._initialized = False # Initialize here @@ -22,31 +22,38 @@ def __new__(cls): def __init__(self, collection: str): if not self._initialized: - try: - self._connection_string = Secrets.get_secret( - "MONGODB_CONNECTION_STRING" - ) - self._client = AsyncIOMotorClient( - self.connection_string, - server_api=ServerApi("1"), - maxIdleTimeMS=5000, - connectTimeoutMS=5000, - serverSelectionTimeoutMS=30000, - ) - self._collection = self.client.rocketpy[collection] - self._initialized = True # Mark as initialized - except Exception as e: - logging.error(f"Failed to initialize MongoDB client: {e}") - raise ConnectionError( - "Could not establish a connection with MongoDB." - ) from e + self._collection_name = collection + self._initialize_connection() + self._initialized = True - def __enter__(self): + def _initialize_connection(self): + try: + self._connection_string = Secrets.get_secret( + "MONGODB_CONNECTION_STRING" + ) + self._client = AsyncIOMotorClient( + self._connection_string, + server_api=ServerApi("1"), + maxIdleTimeMS=5000, + connectTimeoutMS=5000, + serverSelectionTimeoutMS=30000, + ) + self._collection = self._client.rocketpy[self._collection_name] + except Exception as e: + logging.error( + f"Failed to initialize MongoDB client: {e}", exc_info=True + ) + raise ConnectionError( + "Could not establish a connection with MongoDB." + ) from e + + def __aenter__(self): return self - def __exit__(self, exc_type, exc_value, traceback): + def __aexit__(self, exc_type, exc_value, traceback): self.close_connection() - self._instances.pop(self.__class__) + async with self._lock: + self._instances.pop(self.__class__, None) @property def connection_string(self):