db/database.py

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import OperationalError
from config import MYSQL_CONFIG
from fastapi.responses import JSONResponse
from fastapi import HTTPException, status
from dotenv import load_dotenv
import io
import os
import base64

load_dotenv()

SQLALCHEMY_DATABASE_URL = MYSQL_CONFIG.DB_URI_SQL_ALCHEMY

# Retrieve the Base64-encoded CA certificate from the environment variable
ca_cert_base64 = os.getenv("CA_CERT_BASE64")

if ca_cert_base64:
    # Decode the base64 content
    ca_cert_content = base64.b64decode(ca_cert_base64).decode("utf-8")

    # Use in-memory buffer to handle the decoded CA certificate content
    ca_cert_in_memory = io.StringIO(ca_cert_content)

    # Create SQLAlchemy engine with SSL configuration, using in-memory certificate
    engine = create_engine(
        SQLALCHEMY_DATABASE_URL,
        connect_args={
            "ssl": {
                "sslmode": "REQUIRED",
                "sslrootcert": ca_cert_in_memory,  # In-memory CA certificate
                # Add other SSL options like client cert/key if required
            }
        },
    )

else:
    raise ValueError("CA_CERT_BASE64 environment variable is not set")


SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def get_db():
    db = SessionLocal()
    try:
        yield db
    except OperationalError as e:
        # Log the error and raise HTTPException for FastAPI
        print(f"An error occurred in get database sql alchemy.: {e}")
        raise HTTPException(status_code=400, detail="Database connection error")
    # Check if it's an authentication-related error
    except Exception as e:
        # Check if it's an authentication-related error
        if "401" in str(e):
            raise HTTPException(status_code=401, detail="Authentication failed")
        else:
            # For any other type of exception, raise a generic 400 error
            print(f"An error occurred: {e}")
            raise HTTPException(status_code=400, detail="An unexpected error occurred")