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 datetime, UTC 

8from enum import Enum 

9from functools import partial 

10from sqlalchemy import Column, Integer, String, Text, DateTime 

11from sqlalchemy.ext.declarative import declarative_base 

12 

13Base = declarative_base() 

14 

15 

16class FailureType(str, Enum): 

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

18 

19 NOT_FOUND = "not_found" 

20 FORBIDDEN = "forbidden" 

21 GONE = "gone" 

22 RATE_LIMITED = "rate_limited" 

23 SERVER_ERROR = "server_error" 

24 RECAPTCHA_PROTECTION = "recaptcha_protection" 

25 INCOMPATIBLE_FORMAT = "incompatible_format" 

26 TIMEOUT = "timeout" 

27 NETWORK_ERROR = "network_error" 

28 UNKNOWN_ERROR = "unknown_error" 

29 

30 

31class DownloadStatus(str, Enum): 

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

33 

34 AVAILABLE = "available" 

35 TEMPORARILY_FAILED = "temporarily_failed" 

36 PERMANENTLY_FAILED = "permanently_failed" 

37 

38 

39class ResourceDownloadStatus(Base): 

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

41 

42 __tablename__ = "resource_download_status" 

43 

44 id = Column(Integer, primary_key=True) 

45 resource_id = Column(Integer, unique=True, nullable=False, index=True) 

46 

47 # Status tracking 

48 status = Column( 

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

50 ) # available, temporarily_failed, permanently_failed 

51 failure_type = Column(String(100)) # not_found, rate_limited, timeout, etc. 

52 failure_message = Column(Text) 

53 

54 # Retry timing 

55 retry_after_timestamp = Column( 

56 DateTime 

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

58 last_attempt_at = Column(DateTime) 

59 permanent_failure_at = Column(DateTime) # When permanently failed 

60 

61 # Statistics 

62 total_retry_count = Column(Integer, default=0) 

63 today_retry_count = Column(Integer, default=0) 

64 

65 # Timestamps 

66 created_at = Column(DateTime, default=partial(datetime.now, UTC)) 

67 updated_at = Column( 

68 DateTime, 

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

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

71 ) 

72 

73 def __repr__(self): 

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