Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion config/template.json
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
46 changes: 46 additions & 0 deletions src/gtfs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down Expand Up @@ -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):
"""
Expand Down
4 changes: 3 additions & 1 deletion src/s3_upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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")
Expand Down
59 changes: 59 additions & 0 deletions src/tests/test_gtfs.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import datetime
import os
import numpy as np
import pandas as pd
import pathlib
Expand Down Expand Up @@ -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()
Loading