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 

9 

10 

11class MetricsDatabase: 

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

13 

14 def __init__( 

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

16 ): 

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

18 self.username = username 

19 self.password = password 

20 

21 @contextmanager 

22 def get_session( 

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

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

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

26 

27 Args: 

28 username: Override username for this session 

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

30 """ 

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

32 user = username or self.username 

33 pwd = password or self.password 

34 

35 # Try to get username from Flask session if still not available 

36 if not user: 

37 try: 

38 from flask import session as flask_session 

39 

40 user = flask_session.get("username") 

41 except: 

42 pass 

43 

44 if not user: 

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

46 yield None 

47 return 

48 

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

50 if pwd: 

51 from ..database.thread_metrics import metrics_writer 

52 

53 metrics_writer.set_user_password(user, pwd) 

54 with metrics_writer.get_session(user) as session: 

55 yield session 

56 else: 

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

58 with get_user_db_session(user) as session: 

59 yield session 

60 

61 

62# Singleton instance 

63_metrics_db = None 

64 

65 

66def get_metrics_db() -> MetricsDatabase: 

67 """Get the singleton metrics database instance.""" 

68 global _metrics_db 

69 if _metrics_db is None: 

70 _metrics_db = MetricsDatabase() 

71 return _metrics_db