Coverage for src/local_deep_research/metrics/database.py: 100%

1"""Database utilities for metrics module with SQLAlchemy.""" 

2 

3from contextlib import contextmanager 

4from typing import Generator, Optional 

5 

6from sqlalchemy.orm import Session 

7 

8from ..database.session_context import get_user_db_session 

9from ..utilities.request_context import get_current_username # noqa: F401 

10 

11 

12class MetricsDatabase: 

13 """Database manager for metrics using SQLAlchemy.""" 

14 

15 def __init__( 

16 self, username: Optional[str] = None, password: Optional[str] = None 

17 ): 

18 # Store credentials if provided (for testing/background tasks) 

19 self.username = username 

20 self.password = password 

21 

22 @contextmanager 

23 def get_session( 

24 self, username: Optional[str] = None, password: Optional[str] = None 

25 ) -> Generator[Session, None, None]: 

26 """Get a database session with automatic cleanup. 

27 

28 Args: 

29 username: Override username for this session 

30 password: Override password for this session (if needed for thread access) 

31 """ 

32 # Use provided username or fall back to stored/session username 

33 user = username or self.username 

34 pwd = password or self.password 

35 

36 # Try to get username from request context if still not available 

37 if not user: 

38 user = get_current_username() 

39 

40 if not user: 

41 # No username available - can't access user database 

42 yield None 

43 return 

44 

45 # If we have password, use thread-safe access 

46 if pwd: 

47 from ..database.thread_metrics import metrics_writer 

48 

49 metrics_writer.set_user_password(user, pwd) 

50 with metrics_writer.get_session(user) as session: 

51 yield session 

52 else: 

53 # Use the per-user database session with proper context management 

54 with get_user_db_session(user) as session: 

55 yield session