diff --git a/config/template.json b/config/template.json index 70cd31c..bf1a6bd 100644 --- a/config/template.json +++ b/config/template.json @@ -4,7 +4,8 @@ }, "gtfs": { "dir": null, - "refresh_interval_days": 7 + "refresh_interval_days": 7, + "archive_retention_days": 90 }, "DATADOG_TRACE_ENABLED": false, "file_retention_days": 180, diff --git a/src/gtfs.py b/src/gtfs.py index 4c84bef..6a82256 100644 --- a/src/gtfs.py +++ b/src/gtfs.py @@ -27,6 +27,10 @@ GTFS_ARCHIVES_PREFIX = "https://cdn.mbta.com/archive/" GTFS_ARCHIVES_FILENAME = "archived_feeds.txt" +# Number of days to retain downloaded GTFS archives before cleanup. +# The most recent archive is always kept regardless of age. +GTFS_ARCHIVE_RETENTION_DAYS = CONFIG["gtfs"].get("archive_retention_days", 90) + # defining these columns in particular becasue we use them everywhere RTE_DIR_STOP = ["route_id", "direction_id", "stop_id"] @@ -121,6 +125,48 @@ def _find_most_recent_gtfs_archive() -> Optional[pathlib.Path]: return None +@tracer.wrap() +def cleanup_old_gtfs_archives(reference_time: Optional[datetime.datetime] = None) -> int: + """Delete downloaded GTFS archives older than the configured retention period. + + The most recent archive is always retained—even if older than the cutoff—since + it is the active feed and the fallback used when newer archives can't be fetched. + + Args: + reference_time: Datetime used as the reference for the cutoff. + Defaults to the current time. + + Returns: + Number of archive directories deleted. + """ + logger.info("Starting cleanup of old GTFS archives") + if reference_time is None: + reference_time = datetime.datetime.now() + cutoff = reference_time - datetime.timedelta(days=GTFS_ARCHIVE_RETENTION_DAYS) + + try: + archive_dirs = [d for d in MAIN_DIR.iterdir() if d.is_dir()] + except (PermissionError, OSError, IOError) as e: + logger.error(f"Failed to scan GTFS archives directory: {e}") + return 0 + + # Sort by name (date-based) descending and always keep the most recent archive. + archive_dirs.sort(key=lambda d: d.name, reverse=True) + + deleted = 0 + for archive_dir in archive_dirs[1:]: + try: + if datetime.datetime.fromtimestamp(archive_dir.stat().st_mtime) < cutoff: + shutil.rmtree(archive_dir) + deleted += 1 + logger.info(f"Deleted old GTFS archive: {archive_dir}") + except (PermissionError, OSError, IOError) as e: + logger.warning(f"Skipping {archive_dir}: {e}") + + logger.info(f"Completed GTFS archive cleanup — deleted {deleted} archive(s)") + return deleted + + @tracer.wrap() def get_gtfs_archive(dateint: int): """ diff --git a/src/s3_upload.py b/src/s3_upload.py index afce386..27a33bf 100644 --- a/src/s3_upload.py +++ b/src/s3_upload.py @@ -12,6 +12,7 @@ from config import CONFIG from disk import DATA_DIR, cleanup_old_files +from gtfs import cleanup_old_gtfs_archives from logger import set_up_logging from util import EASTERN_TIME, service_date @@ -65,8 +66,9 @@ def upload_todays_events_to_s3(): end_time = time.time() logger.info(f"Uploaded {len(files_updated_today)} files to s3, took {end_time - start_time} seconds.") - # cleanup old files, free up disk space + # cleanup old files and GTFS archives, free up disk space cleanup_old_files(reference_time=start_datetime) + cleanup_old_gtfs_archives(reference_time=start_datetime) @tracer.wrap(service="gobble") diff --git a/src/tests/test_gtfs.py b/src/tests/test_gtfs.py index be00191..bd8cade 100644 --- a/src/tests/test_gtfs.py +++ b/src/tests/test_gtfs.py @@ -1,4 +1,5 @@ import datetime +import os import numpy as np import pandas as pd import pathlib @@ -296,3 +297,61 @@ def test_read_gtfs_date_exists_feed_is_read(self): assert "Harvard" in result.trips_by_route_id("1")["trip_headsign"].values shutil.rmtree(expected_path) + + +class TestCleanupOldGtfsArchives: + """Test cleanup_old_gtfs_archives function""" + + @pytest.fixture + def temp_gtfs_dir(self, tmp_path, monkeypatch): + """Point the GTFS archive dir at a temp directory with a known retention.""" + monkeypatch.setattr(gtfs, "MAIN_DIR", tmp_path) + monkeypatch.setattr(gtfs, "GTFS_ARCHIVE_RETENTION_DAYS", 90) + return tmp_path + + def _make_archive(self, base: pathlib.Path, name: str, age_days: int) -> pathlib.Path: + """Create a fake archive directory whose mtime is age_days in the past.""" + archive_dir = base / name + archive_dir.mkdir() + (archive_dir / "trips.txt").write_text("trip_id\n") + mtime = (datetime.datetime.now() - datetime.timedelta(days=age_days)).timestamp() + os.utime(archive_dir, (mtime, mtime)) + return archive_dir + + def test_deletes_archives_older_than_retention(self, temp_gtfs_dir): + self._make_archive(temp_gtfs_dir, "20200101", age_days=200) + self._make_archive(temp_gtfs_dir, "20260601", age_days=1) + + deleted = gtfs.cleanup_old_gtfs_archives() + + assert deleted == 1 + assert not (temp_gtfs_dir / "20200101").exists() + assert (temp_gtfs_dir / "20260601").exists() + + def test_keeps_archives_within_retention(self, temp_gtfs_dir): + self._make_archive(temp_gtfs_dir, "20260501", age_days=30) + self._make_archive(temp_gtfs_dir, "20260601", age_days=1) + + deleted = gtfs.cleanup_old_gtfs_archives() + + assert deleted == 0 + assert len([d for d in temp_gtfs_dir.iterdir() if d.is_dir()]) == 2 + + def test_always_keeps_most_recent_archive(self, temp_gtfs_dir): + # Only one archive and it is very old—it must still be retained as the active feed. + self._make_archive(temp_gtfs_dir, "20190101", age_days=500) + + deleted = gtfs.cleanup_old_gtfs_archives() + + assert deleted == 0 + assert (temp_gtfs_dir / "20190101").exists() + + def test_ignores_archived_feeds_file(self, temp_gtfs_dir): + (temp_gtfs_dir / gtfs.GTFS_ARCHIVES_FILENAME).write_text("data") + self._make_archive(temp_gtfs_dir, "20200101", age_days=200) + self._make_archive(temp_gtfs_dir, "20260601", age_days=1) + + deleted = gtfs.cleanup_old_gtfs_archives() + + assert deleted == 1 + assert (temp_gtfs_dir / gtfs.GTFS_ARCHIVES_FILENAME).exists()