Coverage for src/local_deep_research/library/download_management/models/__init__.py: 98%

1""" 

2Database Models for Download Management 

3 

4Contains ORM models for tracking resource download status and retry logic. 

5""" 

6 

7from datetime import UTC, datetime 

8from enum import Enum 

9from functools import partial 

10from typing import Optional 

11 

12from sqlalchemy import Integer, String, Text 

13from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column 

14from sqlalchemy_utc import UtcDateTime 

15 

16 

17class Base(DeclarativeBase): 

18 pass 

19 

20 

21class FailureType(str, Enum): 

22 """Enum for failure types - ensures consistency across the codebase""" 

23 

24 PROCESSING = "processing" 

25 NOT_FOUND = "not_found" 

26 FORBIDDEN = "forbidden" 

27 GONE = "gone" 

28 RATE_LIMITED = "rate_limited" 

29 SERVER_ERROR = "server_error" 

30 HTML_NOT_PDF = "html_not_pdf" 

31 RECAPTCHA_PROTECTION = "recaptcha_protection" 

32 INCOMPATIBLE_FORMAT = "incompatible_format" 

33 PAYWALL_OR_LOGIN = "paywall_or_login" 

34 TIMEOUT = "timeout" 

35 NETWORK_ERROR = "network_error" 

36 DOWNLOAD_FAILED = "download_failed" 

37 UNKNOWN_ERROR = "unknown_error" 

38 

39 

40class DownloadStatus(str, Enum): 

41 """Status values for resource download tracking.""" 

42 

43 AVAILABLE = "available" 

44 TEMPORARILY_FAILED = "temporarily_failed" 

45 PERMANENTLY_FAILED = "permanently_failed" 

46 

47 

48class ResourceDownloadStatus(Base): 

49 """Database model for tracking resource download status""" 

50 

51 __tablename__ = "resource_download_status" 

52 

53 id: Mapped[int] = mapped_column(Integer, primary_key=True) 

54 resource_id: Mapped[int] = mapped_column( 

55 Integer, unique=True, nullable=False, index=True 

56 ) 

57 

58 # Status tracking 

59 status: Mapped[str] = mapped_column( 

60 String(50), nullable=False, default="available" 

61 ) # available, temporarily_failed, permanently_failed 

62 failure_type: Mapped[Optional[str]] = mapped_column( 

63 String(100), nullable=True 

64 ) # not_found, rate_limited, timeout, etc. 

65 failure_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True) 

66 

67 # Retry timing. 

68 # 

69 # UtcDateTime, not a bare DateTime: SQLite does not persist tzinfo, so a 

70 # plain DateTime column silently reads back NAIVE even though every writer 

71 # here stores ``datetime.now(UTC)``. That produced a real, silent failure — 

72 # ``get_resource_status()`` returned a naive ``.isoformat()``, and 

73 # ``retry_manager`` then computed ``retry_time - datetime.now(UTC)``, which 

74 # raises TypeError on mixed naive/aware. A bare ``except Exception`` around 

75 # it swallowed the error, so the retry ETA just never appeared in the UI. 

76 # ``can_retry()`` had grown a ``.replace(tzinfo=UTC)`` patch for the same 

77 # root cause at one call site; this fixes it for all of them, including 

78 # rows already written. Matches database/models/library.py, which already 

79 # uses UtcDateTime for its own last_attempt_at. 

80 retry_after_timestamp: Mapped[Optional[datetime]] = mapped_column( 

81 UtcDateTime, nullable=True 

82 ) # When this can be retried (NULL = permanent) 

83 last_attempt_at: Mapped[Optional[datetime]] = mapped_column( 

84 UtcDateTime, nullable=True 

85 ) 

86 permanent_failure_at: Mapped[Optional[datetime]] = mapped_column( 

87 UtcDateTime, nullable=True 

88 ) # When permanently failed 

89 

90 # Statistics 

91 total_retry_count: Mapped[int] = mapped_column(Integer, default=0) 

92 today_retry_count: Mapped[int] = mapped_column(Integer, default=0) 

93 

94 # Timestamps 

95 # UtcDateTime for the same reason as the retry columns above: the 

96 # defaults already produce aware values, but a bare DateTime would read 

97 # them back naive on SQLite. Matches the rest of the app's models. 

98 created_at: Mapped[Optional[datetime]] = mapped_column( 

99 UtcDateTime, default=partial(datetime.now, UTC) 

100 ) 

101 updated_at: Mapped[Optional[datetime]] = mapped_column( 

102 UtcDateTime, 

103 default=partial(datetime.now, UTC), 

104 onupdate=partial(datetime.now, UTC), 

105 ) 

106 

107 def __repr__(self) -> str: 

108 return f"<ResourceDownloadStatus(resource_id={self.resource_id}, status='{self.status}', failure_type='{self.failure_type}')>"