Spaces:
Sleeping
Sleeping
Commit ·
99ac2ef
1
Parent(s): 71c1ad2
fix: initialize beanie with pymongo async client
Browse files- app/db/connection.py +24 -9
app/db/connection.py
CHANGED
|
@@ -1,36 +1,51 @@
|
|
| 1 |
# app/db/connection.py
|
| 2 |
-
# Beanie ODM initialization
|
| 3 |
|
| 4 |
from beanie import init_beanie
|
|
|
|
| 5 |
from app.observability.logging import get_logger
|
| 6 |
|
| 7 |
logger = get_logger(__name__)
|
| 8 |
|
|
|
|
|
|
|
| 9 |
|
| 10 |
async def connect_db() -> None:
|
| 11 |
-
"""Initialize Beanie ODM using
|
| 12 |
-
from
|
| 13 |
from app.db.models.user import UserDocument
|
| 14 |
from app.db.models.scan_result import ScanResultDocument
|
| 15 |
from app.db.models.alert import AlertDocument
|
|
|
|
|
|
|
|
|
|
| 16 |
|
| 17 |
-
if not
|
| 18 |
-
logger.error("beanie_init_skipped", reason="
|
| 19 |
return
|
| 20 |
|
| 21 |
try:
|
| 22 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 23 |
await init_beanie(
|
| 24 |
database=db,
|
| 25 |
document_models=[UserDocument, ScanResultDocument, AlertDocument],
|
| 26 |
allow_index_dropping=False,
|
| 27 |
)
|
| 28 |
-
logger.info("beanie_initialized", db=
|
| 29 |
except Exception as e:
|
| 30 |
logger.error("beanie_init_failed", error=str(e))
|
| 31 |
raise
|
| 32 |
|
| 33 |
|
| 34 |
async def close_db() -> None:
|
| 35 |
-
"""
|
| 36 |
-
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
# app/db/connection.py
|
| 2 |
+
# Beanie ODM initialization using PyMongo async client for Beanie compatibility.
|
| 3 |
|
| 4 |
from beanie import init_beanie
|
| 5 |
+
from app.config import get_settings
|
| 6 |
from app.observability.logging import get_logger
|
| 7 |
|
| 8 |
logger = get_logger(__name__)
|
| 9 |
|
| 10 |
+
_beanie_client = None
|
| 11 |
+
|
| 12 |
|
| 13 |
async def connect_db() -> None:
|
| 14 |
+
"""Initialize Beanie ODM using PyMongo's async client."""
|
| 15 |
+
from pymongo import AsyncMongoClient
|
| 16 |
from app.db.models.user import UserDocument
|
| 17 |
from app.db.models.scan_result import ScanResultDocument
|
| 18 |
from app.db.models.alert import AlertDocument
|
| 19 |
+
global _beanie_client
|
| 20 |
+
|
| 21 |
+
settings = get_settings()
|
| 22 |
|
| 23 |
+
if not settings.mongodb_uri:
|
| 24 |
+
logger.error("beanie_init_skipped", reason="mongodb_uri_missing")
|
| 25 |
return
|
| 26 |
|
| 27 |
try:
|
| 28 |
+
if _beanie_client is None:
|
| 29 |
+
_beanie_client = AsyncMongoClient(
|
| 30 |
+
settings.mongodb_uri,
|
| 31 |
+
serverSelectionTimeoutMS=5000,
|
| 32 |
+
)
|
| 33 |
+
|
| 34 |
+
db = _beanie_client[settings.mongodb_db_name]
|
| 35 |
await init_beanie(
|
| 36 |
database=db,
|
| 37 |
document_models=[UserDocument, ScanResultDocument, AlertDocument],
|
| 38 |
allow_index_dropping=False,
|
| 39 |
)
|
| 40 |
+
logger.info("beanie_initialized", db=settings.mongodb_db_name)
|
| 41 |
except Exception as e:
|
| 42 |
logger.error("beanie_init_failed", error=str(e))
|
| 43 |
raise
|
| 44 |
|
| 45 |
|
| 46 |
async def close_db() -> None:
|
| 47 |
+
"""Close the Beanie PyMongo async client."""
|
| 48 |
+
global _beanie_client
|
| 49 |
+
if _beanie_client is not None:
|
| 50 |
+
await _beanie_client.close()
|
| 51 |
+
_beanie_client = None
|