Coverage for src / quber / db / connection.py: 100%

1"""Database connection utilities for Quber RAG system.""" 

2 

3from contextlib import contextmanager 

4from typing import Generator 

5 

6from loguru import logger 

7from sqlalchemy import create_engine 

8from sqlalchemy.engine import Engine 

9from sqlalchemy.orm import Session, sessionmaker 

10 

11from quber.db.models import Base 

12from quber.settings import get_settings 

13 

14 

15def get_database_url() -> str: 

16 """ 

17 Construct database URL from environment variables. 

18 

19 Returns: 

20 Database connection URL string. 

21 

22 Raises: 

23 ValueError: If required environment variables are missing. 

24 """ 

25 db_settings = get_settings().db 

26 user = db_settings.user 

27 password = db_settings.password 

28 host = db_settings.host 

29 port = db_settings.port 

30 db = db_settings.db 

31 

32 url = f"postgresql+psycopg://{user}:{password}@{host}:{port}/{db}" 

33 logger.debug(f"Database URL: postgresql+psycopg://{user}:***@{host}:{port}/{db}") 

34 return url 

35 

36 

37def get_engine(echo: bool = False) -> Engine: 

38 """ 

39 Create and return a SQLAlchemy engine. 

40 

41 Args: 

42 echo: If True, log all SQL statements (useful for debugging). 

43 

44 Returns: 

45 SQLAlchemy Engine instance. 

46 """ 

47 url = get_database_url() 

48 engine = create_engine(url, echo=echo, pool_pre_ping=True) 

49 logger.info(f"Database engine created: {engine.url.database}@{engine.url.host}") 

50 return engine 

51 

52 

53def init_db(engine: Engine | None = None, drop_all: bool = False) -> None: 

54 """ 

55 Initialize database schema. 

56 

57 Args: 

58 engine: SQLAlchemy engine. If None, creates a new one. 

59 drop_all: If True, drop all tables before creating (WARNING: data loss). 

60 """ 

61 if engine is None: 

62 engine = get_engine() 

63 

64 if drop_all: 

65 logger.warning("Dropping all tables...") 

66 Base.metadata.drop_all(engine) 

67 

68 logger.info("Creating database tables...") 

69 Base.metadata.create_all(engine) 

70 logger.success("Database initialized successfully") 

71 

72 

73def get_session_factory(engine: Engine | None = None) -> sessionmaker[Session]: 

74 """ 

75 Create a session factory. 

76 

77 Args: 

78 engine: SQLAlchemy engine. If None, creates a new one. 

79 

80 Returns: 

81 SQLAlchemy sessionmaker. 

82 """ 

83 if engine is None: 

84 engine = get_engine() 

85 

86 return sessionmaker(bind=engine, expire_on_commit=False) 

87 

88 

89@contextmanager 

90def get_session(engine: Engine | None = None) -> Generator[Session, None, None]: 

91 """ 

92 Context manager for database sessions. 

93 

94 Args: 

95 engine: SQLAlchemy engine. If None, creates a new one. 

96 

97 Yields: 

98 Database session. 

99 

100 Example: 

101 ```python 

102 with get_session() as session: 

103 documents = session.query(Document).all() 

104 ``` 

105 """ 

106 if engine is None: 

107 engine = get_engine() 

108 

109 SessionLocal = sessionmaker(bind=engine, expire_on_commit=False) 

110 session = SessionLocal() 

111 try: 

112 yield session 

113 session.commit() 

114 except Exception as e: 

115 session.rollback() 

116 logger.error(f"Session error: {e}") 

117 raise 

118 finally: 

119 session.close()