Coverage for backend / app / database.py: 97%

1"""Database connection functions""" 

2 

3import os 

4from contextvars import ContextVar 

5from typing import Generator, Any 

6 

7from sqlalchemy import create_engine 

8from sqlalchemy.ext.declarative import declarative_base 

9from sqlalchemy.orm import sessionmaker, Session 

10 

11from app.config import settings 

12 

13# Context variable to track demo mode per-request 

14demo_mode = ContextVar("demo_mode", default=False) 

15 

16 

17def create_db_url(db_name: str = "") -> str: 

18 """Create a database URL from the database settings.""" 

19 

20 if not db_name: 

21 db_name = settings.database_name 

22 return ( 

23 f"postgresql://{settings.database_username}:{settings.database_password}@" 

24 f"{settings.database_hostname}:{settings.database_port}/{db_name}" 

25 ) 

26 

27 

28SQLALCHEMY_DATABASE_URL = os.getenv("SQLALCHEMY_DATABASE_URL", create_db_url()) 

29engine = create_engine(SQLALCHEMY_DATABASE_URL) 

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

31Base = declarative_base() 

32 

33# Demo engine with search_path set to demo schema 

34demo_engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"options": "-c search_path=demo"}) 

35demo_session_local = sessionmaker(autocommit=False, autoflush=False, bind=demo_engine) 

36 

37 

38def get_demo_db() -> Generator[Session, Any, None]: 

39 """Get a demo schema database session. 

40 :return: the demo database session.""" 

41 

42 db = demo_session_local() 

43 try: 

44 yield db 

45 finally: 

46 db.close() 

47 

48 

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

50 """Get the database session. Uses demo schema if demo_mode context var is set. 

51 :return: the database session.""" 

52 

53 if demo_mode.get(False): 

54 db = demo_session_local() 

55 else: 

56 db = session_local() 

57 try: 

58 yield db 

59 finally: 

60 db.close()