Coverage for app/infrastructure/database/db.py: 88%

1from collections.abc import Generator 

2from contextlib import contextmanager 

3 

4# 設定モデルを定義 

5from pydantic_settings import BaseSettings 

6from sqlmodel import Session, create_engine 

7 

8 

9class Settings(BaseSettings): 

10 SQLALCHEMY_DATABASE: str 

11 SQLALCHEMY_DATABASE_SQLITE_URI: str 

12 

13 @property 

14 def database_url(self) -> str: 

15 """SQLALCHEMY_DATABASE に基づいて適切な URI を返す.""" 

16 db_urls = { 

17 "sqlite": self.SQLALCHEMY_DATABASE_SQLITE_URI, 

18 } 

19 try: 

20 return db_urls[self.SQLALCHEMY_DATABASE] 

21 except KeyError as e: 

22 msg = ( 

23 f"Invalid database type: {self.SQLALCHEMY_DATABASE}. " 

24 f"Supported types: {', '.join(db_urls.keys())}" 

25 ) 

26 raise ValueError(msg) from e 

27 

28 

29settings = Settings() 

30 

31url = settings.database_url 

32 

33db_engine = create_engine( 

34 url=url, 

35 echo=False, # SQL見る場合に有効化する 

36) 

37 

38# 手動実行用途で分離 

39@contextmanager 

40def get_session_context() -> Generator[Session]: 

41 with Session(db_engine) as session: 

42 yield session 

43 

44 

45def get_session() -> Generator[Session]: 

46 with get_session_context() as session: 

47 yield session