Coverage for src/local_deep_research/library/download_management/status_tracker.py: 94%
146 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Resource Status Tracker
4Tracks download attempts, failures, and cooldowns in the database.
5Provides persistent storage for failure classifications and retry eligibility.
6"""
8from datetime import datetime, timedelta, UTC
9from typing import Optional, Dict, Any
11from loguru import logger
12from sqlalchemy.orm import sessionmaker, Session
14from .models import ResourceDownloadStatus, Base
15from .failure_classifier import BaseFailure
17MAX_TOTAL_RETRIES = 5
20def compute_retry_cooldown(
21 attempt: int, default_cooldown: timedelta
22) -> Optional[timedelta]:
23 """Return cooldown for this attempt number, or None for permanent failure.
25 Schedule:
26 attempt 1: default_cooldown (from failure type)
27 attempt 2: 1 day
28 attempt 3-4: 30 days
29 attempt >= 5: None (permanent failure)
30 """
31 if attempt >= MAX_TOTAL_RETRIES:
32 return None
33 if attempt >= 3:
34 return timedelta(days=30)
35 if attempt == 2:
36 return timedelta(days=1)
37 return default_cooldown
40class ResourceStatusTracker:
41 """Track download attempts, failures, and cooldowns in database"""
43 def __init__(self, username: str, password: Optional[str] = None):
44 """
45 Initialize the status tracker for a user.
47 Args:
48 username: Username for database access
49 password: Optional password for encrypted database
50 """
51 self.username = username
52 self.password = password
54 # Use the global db_manager singleton to share cached connections
55 from ...database.encrypted_db import (
56 DatabaseInitializationError,
57 db_manager,
58 )
60 self.db_manager = db_manager
61 try:
62 self.engine = db_manager.open_user_database(username, password)
63 except DatabaseInitializationError:
64 # Surface init failures from the schedulers/library-init
65 # callers as a plain RuntimeError — they all wrap construction
66 # in try/except already, and propagating the typed exception
67 # would couple every caller to encrypted_db's internals.
68 # ``logger.warning`` (no traceback) rather than
69 # ``logger.exception``: ``password`` is a live local in this
70 # frame, so rendering a traceback under ``diagnose=True``
71 # would dump the plaintext SQLCipher master password
72 # (unrecoverable — TRUST.md §5). The redacted failure detail
73 # is already logged at the raise site in
74 # ``open_user_database`` (#4182).
75 logger.warning(
76 f"[STATUS_TRACKER] Database init failed for user: {username}"
77 )
78 raise RuntimeError(
79 f"Database initialisation failed for user {username}"
80 ) from None
81 self.Session = sessionmaker(bind=self.engine)
83 # Create tables if they don't exist
84 Base.metadata.create_all(self.engine)
85 logger.info(
86 f"[STATUS_TRACKER] Initialized for user: {username} with encrypted database"
87 )
89 def _get_session(self) -> Session:
90 """Get a database session"""
91 return self.Session()
93 def mark_failure(
94 self,
95 resource_id: int,
96 failure: BaseFailure,
97 session: Optional[Session] = None,
98 commit: bool = True,
99 ) -> None:
100 """
101 Mark a resource as failed with classification.
103 Args:
104 resource_id: Resource identifier
105 failure: Classified failure object
106 session: Optional existing database session to reuse
107 commit: If True (default) and ``session`` is provided, commit
108 after applying the status update. Set to False when the
109 caller manages the transaction (e.g. batching multiple
110 status updates). The internal-session path always commits.
111 """
112 if session is not None: 112 ↛ 118line 112 didn't jump to line 118 because the condition on line 112 was always true
113 self._apply_failure(session, resource_id, failure)
114 if commit:
115 session.commit()
116 return
118 with self._get_session() as session:
119 self._apply_failure(session, resource_id, failure)
120 session.commit()
122 def _apply_failure(
123 self, session: Session, resource_id: int, failure: BaseFailure
124 ) -> None:
125 """Apply failure status updates to a session (does not commit)."""
126 # Get or create status record
127 status = (
128 session.query(ResourceDownloadStatus)
129 .filter_by(resource_id=resource_id)
130 .first()
131 )
133 if not status:
134 status = ResourceDownloadStatus(resource_id=resource_id)
135 session.add(status)
137 # Update status information
138 if failure.is_permanent():
139 status.status = "permanently_failed"
140 status.retry_after_timestamp = None
141 status.failure_type = failure.error_type
142 status.failure_message = failure.message
143 status.permanent_failure_at = datetime.now(UTC)
144 logger.info(
145 f"[STATUS_TRACKER] Marked resource {resource_id} as permanently failed: {failure.error_type}"
146 )
147 else:
148 status.failure_type = failure.error_type
149 status.failure_message = failure.message
151 attempt = (status.total_retry_count or 0) + 1
152 cooldown = compute_retry_cooldown(attempt, failure.retry_after)
154 if cooldown is None:
155 status.status = "permanently_failed"
156 status.permanent_failure_at = datetime.now(UTC)
157 status.retry_after_timestamp = None
158 logger.info(
159 f"[STATUS_TRACKER] Auto-promoted resource {resource_id} to "
160 f"permanently failed after {attempt} attempts"
161 )
162 else:
163 status.status = "temporarily_failed"
164 status.retry_after_timestamp = datetime.now(UTC) + cooldown
165 logger.info(
166 f"[STATUS_TRACKER] Marked resource {resource_id} as temporarily failed: "
167 f"{failure.error_type}, attempt {attempt}, retry after: {cooldown}"
168 )
170 # Update retry statistics
171 # Ensure total_retry_count is initialized (handle None from legacy data)
172 if status.total_retry_count is None:
173 status.total_retry_count = 0
174 status.total_retry_count += 1
176 # Check if last attempt was today (before overwriting last_attempt_at)
177 today = datetime.now(UTC).date()
178 last_attempt = (
179 status.last_attempt_at.date() if status.last_attempt_at else None
180 )
181 status.last_attempt_at = datetime.now(UTC)
182 if last_attempt == today:
183 # Ensure today_retry_count is initialized (handle None from legacy data)
184 if status.today_retry_count is None:
185 status.today_retry_count = 0
186 status.today_retry_count += 1
187 else:
188 status.today_retry_count = 1
190 logger.debug(
191 f"[STATUS_TRACKER] Updated failure status for resource {resource_id}"
192 )
194 def mark_success(
195 self,
196 resource_id: int,
197 session: Optional[Session] = None,
198 commit: bool = True,
199 ) -> None:
200 """
201 Mark a resource as successfully downloaded.
203 Args:
204 resource_id: Resource identifier
205 session: Optional existing database session to reuse
206 commit: If True (default) and ``session`` is provided, commit
207 after applying the status update. Set to False when the
208 caller manages the transaction (e.g. batching multiple
209 status updates). The internal-session path always commits.
210 """
211 if session is not None: 211 ↛ 217line 211 didn't jump to line 217 because the condition on line 211 was always true
212 self._apply_success(session, resource_id)
213 if commit:
214 session.commit()
215 return
217 with self._get_session() as session:
218 self._apply_success(session, resource_id)
219 session.commit()
221 def _apply_success(self, session: Session, resource_id: int) -> None:
222 """Apply success status updates to a session (does not commit)."""
223 status = (
224 session.query(ResourceDownloadStatus)
225 .filter_by(resource_id=resource_id)
226 .first()
227 )
229 if status:
230 status.status = "completed"
231 status.failure_type = None
232 status.failure_message = None
233 status.retry_after_timestamp = None
234 status.updated_at = datetime.now(UTC)
235 logger.info(
236 f"[STATUS_TRACKER] Marked resource {resource_id} as successfully completed"
237 )
239 def can_retry(self, resource_id: int) -> tuple[bool, Optional[str]]:
240 """
241 Check if a resource can be retried right now.
243 Args:
244 resource_id: Resource identifier
246 Returns:
247 Tuple of (can_retry, reason_if_not)
248 """
249 with self._get_session() as session:
250 status = (
251 session.query(ResourceDownloadStatus)
252 .filter_by(resource_id=resource_id)
253 .first()
254 )
256 if not status:
257 # No status record, can retry
258 return True, None
260 # Check if permanently failed
261 if status.status == "permanently_failed":
262 return (
263 False,
264 f"Permanently failed: {status.failure_message or status.failure_type}",
265 )
267 # Check if temporarily failed and cooldown not expired
268 if (
269 status.status == "temporarily_failed"
270 and status.retry_after_timestamp
271 ):
272 # Ensure retry_after_timestamp is timezone-aware (handle legacy data)
273 retry_timestamp = status.retry_after_timestamp
274 if retry_timestamp.tzinfo is None:
275 # Assume UTC for naive timestamps
276 retry_timestamp = retry_timestamp.replace(tzinfo=UTC)
278 if datetime.now(UTC) < retry_timestamp:
279 return (
280 False,
281 f"Cooldown active, retry available at {retry_timestamp.strftime('%Y-%m-%d %H:%M:%S')}",
282 )
284 # Check daily retry limit (max 3 retries per day)
285 # today_retry_count is only reset inside _apply_failure, so check
286 # whether the stored count is actually from today before using it.
287 today = datetime.now(UTC).date()
288 last_attempt_date = (
289 status.last_attempt_at.date()
290 if status.last_attempt_at
291 else None
292 )
293 today_count = (
294 status.today_retry_count if last_attempt_date == today else 0
295 )
296 if today_count >= 3:
297 return (
298 False,
299 f"Daily retry limit exceeded ({today_count}/3). Retry available tomorrow.",
300 )
302 # Check total retry limit (safety net for records not yet auto-promoted)
303 if (status.total_retry_count or 0) >= MAX_TOTAL_RETRIES:
304 return (
305 False,
306 f"Permanently failed after {status.total_retry_count} attempts. Will not retry.",
307 )
309 # Can retry
310 return True, None
312 def get_resource_status(self, resource_id: int) -> Optional[Dict[str, Any]]:
313 """
314 Get the current status of a resource.
316 Args:
317 resource_id: Resource identifier
319 Returns:
320 Status information dictionary or None if not found
321 """
322 with self._get_session() as session:
323 status = (
324 session.query(ResourceDownloadStatus)
325 .filter_by(resource_id=resource_id)
326 .first()
327 )
329 if not status:
330 return None
332 return {
333 "resource_id": status.resource_id,
334 "status": status.status,
335 "failure_type": status.failure_type,
336 "failure_message": status.failure_message,
337 "retry_after_timestamp": status.retry_after_timestamp.isoformat()
338 if status.retry_after_timestamp
339 else None,
340 "last_attempt_at": status.last_attempt_at.isoformat()
341 if status.last_attempt_at
342 else None,
343 "total_retry_count": status.total_retry_count,
344 "today_retry_count": status.today_retry_count,
345 "created_at": status.created_at.isoformat(),
346 "updated_at": status.updated_at.isoformat(),
347 }
349 def get_failed_resources_count(self) -> Dict[str, int]:
350 """
351 Get counts of resources by failure type.
353 Returns:
354 Dictionary mapping failure types to counts
355 """
356 with self._get_session() as session:
357 failed_resources = (
358 session.query(ResourceDownloadStatus)
359 .filter(
360 ResourceDownloadStatus.status.in_(
361 ["temporarily_failed", "permanently_failed"]
362 )
363 )
364 .all()
365 )
367 counts = {}
368 for resource in failed_resources:
369 failure_type = resource.failure_type or "unknown"
370 counts[failure_type] = counts.get(failure_type, 0) + 1
372 return counts
374 def clear_permanent_failures(self, older_than_days: int = 30) -> int:
375 """
376 Clear permanent failure statuses for old records.
378 Args:
379 older_than_days: Clear failures older than this many days
381 Returns:
382 Number of records cleared
383 """
384 cutoff_date = datetime.now(UTC) - timedelta(days=older_than_days)
386 with self._get_session() as session:
387 old_failures = (
388 session.query(ResourceDownloadStatus)
389 .filter(
390 ResourceDownloadStatus.status == "permanently_failed",
391 ResourceDownloadStatus.created_at < cutoff_date,
392 )
393 .all()
394 )
396 count = len(old_failures)
397 for failure in old_failures:
398 failure.status = "available"
399 failure.failure_type = None
400 failure.failure_message = None
401 failure.retry_after_timestamp = None
402 failure.permanent_failure_at = None
403 failure.total_retry_count = 0
404 failure.today_retry_count = 0
405 failure.updated_at = datetime.now(UTC)
407 session.commit()
408 logger.info(
409 f"[STATUS_TRACKER] Cleared {count} old permanent failure records"
410 )
411 return count