Coverage for src / local_deep_research / web / auth / cleanup_middleware.py: 96%

1""" 

2Middleware to clean up completed research records. 

3Runs in request context where we have database access. 

4""" 

5 

6import random 

7 

8from flask import session 

9from loguru import logger 

10from sqlalchemy.exc import OperationalError, PendingRollbackError, TimeoutError 

11 

12from ...database.models import UserActiveResearch 

13from ...database.session_context import get_g_db_session 

14from .middleware_optimizer import should_skip_database_middleware 

15 

16 

17def cleanup_completed_research(): 

18 """ 

19 Clean up completed research records for the current user. 

20 Called as a before_request handler. 

21 

22 Runs on only ~1% of requests to avoid unnecessary DB queries on 

23 every authenticated request (same sampling pattern used by 

24 session cleanup in middleware_optimizer.should_skip_session_cleanup). 

25 """ 

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

27 if should_skip_database_middleware(): 

28 return 

29 

30 # Only run cleanup on ~1% of requests to reduce per-request overhead 

31 if random.randint(1, 100) > 1: # noqa: S311 — not cryptographic, just sampling 

32 return 

33 

34 username = session.get("username") 

35 if not username: 

36 return 

37 

38 db_session = get_g_db_session() 

39 if db_session: 

40 try: 

41 # Find completed researches that haven't been cleaned up 

42 from ..routes.globals import is_research_active 

43 

44 # Get all active records for this user with limit and better error handling 

45 active_records = ( 

46 db_session.query(UserActiveResearch) 

47 .filter_by(username=username) 

48 .limit(50) # Limit to prevent excessive memory usage 

49 .all() 

50 ) 

51 

52 cleaned_count = 0 

53 for record in active_records: 

54 # Check if this research is still active 

55 if not is_research_active(record.research_id): 

56 # Research has completed, clean up the record 

57 db_session.delete(record) 

58 cleaned_count += 1 

59 logger.debug( 

60 f"Cleaned up completed research {record.research_id} " 

61 f"for user {username}" 

62 ) 

63 

64 if cleaned_count > 0: 

65 db_session.commit() 

66 logger.info( 

67 f"Cleaned up {cleaned_count} completed research records " 

68 f"for user {username}" 

69 ) 

70 

71 except (OperationalError, PendingRollbackError, TimeoutError) as e: 

72 # Handle connection pool exhaustion and timeout errors gracefully 

73 logger.warning( 

74 f"Database connection issue during cleanup - skipping: {type(e).__name__}" 

75 ) 

76 try: 

77 db_session.rollback() 

78 except Exception: 

79 # Even rollback might fail if connections are exhausted 

80 logger.warning( 

81 "Could not rollback database session during cleanup error" 

82 ) 

83 except Exception: 

84 # Don't let cleanup errors break the request 

85 logger.exception("Error cleaning up completed research") 

86 try: 

87 db_session.rollback() 

88 except Exception: 

89 logger.warning( 

90 "Could not rollback database session during cleanup error" 

91 )