Fix DB engine leak by lazy singleton and Add CORS

Author: BhaveshKukreja29

mxgo/api.py

@@ -11,6 +11,7 @@
 import redis.asyncio as aioredis
 from dotenv import load_dotenv
 from fastapi import Depends, FastAPI, File, Form, HTTPException, Response, UploadFile, status
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.security import APIKeyHeader, HTTPBearer
 from sqlalchemy import text
 
@@ -127,6 +128,30 @@ async def lifespan(_app: FastAPI):
 
 
 app = FastAPI(lifespan=lifespan)
+
+IS_PROD = os.getenv("IS_PROD", "false").lower() == "true"
+
+ALLOWED_ORIGINS_PROD = [
+    "https://mxgo.ai",
+]
+
+
+ALLOWED_ORIGINS_DEV = [
+    "http://localhost",
+    "http://localhost:8080",
+    "http://127.0.0.1",
+    "http://127.0.0.1:8080",
+]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=ALLOWED_ORIGINS_PROD if IS_PROD else ALLOWED_ORIGINS_DEV,
+    allow_credentials=True,
+    allow_methods=["GET", "POST"],
+    allow_headers=["*"],
+)
+
+
 if os.getenv("IS_PROD", "false").lower() == "true":
     app.openapi_url = None
 

mxgo/db/__init__.py

@@ -14,36 +14,35 @@ class DbConnection:
     _engine = None
 
     def __init__(self) -> None:
-        self.db_uri = self.get_db_uri_from_env()
+        pass
 
     @classmethod
     def get_db_uri_from_env(cls) -> str:
         return f"postgresql://{os.environ['DB_USER']}:{os.environ['DB_PASSWORD']}@{os.environ['DB_HOST']}:{os.environ['DB_PORT']}/{os.environ['DB_NAME']}"
 
     def init_connection(self) -> None:
         """Initialize the database connection."""
-        if self._engine is None:
+        if DbConnection._engine is None:
             self._create_engine()
 
     def close_connection(self) -> None:
         """Close the database connection."""
-        if self._engine is not None:
-            self._engine.dispose()
+        if DbConnection._engine is not None:
+            DbConnection._engine.dispose()
+            DbConnection._engine = None
 
     def get_connection(self):
         """Get the database connection engine."""
-        if self._engine is None:
-            msg = "DB session isn't initialized"
-            raise ConnectionError(msg)
-        return self._engine
+        if DbConnection._engine is None:
+            self.init_connection()
+        return DbConnection._engine
 
     @contextmanager
     def get_session(self) -> Generator[Session]:
         """Get a synchronous database session."""
-        if self._engine is None:
-            self.init_connection()
+        engine = self.get_connection()
 
-        session = Session(self._engine)
+        session = Session(engine)
         try:
             yield session
             session.commit()
@@ -55,9 +54,9 @@ def get_session(self) -> Generator[Session]:
 
     def _create_engine(self):
         """Create a synchronous SQLAlchemy engine."""
-        db_url = self.db_uri
-        self._engine = create_engine(db_url, pool_pre_ping=True, echo=False)
-        return self._engine
+        db_url = self.get_db_uri_from_env()
+        DbConnection._engine = create_engine(db_url, pool_pre_ping=True, echo=False)
+        return DbConnection._engine
 
 
 class AsyncDbConnection:
@@ -73,18 +72,19 @@ def get_db_uri_from_env(cls) -> str:
         return f"postgresql+asyncpg://{os.environ['DB_USER']}:{os.environ['DB_PASSWORD']}@{os.environ['DB_HOST']}:{os.environ['DB_PORT']}/{os.environ['DB_NAME']}"
 
     async def init_connection(self) -> None:
-        if self._engine is None:
+        if AsyncDbConnection._engine is None:
             await self._create_engine()
 
     async def close_connection(self) -> None:
-        if self._engine is not None:
-            await self._engine.dispose()
+        if AsyncDbConnection._engine is not None:
+            await AsyncDbConnection._engine.dispose()
+            AsyncDbConnection._engine = None
 
     def get_connection(self) -> AsyncEngine:
-        if self._engine is None:
+        if AsyncDbConnection._engine is None:
             msg = "Async DB session isn't initialized"
             raise ConnectionError(msg)
-        return self._engine
+        return AsyncDbConnection._engine
 
     @asynccontextmanager
     async def get_session(self) -> AsyncGenerator[AsyncSession]:
@@ -106,19 +106,21 @@ async def get_session(self) -> AsyncGenerator[AsyncSession]:
 
     async def _create_engine(self) -> AsyncEngine:
         db_url = self.db_uri
-        self._engine = create_async_engine(db_url, pool_pre_ping=True, echo=False)
-        return self._engine
+        AsyncDbConnection._engine = create_async_engine(db_url, pool_pre_ping=True, echo=False)
+        return AsyncDbConnection._engine
+
+
+db_connection = DbConnection()
+async_db_connection = AsyncDbConnection()
 
 
 def init_db_connection() -> DbConnection:
     """Initialize a synchronous database connection."""
-    db_connection = DbConnection()
     db_connection.init_connection()
     return db_connection
 
 
 async def init_async_db_connection() -> AsyncDbConnection:
     """Initialize an asynchronous database connection for future use."""
-    db_connection = AsyncDbConnection()
-    await db_connection.init_connection()
-    return db_connection
+    await async_db_connection.init_connection()
+    return async_db_connection