ryan
0f903f57bb
Chore App Build, Test, and Push Docker Images / build-and-push (push) Failing after 2m22s
116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
# python
|
|
"""Lazy MongoDB client factory and database-name helpers.
|
|
|
|
The client is intentionally **not** created at module import. Use
|
|
``get_mongo_client()`` to obtain a cached singleton. For Gunicorn multi-worker
|
|
deployments call ``init_mongo_client()`` from a ``post_fork`` hook so each
|
|
worker process owns its own connection pool rather than inheriting the parent
|
|
process's client.
|
|
"""
|
|
import os
|
|
import threading
|
|
from urllib.parse import urlparse
|
|
|
|
from pymongo import MongoClient
|
|
from pymongo.uri_parser import parse_uri
|
|
|
|
|
|
_mongo_client_lock = threading.Lock()
|
|
_mongo_client = None
|
|
|
|
|
|
def _create_mongo_client():
|
|
"""Build a fail-fast MongoClient from environment variables."""
|
|
uri = os.environ.get('MONGO_URI')
|
|
|
|
|
|
if not uri:
|
|
raise RuntimeError(
|
|
'MONGO_URI environment variable is required when USE_MONGODB=true.'
|
|
)
|
|
|
|
# mongomock is used for unit/integration tests without a real server.
|
|
if uri.lower().startswith('mongomock') or uri.lower() == 'mongomock':
|
|
try:
|
|
import mongomock
|
|
except ImportError as exc: # pragma: no cover - test dependency
|
|
raise RuntimeError(
|
|
'mongomock is required for test MongoDB mode. '
|
|
'Install it with: pip install mongomock'
|
|
) from exc
|
|
return mongomock.MongoClient()
|
|
|
|
return MongoClient(
|
|
uri,
|
|
serverSelectionTimeoutMS=5000,
|
|
connectTimeoutMS=5000,
|
|
maxPoolSize=20,
|
|
)
|
|
|
|
|
|
def init_mongo_client():
|
|
"""Create a fresh MongoClient and store it as the process singleton.
|
|
|
|
Call this from a Gunicorn ``post_fork`` hook so each worker process gets
|
|
its own client after forking. It can also be called in tests to reset the
|
|
shared client to a known state.
|
|
"""
|
|
global _mongo_client
|
|
with _mongo_client_lock:
|
|
_mongo_client = _create_mongo_client()
|
|
return _mongo_client
|
|
|
|
|
|
def get_mongo_client():
|
|
"""Return the cached process-level MongoClient, creating it lazily once."""
|
|
global _mongo_client
|
|
if _mongo_client is None:
|
|
with _mongo_client_lock:
|
|
if _mongo_client is None:
|
|
_mongo_client = _create_mongo_client()
|
|
return _mongo_client
|
|
|
|
|
|
def _db_name_from_uri(uri: str) -> str | None:
|
|
"""Extract the database name from a MongoDB connection URI, if present."""
|
|
if not uri or uri.lower().startswith('mongomock'):
|
|
return None
|
|
try:
|
|
parsed = parse_uri(uri)
|
|
return parsed.get('database') or None
|
|
except Exception:
|
|
# Fallback to a simple path-based parse for non-standard URIs.
|
|
try:
|
|
path = urlparse(uri).path
|
|
return path.lstrip('/') or None
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def get_mongo_db_name() -> str:
|
|
"""Resolve the MongoDB database name from env vars or the connection URI.
|
|
|
|
Precedence:
|
|
1. ``MONGO_DB_NAME`` environment variable.
|
|
2. Database name parsed from ``MONGO_URI``.
|
|
3. Default based on ``DATA_ENV`` / ``DB_ENV``:
|
|
* prod -> ``chore_db``
|
|
* e2e -> ``chore_db_e2e``
|
|
* test -> ``chore_db_test``
|
|
"""
|
|
env_name = os.environ.get('MONGO_DB_NAME')
|
|
if env_name:
|
|
return env_name
|
|
|
|
uri = os.environ.get('MONGO_URI', '')
|
|
db_name = _db_name_from_uri(uri)
|
|
if db_name:
|
|
return db_name
|
|
|
|
env = (os.environ.get('DATA_ENV') or os.environ.get('DB_ENV', 'prod')).lower()
|
|
if env == 'prod':
|
|
return 'chore_db'
|
|
if env == 'e2e':
|
|
return 'chore_db_e2e'
|
|
return 'chore_db_test'
|