Coverage for services/data-ingestion-web/src/database.py: 0%

1"""Database connection and session management""" 

2import os 

3from typing import Generator 

4from sqlalchemy import create_engine 

5from sqlalchemy.orm import sessionmaker, Session 

6from sqlalchemy.pool import NullPool 

7# Create this file: 

8# /c:/Users/aless/Documents/Projects/heimdall/services/data-ingestion-web/src/models/__init__.py 

9 

10from sqlalchemy.orm import declarative_base 

11 

12Base = declarative_base() 

13 

14# Connection string from environment 

15POSTGRES_USER = os.getenv("POSTGRES_USER", "heimdall_user") 

16POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "changeme") 

17POSTGRES_HOST = os.getenv("POSTGRES_HOST", "postgres") 

18POSTGRES_PORT = os.getenv("POSTGRES_PORT", "5432") 

19POSTGRES_DB = os.getenv("POSTGRES_DB", "heimdall") 

20 

21DATABASE_URL = f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{POSTGRES_HOST}:{POSTGRES_PORT}/{POSTGRES_DB}" 

22 

23# Engine configuration 

24engine = create_engine( 

25 DATABASE_URL, 

26 poolclass=NullPool, # No pooling (better for async) 

27 echo=False, 

28) 

29 

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

31 

32 

33def get_db() -> Generator[Session, None, None]: 

34 """Dependency for FastAPI to get a DB session""" 

35 db = SessionLocal() 

36 try: 

37 yield db 

38 finally: 

39 db.close() 

40 

41 

42def init_db(): 

43 """Initialize database tables""" 

44 Base.metadata.create_all(bind=engine)