57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
|
import re
|
||
|
|
from urllib.parse import urlparse
|
||
|
|
from sqlalchemy import create_engine
|
||
|
|
from sqlalchemy.ext.declarative import declarative_base
|
||
|
|
from sqlalchemy.orm import sessionmaker
|
||
|
|
from app.core.config import settings
|
||
|
|
|
||
|
|
def ensure_database_exists(db_url: str):
|
||
|
|
"""
|
||
|
|
If using MySQL, automatically creates the target database if it does not exist
|
||
|
|
(mirrors spring.datasource.url createDatabaseIfNotExist=true).
|
||
|
|
"""
|
||
|
|
if db_url.startswith("mysql"):
|
||
|
|
try:
|
||
|
|
import pymysql
|
||
|
|
# Parse connection URL safely
|
||
|
|
clean_url = db_url.replace("mysql+pymysql://", "http://").replace("mysql://", "http://")
|
||
|
|
parsed = urlparse(clean_url)
|
||
|
|
db_name = parsed.path.lstrip("/").split("?")[0]
|
||
|
|
if db_name:
|
||
|
|
conn = pymysql.connect(
|
||
|
|
host=parsed.hostname,
|
||
|
|
port=parsed.port or 3306,
|
||
|
|
user=parsed.username,
|
||
|
|
password=parsed.password,
|
||
|
|
connect_timeout=5
|
||
|
|
)
|
||
|
|
with conn.cursor() as cur:
|
||
|
|
cur.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}`;")
|
||
|
|
conn.close()
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Note on DB creation: {e}")
|
||
|
|
|
||
|
|
ensure_database_exists(settings.DATABASE_URL)
|
||
|
|
|
||
|
|
connect_args = {}
|
||
|
|
if settings.DATABASE_URL.startswith("sqlite"):
|
||
|
|
connect_args = {"check_same_thread": False}
|
||
|
|
|
||
|
|
engine = create_engine(
|
||
|
|
settings.DATABASE_URL,
|
||
|
|
connect_args=connect_args,
|
||
|
|
pool_pre_ping=True
|
||
|
|
)
|
||
|
|
|
||
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||
|
|
|
||
|
|
Base = declarative_base()
|
||
|
|
|
||
|
|
def get_db():
|
||
|
|
"""Dependency for obtaining database session per request."""
|
||
|
|
db = SessionLocal()
|
||
|
|
try:
|
||
|
|
yield db
|
||
|
|
finally:
|
||
|
|
db.close()
|