Coverage for core/database.py: 72.00%

1import logging 

2 

3from sqlalchemy import create_engine 

4from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine 

5from sqlalchemy.orm import declarative_base, sessionmaker 

6 

7from core.config import settings 

8 

9logger = logging.getLogger("database") 

10 

11 

12def make_async_url(url: str) -> str: 

13 if url.startswith("mysql://"): 

14 return url.replace("mysql://", "mysql+aiomysql://", 1) 

15 if url.startswith("mysql+pymysql://"): 

16 return url.replace("mysql+pymysql://", "mysql+aiomysql://", 1) 

17 return url 

18 

19 

20# Async engine/session for API 

21async_engine = create_async_engine( 

22 make_async_url(settings.DATABASE_URL), 

23 echo=settings.DEBUG_MODE, 

24 future=True, 

25 pool_pre_ping=True, 

26 pool_recycle=settings.DB_POOL_RECYCLE, 

27 pool_size=settings.DB_POOL_SIZE, 

28 pool_timeout=settings.DB_POOL_TIMEOUT, 

29 max_overflow=settings.DB_MAX_OVERFLOW, 

30 connect_args={ 

31 "charset": "utf8mb4", 

32 }, 

33) 

34AsyncSessionLocal = async_sessionmaker( 

35 bind=async_engine, 

36 class_=AsyncSession, 

37 expire_on_commit=False, 

38 autoflush=False, 

39 autocommit=False, 

40) 

41 

42# Sync engine/session for migration and schedule 

43engine = create_engine( 

44 settings.DATABASE_URL, 

45 echo=settings.DEBUG_MODE, 

46 future=True, 

47 pool_pre_ping=True, 

48 pool_recycle=settings.DB_POOL_RECYCLE, 

49 pool_size=settings.DB_POOL_SIZE, 

50 pool_timeout=settings.DB_POOL_TIMEOUT, 

51 max_overflow=settings.DB_MAX_OVERFLOW, 

52 connect_args={ 

53 "charset": "utf8mb4", 

54 "connect_timeout": settings.DB_CONNECT_TIMEOUT, 

55 "read_timeout": settings.DB_READ_TIMEOUT, 

56 "write_timeout": settings.DB_WRITE_TIMEOUT, 

57 }, 

58) 

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

60 

61# Base class for declarative models 

62Base = declarative_base() 

63 

64 

65async def init_db(): 

66 """Initialize the database with default data""" 

67 try: 

68 # Import here to avoid circular import 

69 from core.init_db import init_database 

70 

71 await init_database() 

72 logger.info("Database initialization completed") 

73 except Exception as e: 

74 logger.error(f"Database initialization failed: {str(e)}") 

75 raise