Coverage for backend/django/core/auxiliary/services/csv_lifecycle.py: 89%

1from datetime import datetime, timedelta 

2 

3from django.conf import settings 

4from django.utils import timezone 

5 

6from core.auxiliary.models.UploadSession import UploadSessionPurpose 

7from core.auxiliary.services.object_storage import s3 as s3_storage 

8 

9 

10# Initial browser-managed uploads get a short retention window so abandoned 

11# multipart sessions stop consuming object-storage space quickly. 

12SHORT_CSV_LIFECYCLE_TTL = timedelta(hours=1) 

13# Completed CSV uploads expire after a week by default. Scenario and pinch CSVs 

14# also use this window so delayed follow-up jobs do not depend on faster object 

15# storage cleanup. 

16STANDARD_COMPLETED_CSV_LIFECYCLE_TTL = timedelta(days=7) 

17CSV_UPLOAD_PREFIX = "csv-uploads/" 

18CSV_INCOMPLETE_MULTIPART_ABORT_DAYS = 1 

19 

20 

21def completed_csv_lifecycle_ttl(purpose: str) -> timedelta: 

22 """Return the completed-object retention window for the given CSV purpose.""" 

23 if purpose == UploadSessionPurpose.ML_TRAINING_CSV: 

24 return timedelta(days=settings.SEAWEED_ML_UPLOAD_RETENTION_DAYS) 

25 if purpose in { 25 ↛ 30line 25 didn't jump to line 30 because the condition on line 25 was always true

26 UploadSessionPurpose.SCENARIO_CSV, 

27 UploadSessionPurpose.PINCH_UTILITY_CSV, 

28 }: 

29 return STANDARD_COMPLETED_CSV_LIFECYCLE_TTL 

30 raise ValueError(f"Unsupported CSV upload purpose: {purpose}") 

31 

32 

33def _completed_csv_lifecycle_days(purpose: str) -> int: 

34 """Return the whole-day storage lifecycle window for completed CSV objects.""" 

35 ttl = completed_csv_lifecycle_ttl(purpose) 

36 days = ttl.days 

37 if ttl.seconds or ttl.microseconds: 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true

38 days += 1 

39 return max(days, 1) 

40 

41 

42def csv_upload_lifecycle_rules() -> list[dict]: 

43 """Build stable prefix lifecycle rules for CSV upload storage cleanup.""" 

44 rules = [ 

45 { 

46 "ID": f"ahuora-csv-expire-{purpose}", 

47 "Status": "Enabled", 

48 "Filter": {"Prefix": f"{CSV_UPLOAD_PREFIX}{purpose}/"}, 

49 "Expiration": {"Days": _completed_csv_lifecycle_days(purpose)}, 

50 } 

51 for purpose in ( 

52 UploadSessionPurpose.ML_TRAINING_CSV, 

53 UploadSessionPurpose.SCENARIO_CSV, 

54 UploadSessionPurpose.PINCH_UTILITY_CSV, 

55 ) 

56 ] 

57 rules.append( 

58 { 

59 "ID": "ahuora-csv-abort-incomplete-multipart", 

60 "Status": "Enabled", 

61 "Filter": {"Prefix": CSV_UPLOAD_PREFIX}, 

62 "AbortIncompleteMultipartUpload": { 

63 "DaysAfterInitiation": CSV_INCOMPLETE_MULTIPART_ABORT_DAYS, 

64 }, 

65 } 

66 ) 

67 return rules 

68 

69 

70def ensure_csv_upload_lifecycle_configuration(bucket: str) -> None: 

71 """Install stable CSV lifecycle rules without adding per-object rules.""" 

72 s3_storage.ensure_lifecycle_rules( 

73 bucket, 

74 csv_upload_lifecycle_rules(), 

75 obsolete_expiration_prefixes=(CSV_UPLOAD_PREFIX,), 

76 ) 

77 

78 

79def expires_at_from_ttl(ttl: timedelta, *, reference_time: datetime | None = None) -> datetime: 

80 """Convert a retention delta into an absolute expiry timestamp.""" 

81 return (reference_time or timezone.now()) + ttl