Coverage for src / local_deep_research / web / auth / database_middleware.py: 97%

1""" 

2Middleware to ensure database connections are available for authenticated users. 

3""" 

4 

5from flask import g, session 

6from loguru import logger 

7 

8from ...database.encrypted_db import db_manager 

9from .middleware_optimizer import should_skip_database_middleware 

10 

11 

12def ensure_user_database(): 

13 """ 

14 Ensure the user's database is open for the current request. 

15 This is called as a before_request handler. 

16 """ 

17 # Skip for requests that don't need database access 

18 if should_skip_database_middleware(): 

19 return 

20 

21 # Check if we already ran for this request 

22 if getattr(g, "_db_initialized", False): 

23 return # Already set up 

24 

25 username = session.get("username") 

26 if username: 

27 # Try to get password from various sources 

28 password = None 

29 

30 # Check for temporary auth token (post-registration/login) 

31 temp_auth_token = session.get("temp_auth_token") 

32 if temp_auth_token: 

33 from ...database.temp_auth import temp_auth_store 

34 

35 auth_data = temp_auth_store.retrieve_auth(temp_auth_token) 

36 if auth_data: 36 ↛ 54line 36 didn't jump to line 54 because the condition on line 36 was always true

37 stored_username, password = auth_data 

38 if stored_username == username: 

39 # Remove token from session after use 

40 session.pop("temp_auth_token", None) 

41 

42 # Store in session password store for future requests 

43 session_id = session.get("session_id") 

44 if session_id: 44 ↛ 54line 44 didn't jump to line 54 because the condition on line 44 was always true

45 from ...database.session_passwords import ( 

46 session_password_store, 

47 ) 

48 

49 session_password_store.store_session_password( 

50 username, session_id, password 

51 ) 

52 

53 # If no password from temp auth, try session password store 

54 if not password: 

55 session_id = session.get("session_id") 

56 if session_id: 

57 from ...database.session_passwords import session_password_store 

58 

59 password = session_password_store.get_session_password( 

60 username, session_id 

61 ) 

62 

63 # For unencrypted databases, use dummy password 

64 if not password and not db_manager.has_encryption: 

65 password = "dummy" 

66 

67 # If we have a password, ensure the user's database is open 

68 if password: 

69 try: 

70 if not db_manager.is_user_connected(username): 

71 engine = db_manager.open_user_database(username, password) 

72 if not engine: 

73 logger.warning( 

74 f"open_user_database returned None for user {username}" 

75 ) 

76 return 

77 g._db_initialized = True 

78 g.user_password = password 

79 g.username = username 

80 except Exception: 

81 logger.warning( 

82 f"Failed to open database for user {username}", 

83 )