Coverage for app / db / database.py: 73%

1import os 

2from pathlib import Path 

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

4from sqlalchemy.orm import DeclarativeBase 

5 

6from app.core.config import settings 

7 

8 

9class Base(DeclarativeBase): 

10 pass 

11 

12 

13# Ensure data directory exists for SQLite 

14if "sqlite" in settings.database_url: 

15 db_path = settings.database_url.split("///")[-1] if "///" in settings.database_url else settings.database_url.split(":///")[-1] 

16 db_dir = os.path.dirname(db_path) 

17 if db_dir: 

18 Path(db_dir).mkdir(parents=True, exist_ok=True) 

19 

20 

21engine = create_async_engine( 

22 settings.database_url, 

23 echo=False, 

24 future=True, 

25) 

26 

27AsyncSessionLocal = async_sessionmaker( 

28 engine, 

29 class_=AsyncSession, 

30 expire_on_commit=False, 

31 autocommit=False, 

32 autoflush=False, 

33) 

34 

35 

36async def get_db() -> AsyncSession: 

37 async with AsyncSessionLocal() as session: 

38 try: 

39 yield session 

40 finally: 

41 await session.close() 

42 

43 

44async def init_db(): 

45 """Initialize database tables.""" 

46 async with engine.begin() as conn: 

47 await conn.run_sync(Base.metadata.create_all)