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 NOT_FOUND = "not_found" 

25 FORBIDDEN = "forbidden" 

26 GONE = "gone" 

27 RATE_LIMITED = "rate_limited" 

28 SERVER_ERROR = "server_error" 

29 RECAPTCHA_PROTECTION = "recaptcha_protection" 

30 INCOMPATIBLE_FORMAT = "incompatible_format" 

31 TIMEOUT = "timeout" 

32 NETWORK_ERROR = "network_error" 

33 UNKNOWN_ERROR = "unknown_error" 

34 

35 

36class DownloadStatus(str, Enum): 

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

38 

39 AVAILABLE = "available" 

40 TEMPORARILY_FAILED = "temporarily_failed" 

41 PERMANENTLY_FAILED = "permanently_failed" 

42 

43 

44class ResourceDownloadStatus(Base): 

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

46 

47 __tablename__ = "resource_download_status" 

48 

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

50 resource_id: Mapped[int] = mapped_column( 

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

52 ) 

53 

54 # Status tracking 

55 status: Mapped[str] = mapped_column( 

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

57 ) # available, temporarily_failed, permanently_failed 

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

59 String(100), nullable=True 

60 ) # not_found, rate_limited, timeout, etc. 

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

62 

63 # Retry timing. 

64 # 

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

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

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

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

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

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

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

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

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

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

75 # uses UtcDateTime for its own last_attempt_at. 

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

77 UtcDateTime, nullable=True 

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

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

80 UtcDateTime, nullable=True 

81 ) 

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

83 UtcDateTime, nullable=True 

84 ) # When permanently failed 

85 

86 # Statistics 

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

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

89 

90 # Timestamps 

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

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

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

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

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

96 ) 

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

98 UtcDateTime, 

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

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

101 ) 

102 

103 def __repr__(self) -> str: 

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