Add routine to automatically cleanup gtfs directories (#139)

Author: ankoure

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,

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):
     """

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")

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()