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

1""" 

2Middleware to clean up completed research records. 

3Runs in request context where we have database access. 

4""" 

5 

6from flask import g, session 

7from loguru import logger 

8from sqlalchemy.exc import OperationalError, TimeoutError 

9 

10from ...database.models import UserActiveResearch 

11from .middleware_optimizer import should_skip_database_middleware 

12 

13 

14def cleanup_completed_research(): 

15 """ 

16 Clean up completed research records for the current user. 

17 Called as a before_request handler. 

18 """ 

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

20 if should_skip_database_middleware(): 

21 return 

22 

23 username = session.get("username") 

24 

25 if username and hasattr(g, "db_session") and g.db_session: 

26 try: 

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

28 from ..routes.globals import active_research 

29 

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

31 active_records = ( 

32 g.db_session.query(UserActiveResearch) 

33 .filter_by(username=username) 

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

35 .all() 

36 ) 

37 

38 cleaned_count = 0 

39 for record in active_records: 

40 # Check if this research is still active 

41 if record.research_id not in active_research: 

42 # Research has completed, clean up the record 

43 g.db_session.delete(record) 

44 cleaned_count += 1 

45 logger.debug( 

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

47 f"for user {username}" 

48 ) 

49 

50 if cleaned_count > 0: 

51 g.db_session.commit() 

52 logger.info( 

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

54 f"for user {username}" 

55 ) 

56 

57 except (OperationalError, TimeoutError) as e: 

58 # Handle connection pool exhaustion and timeout errors gracefully 

59 logger.warning( 

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

61 ) 

62 try: 

63 g.db_session.rollback() 

64 except Exception: 

65 # Even rollback might fail if connections are exhausted 

66 logger.warning( 

67 "Could not rollback database session during cleanup error" 

68 ) 

69 except Exception: 

70 # Don't let cleanup errors break the request 

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

72 try: 

73 g.db_session.rollback() 

74 except Exception: 

75 logger.warning( 

76 "Could not rollback database session during cleanup error" 

77 )