mongodb_keep_alive_service.py

import os
import logging
from urllib.parse import parse_qs, urlparse
from dotenv import load_dotenv
from pymongo import MongoClient
from datetime import datetime
from pydantic import BaseModel

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

load_dotenv()

mongodb_uris = os.getenv('MONGODB_URIS')


# ---------- Pydantic Response Model ----------
class MongoPingResponse(BaseModel):
    service_name: str
    success: bool
    error: str | None = None
    time: str

class MongoPingAllResponse(BaseModel):
    mongodb_services: list[MongoPingResponse]

# ---------- Helpers ----------
def extract_app_name_from_uri(uri: str) -> str:
    """Extracts the appName from MongoDB URI if present."""
    try:
        parsed_uri = urlparse(uri)
        query_params = parse_qs(parsed_uri.query)
        app_name = query_params.get('appName', ['MongoDB Ping Service'])[0]
        return app_name
    except Exception as e:
        logger.debug(f"Failed to extract appName from URI: {e}")
        return "MongoDB Ping Service"  # Fallback if parsing fails


# ---------- MongoDB Ping ----------
def ping_mongodb() -> list[MongoPingResponse]:
    """Ping MongoDB and return structured responses."""
    responses: list[MongoPingResponse] = []

    if not mongodb_uris:
        logger.warning("No MongoDB URIs found in environment variables")
        return responses

    for uri in mongodb_uris.split(','):
        uri = uri.strip()
        if not uri:
            continue

        service_name = extract_app_name_from_uri(uri)
        now = datetime.utcnow().isoformat()

        try:
            client = MongoClient(uri, serverSelectionTimeoutMS=3000)  # 3s timeout
            client.admin.command('ping')
            logger.info(f"{service_name}: Successfully pinged MongoDB at {uri[:50]}...")
            responses.append(
                MongoPingResponse(
                    service_name=service_name,
                    success=True,
                    error=None,
                    time=now
                )
            )
            client.close()
        except Exception as e:
            logger.error(f"{service_name}: Error pinging MongoDB at {uri[:50]}... - {str(e)}")
            responses.append(
                MongoPingResponse(
                    service_name=service_name,
                    success=False,
                    error=str(e),
                    time=now
                )
            )

    return responses


def ping_mongodb_projects():
    logger.info("Starting MongoDB ping checks...")
    results = ping_mongodb()
    logger.info("Completed MongoDB ping checks.")
    final_response = MongoPingAllResponse(mongodb_services=results)
    return final_response