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

1""" 

2Queue service for managing tasks using encrypted user databases. 

3Replaces the service_db approach with direct access to user databases. 

4""" 

5 

6from datetime import datetime, timedelta, UTC 

7from typing import Any, Dict, List, Optional 

8 

9from loguru import logger 

10from sqlalchemy.orm import Session 

11 

12from .models import QueueStatus, TaskMetadata 

13 

14 

15class UserQueueService: 

16 """Manages queue operations within a user's encrypted database.""" 

17 

18 def __init__(self, session: Session): 

19 """ 

20 Initialize with a database session. 

21 

22 Args: 

23 session: SQLAlchemy session for the user's encrypted database 

24 """ 

25 self.session = session 

26 

27 def _safe_commit(self) -> None: 

28 """Commit the current transaction, rolling back on failure.""" 

29 try: 

30 self.session.commit() 

31 except Exception: 

32 self.session.rollback() 

33 logger.exception("Database commit failed, transaction rolled back") 

34 raise 

35 

36 def update_queue_status( 

37 self, 

38 active_tasks: int, 

39 queued_tasks: int, 

40 last_task_id: Optional[str] = None, 

41 ) -> None: 

42 """Update queue status for the user.""" 

43 status = self.session.query(QueueStatus).first() 

44 

45 if status: 

46 status.active_tasks = active_tasks 

47 status.queued_tasks = queued_tasks 

48 status.last_checked = datetime.now(UTC) 

49 if last_task_id: 

50 status.last_task_id = last_task_id 

51 else: 

52 status = QueueStatus( 

53 active_tasks=active_tasks, 

54 queued_tasks=queued_tasks, 

55 last_task_id=last_task_id, 

56 ) 

57 self.session.add(status) 

58 

59 self._safe_commit() 

60 

61 def get_queue_status(self) -> Optional[Dict[str, Any]]: 

62 """Get queue status for the user.""" 

63 status = self.session.query(QueueStatus).first() 

64 

65 if status: 

66 return { 

67 "active_tasks": status.active_tasks, 

68 "queued_tasks": status.queued_tasks, 

69 "last_checked": status.last_checked, 

70 "last_task_id": status.last_task_id, 

71 } 

72 return None 

73 

74 def add_task_metadata( 

75 self, 

76 task_id: str, 

77 task_type: str, 

78 priority: int = 0, 

79 ) -> None: 

80 """Add metadata for a new task.""" 

81 task = TaskMetadata( 

82 task_id=task_id, 

83 status="queued", 

84 task_type=task_type, 

85 priority=priority, 

86 ) 

87 self.session.add(task) 

88 

89 if task_type == "research": 

90 # Update queue counts 

91 self._increment_queue_count() 

92 

93 self._safe_commit() 

94 

95 def update_task_status( 

96 self, task_id: str, status: str, error_message: Optional[str] = None 

97 ) -> None: 

98 """Update task status.""" 

99 task = ( 

100 self.session.query(TaskMetadata).filter_by(task_id=task_id).first() 

101 ) 

102 

103 if task: 

104 old_status = task.status 

105 task.status = status 

106 task.error_message = error_message 

107 if status in ["completed", "failed"]: 

108 task.completed_at = datetime.now(UTC) 

109 

110 if task.task_type == "research": 

111 if status == "processing" and old_status == "queued": 

112 task.started_at = datetime.now(UTC) 

113 self._update_queue_counts(-1, 1) # -1 queued, +1 active 

114 

115 elif status == "queued" and old_status == "processing": 

116 # Dispatch was rolled back (e.g. global capacity reject left 

117 # the research queued for the next tick). Revert the claim so 

118 # the counters stay balanced instead of leaking a slot into 

119 # active_tasks on every retry. 

120 task.started_at = None 

121 self._update_queue_counts(1, -1) # +1 queued, -1 active 

122 

123 elif status in ["completed", "failed"]: 

124 self._update_queue_counts(0, -1) # 0 queued, -1 active 

125 

126 self._safe_commit() 

127 

128 def get_pending_tasks(self, limit: int = 50) -> List[Dict[str, Any]]: 

129 """Get pending tasks for the user.""" 

130 tasks = ( 

131 self.session.query(TaskMetadata) 

132 .filter_by(status="queued") 

133 .order_by(TaskMetadata.priority.desc(), TaskMetadata.created_at) 

134 .limit(limit) 

135 .all() 

136 ) 

137 

138 return [ 

139 { 

140 "task_id": t.task_id, 

141 "task_type": t.task_type, 

142 "created_at": t.created_at, 

143 "priority": t.priority, 

144 } 

145 for t in tasks 

146 ] 

147 

148 def cleanup_old_tasks(self, days: int = 7) -> int: 

149 """Clean up old completed/failed tasks.""" 

150 cutoff_date = datetime.now(UTC) - timedelta(days=days) 

151 

152 deleted = ( 

153 self.session.query(TaskMetadata) 

154 .filter( 

155 TaskMetadata.status.in_(["completed", "failed"]), 

156 TaskMetadata.completed_at < cutoff_date, 

157 ) 

158 .delete() 

159 ) 

160 

161 self._safe_commit() 

162 return deleted 

163 

164 def _get_or_create_status(self) -> QueueStatus: 

165 """Get existing queue status or create a new one with zero counts.""" 

166 status = self.session.query(QueueStatus).first() 

167 if status is None: 

168 status = QueueStatus(queued_tasks=0, active_tasks=0) 

169 self.session.add(status) 

170 return status 

171 

172 def _increment_queue_count(self): 

173 """Increment the queued task count.""" 

174 status = self._get_or_create_status() 

175 status.queued_tasks += 1 

176 status.last_checked = datetime.now(UTC) 

177 

178 def _update_queue_counts(self, queued_delta: int, active_delta: int): 

179 """Update queue counts by deltas.""" 

180 status = self._get_or_create_status() 

181 status.queued_tasks = max(0, status.queued_tasks + queued_delta) 

182 status.active_tasks = max(0, status.active_tasks + active_delta) 

183 status.last_checked = datetime.now(UTC)