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

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 DateTime, Integer, String, Text 

13from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column 

14 

15 

16class Base(DeclarativeBase): 

17 pass 

18 

19 

20class FailureType(str, Enum): 

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

22 

23 NOT_FOUND = "not_found" 

24 FORBIDDEN = "forbidden" 

25 GONE = "gone" 

26 RATE_LIMITED = "rate_limited" 

27 SERVER_ERROR = "server_error" 

28 RECAPTCHA_PROTECTION = "recaptcha_protection" 

29 INCOMPATIBLE_FORMAT = "incompatible_format" 

30 TIMEOUT = "timeout" 

31 NETWORK_ERROR = "network_error" 

32 UNKNOWN_ERROR = "unknown_error" 

33 

34 

35class DownloadStatus(str, Enum): 

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

37 

38 AVAILABLE = "available" 

39 TEMPORARILY_FAILED = "temporarily_failed" 

40 PERMANENTLY_FAILED = "permanently_failed" 

41 

42 

43class ResourceDownloadStatus(Base): 

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

45 

46 __tablename__ = "resource_download_status" 

47 

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

49 resource_id: Mapped[int] = mapped_column( 

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

51 ) 

52 

53 # Status tracking 

54 status: Mapped[str] = mapped_column( 

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

56 ) # available, temporarily_failed, permanently_failed 

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

58 String(100), nullable=True 

59 ) # not_found, rate_limited, timeout, etc. 

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

61 

62 # Retry timing 

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

64 DateTime, nullable=True 

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

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

67 DateTime, nullable=True 

68 ) 

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

70 DateTime, nullable=True 

71 ) # When permanently failed 

72 

73 # Statistics 

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

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

76 

77 # Timestamps 

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

79 DateTime, default=partial(datetime.now, UTC) 

80 ) 

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

82 DateTime, 

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

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

85 ) 

86 

87 def __repr__(self) -> str: 

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