diff --git a/functions-python/batch_process_dataset/src/pipeline_tasks.py b/functions-python/batch_process_dataset/src/pipeline_tasks.py index 17eaced47..146fa719d 100644 --- a/functions-python/batch_process_dataset/src/pipeline_tasks.py +++ b/functions-python/batch_process_dataset/src/pipeline_tasks.py @@ -14,6 +14,10 @@ create_http_gtfs_datasets_comparer_task, ) +# GTFS files that have a registered extractor in the gtfs_file_data_extractor +# function. Keep in sync with that function's src/extractors/registry.py. +EXTRACTABLE_FILES = {"feed_info.txt"} + def create_http_reverse_geolocation_processor_task( stable_id: str, @@ -45,6 +49,39 @@ def create_http_reverse_geolocation_processor_task( ) +def create_http_gtfs_file_data_extractor_task( + stable_id: str, + dataset_stable_id: str, + file_name: str, + file_url: str, +) -> None: + """ + Create a task to extract structured data from a single GTFS file + (handled by the gtfs_file_data_extractor function). + """ + client = tasks_v2.CloudTasksClient() + body = json.dumps( + { + "stable_id": stable_id, + "dataset_id": dataset_stable_id, + "file_name": file_name, + "file_url": file_url, + } + ).encode() + queue_name = os.getenv("GTFS_FILE_DATA_EXTRACTOR_QUEUE") + project_id = os.getenv("PROJECT_ID") + gcp_region = os.getenv("GCP_REGION") + + create_http_task( + client, + body, + f"https://{gcp_region}-{project_id}.cloudfunctions.net/gtfs-file-data-extractor", + project_id, + gcp_region, + queue_name, + ) + + @with_db_session def get_changed_files( dataset: Gtfsdataset, @@ -109,6 +146,15 @@ def create_pipeline_tasks(dataset: Gtfsdataset, db_session: Session) -> None: stable_id, dataset_stable_id, stops_url ) + # Create GTFS file data extraction tasks for changed, extractable files. + files_by_name = {file.file_name: file for file in gtfs_files} + for file_name in EXTRACTABLE_FILES: + gtfs_file = files_by_name.get(file_name) + if gtfs_file and gtfs_file.hosted_url and file_name in changed_files: + create_http_gtfs_file_data_extractor_task( + stable_id, dataset_stable_id, file_name, gtfs_file.hosted_url + ) + routes_file = next( (file for file in gtfs_files if file.file_name == "routes.txt"), None ) diff --git a/functions-python/batch_process_dataset/tests/test_pipeline_tasks.py b/functions-python/batch_process_dataset/tests/test_pipeline_tasks.py index 48d91aa34..d5dab104d 100644 --- a/functions-python/batch_process_dataset/tests/test_pipeline_tasks.py +++ b/functions-python/batch_process_dataset/tests/test_pipeline_tasks.py @@ -5,6 +5,7 @@ from pipeline_tasks import ( create_http_reverse_geolocation_processor_task, + create_http_gtfs_file_data_extractor_task, get_changed_files, create_pipeline_tasks, ) @@ -85,6 +86,113 @@ def test_create_http_reverse_geolocation_processor_task( self.assertEqual(args[4], "northamerica-northeast1") self.assertEqual(args[5], "rev-geo-queue") + @patch.dict( + os.environ, + { + "GTFS_FILE_DATA_EXTRACTOR_QUEUE": "file-data-queue", + "PROJECT_ID": "my-project", + "GCP_REGION": "northamerica-northeast1", + }, + clear=False, + ) + @patch("pipeline_tasks.create_http_task") + @patch("pipeline_tasks.tasks_v2.CloudTasksClient") + def test_create_http_gtfs_file_data_extractor_task( + self, mock_client_cls, mock_create_http_task + ): + client_instance = MagicMock() + mock_client_cls.return_value = client_instance + + create_http_gtfs_file_data_extractor_task( + stable_id="feed-123", + dataset_stable_id="dataset-abc", + file_name="feed_info.txt", + file_url="https://example.com/feed_info.txt", + ) + + self.assertEqual(mock_create_http_task.call_count, 1) + args, _ = mock_create_http_task.call_args + self.assertIs(args[0], client_instance) + payload = json.loads(args[1].decode("utf-8")) + self.assertEqual( + payload, + { + "stable_id": "feed-123", + "dataset_id": "dataset-abc", + "file_name": "feed_info.txt", + "file_url": "https://example.com/feed_info.txt", + }, + ) + self.assertEqual( + args[2], + "https://northamerica-northeast1-my-project.cloudfunctions.net/gtfs-file-data-extractor", + ) + self.assertEqual(args[5], "file-data-queue") + + +class TestCreatePipelineTasksFeedInfo(unittest.TestCase): + """Covers the gtfs_file_data_extractor enqueue gating in create_pipeline_tasks.""" + + def _mock_session_no_base_dataset(self): + mock_session = MagicMock() + mock_session.query.return_value.filter.return_value.order_by.return_value.first.return_value = ( + None + ) + return mock_session + + def _run(self, files, changed_files): + dataset = SimpleDataset( + feed_id=1, + dataset_id=10, + feed_stable_id="feed-A", + dataset_stable_id="ds-1", + files=files, + ) + with patch( + "pipeline_tasks.get_changed_files", return_value=changed_files + ), patch( + "pipeline_tasks.create_http_reverse_geolocation_processor_task" + ), patch( + "pipeline_tasks.create_http_pmtiles_builder_task" + ), patch( + "pipeline_tasks.create_http_gtfs_datasets_comparer_task" + ), patch( + "pipeline_tasks.create_http_gtfs_file_data_extractor_task" + ) as mock_extractor_task: + create_pipeline_tasks( + dataset, db_session=self._mock_session_no_base_dataset() + ) + return mock_extractor_task + + def test_enqueues_when_feed_info_present_and_changed(self): + files = [ + SimpleFile( + "feed_info.txt", + hosted_url="https://x.com/feed_info.txt", + file_hash="h1", + ) + ] + mock_task = self._run(files, changed_files=["feed_info.txt"]) + mock_task.assert_called_once_with( + "feed-A", "ds-1", "feed_info.txt", "https://x.com/feed_info.txt" + ) + + def test_skips_when_feed_info_not_changed(self): + files = [ + SimpleFile( + "feed_info.txt", + hosted_url="https://x.com/feed_info.txt", + file_hash="h1", + ) + ] + mock_task = self._run(files, changed_files=["stops.txt"]) + mock_task.assert_not_called() + + def test_skips_when_feed_info_absent(self): + files = [SimpleFile("stops.txt", hosted_url="https://x.com/stops.txt")] + mock_task = self._run(files, changed_files=["stops.txt"]) + mock_task.assert_not_called() + class TestHasFileChanged(unittest.TestCase): def _make_mock_session_chain(self, previous_dataset): diff --git a/functions-python/gtfs_file_data_extractor/README.md b/functions-python/gtfs_file_data_extractor/README.md new file mode 100644 index 000000000..1050d9034 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/README.md @@ -0,0 +1,32 @@ +# GTFS File Data Extractor + +HTTP Cloud Function that extracts structured data from a single GTFS file and +persists it to the database. It is enqueued as a Cloud Task by +`batch_process_dataset` once a dataset has been processed (see +`batch_process_dataset/src/pipeline_tasks.py`). + +## How it works + +The task payload identifies one GTFS file: + +```json +{ + "stable_id": "", + "dataset_id": "", + "file_name": "feed_info.txt", + "file_url": "" +} +``` + +`processor.py` downloads the file, looks up the matching extractor in +`extractors/registry.py`, and lets it write to the database. + +## Adding a new file extractor + +1. Implement a `FileDataExtractor` subclass under `src/extractors/`. +2. Register it in `src/extractors/registry.py`. +3. Add its `file_name` to `EXTRACTABLE_FILES` in + `batch_process_dataset/src/pipeline_tasks.py` so the producer enqueues it. + +Currently registered: `feed_info.txt` -> `FeedInfoExtractor` (writes the +`feedinfo` table). diff --git a/functions-python/gtfs_file_data_extractor/function_config.json b/functions-python/gtfs_file_data_extractor/function_config.json new file mode 100644 index 000000000..680212add --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/function_config.json @@ -0,0 +1,21 @@ +{ + "name": "gtfs-file-data-extractor", + "description": "Extracts structured data from GTFS files (e.g. feed_info.txt) into the database", + "entry_point": "gtfs_file_data_extractor", + "timeout": 540, + "available_memory": "512Mi", + "trigger_http": true, + "include_folders": ["helpers"], + "include_api_folders": ["database_gen", "database", "common"], + "environment_variables": [], + "secret_environment_variables": [ + { + "key": "FEEDS_DATABASE_URL" + } + ], + "ingress_settings": "ALLOW_ALL", + "max_instance_request_concurrency": 1, + "max_instance_count": 5, + "min_instance_count": 0, + "available_cpu": 1 +} diff --git a/functions-python/gtfs_file_data_extractor/requirements.txt b/functions-python/gtfs_file_data_extractor/requirements.txt new file mode 100644 index 000000000..420bfad44 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/requirements.txt @@ -0,0 +1,20 @@ +# Common packages +functions-framework==3.* +google-cloud-logging +psycopg2-binary==2.9.6 +requests~=2.33.1 +certifi~=2025.8.3 + +# SQL Alchemy and Geo Alchemy (database_gen models depend on geoalchemy2) +SQLAlchemy==2.0.23 +geoalchemy2==0.14.7 + +# Google specific packages (shared helpers) +google-cloud-tasks +google-auth==2.29.0 + +# Additional packages for this function +pandas + +# Configuration +python-dotenv==1.2.2 diff --git a/functions-python/gtfs_file_data_extractor/requirements_dev.txt b/functions-python/gtfs_file_data_extractor/requirements_dev.txt new file mode 100644 index 000000000..e7ff72aea --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/requirements_dev.txt @@ -0,0 +1,3 @@ +Faker +pytest~=7.4.3 +requests-mock diff --git a/functions-python/gtfs_file_data_extractor/src/extractors/__init__.py b/functions-python/gtfs_file_data_extractor/src/extractors/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/functions-python/gtfs_file_data_extractor/src/extractors/base.py b/functions-python/gtfs_file_data_extractor/src/extractors/base.py new file mode 100644 index 000000000..32925d113 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/src/extractors/base.py @@ -0,0 +1,34 @@ +from abc import ABC, abstractmethod + +import pandas as pd +from sqlalchemy.orm import Session + +from shared.database_gen.sqlacodegen_models import Gtfsdataset + + +class FileDataExtractor(ABC): + """ + Base class for extracting structured data from a single GTFS file. + + Each subclass handles exactly one GTFS file (identified by ``file_name``) + and persists the extracted values to the database. + + To add support for a new file: + 1. Implement a subclass here and register it in ``extractors/registry.py``. + 2. Add the file name to ``EXTRACTABLE_FILES`` in batch_process_dataset's + ``pipeline_tasks.py`` so the producer enqueues a task for it. + """ + + #: The GTFS file this extractor handles, e.g. "feed_info.txt". + file_name: str + + @abstractmethod + def extract( + self, df: pd.DataFrame, dataset: Gtfsdataset, db_session: Session + ) -> None: + """ + Parse ``df`` (the parsed contents of ``file_name``) and persist the + extracted data for ``dataset``. Implementations must be idempotent: + re-running for the same dataset should update, not duplicate. + """ + raise NotImplementedError diff --git a/functions-python/gtfs_file_data_extractor/src/extractors/feed_info_extractor.py b/functions-python/gtfs_file_data_extractor/src/extractors/feed_info_extractor.py new file mode 100644 index 000000000..44244fe58 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/src/extractors/feed_info_extractor.py @@ -0,0 +1,88 @@ +import logging +from datetime import date, datetime +from typing import Optional + +import pandas as pd +from sqlalchemy.orm import Session + +from extractors.base import FileDataExtractor +from shared.database_gen.sqlacodegen_models import Feedinfo, Gtfsdataset + +# feed_info.txt text columns (GTFS spec) mapped 1:1 onto Feedinfo string columns. +STRING_FIELDS = ( + "feed_publisher_name", + "feed_publisher_url", + "feed_lang", + "default_lang", + "feed_version", + "feed_contact_email", + "feed_contact_url", +) +# feed_info.txt date columns, stored as plain DATE (YYYYMMDD, no timezone). +DATE_FIELDS = ("feed_start_date", "feed_end_date") + + +def clean_str(value) -> Optional[str]: + """Return a trimmed string, or None for empty/NaN/missing values.""" + if value is None or (isinstance(value, float) and pd.isna(value)): + return None + text = str(value).strip() + return text or None + + +def parse_gtfs_date(value) -> Optional[date]: + """Parse a GTFS YYYYMMDD date. Returns None when missing or unparseable.""" + text = clean_str(value) + if text is None: + return None + # pandas may read a purely numeric column as int/float, yielding "20240101.0". + if text.endswith(".0"): + text = text[:-2] + try: + return datetime.strptime(text, "%Y%m%d").date() + except ValueError: + logging.warning("Unparseable feed_info date value: %r", value) + return None + + +class FeedInfoExtractor(FileDataExtractor): + """Extracts feed_info.txt into a Feedinfo row (one per dataset).""" + + file_name = "feed_info.txt" + + def extract( + self, df: pd.DataFrame, dataset: Gtfsdataset, db_session: Session + ) -> None: + if df is None or df.empty: + logging.info( + "feed_info.txt is empty for dataset %s; nothing to extract.", + dataset.stable_id, + ) + return + + # feed_info.txt holds a single record. + row = df.iloc[0] + + values = {field: clean_str(row.get(field)) for field in STRING_FIELDS} + for field in DATE_FIELDS: + values[field] = parse_gtfs_date(row.get(field)) + + # Upsert keyed by dataset so reprocessing updates in place. + feed_info = ( + db_session.query(Feedinfo) + .filter(Feedinfo.gtfs_dataset_id == dataset.id) + .one_or_none() + ) + if feed_info is None: + feed_info = Feedinfo(gtfs_dataset_id=dataset.id) + db_session.add(feed_info) + + for field, value in values.items(): + setattr(feed_info, field, value) + + logging.info( + "Extracted feed_info for dataset %s: start=%s end=%s", + dataset.stable_id, + values["feed_start_date"], + values["feed_end_date"], + ) diff --git a/functions-python/gtfs_file_data_extractor/src/extractors/registry.py b/functions-python/gtfs_file_data_extractor/src/extractors/registry.py new file mode 100644 index 000000000..271306b09 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/src/extractors/registry.py @@ -0,0 +1,16 @@ +from typing import Dict, Optional + +from extractors.base import FileDataExtractor +from extractors.feed_info_extractor import FeedInfoExtractor + +# Registry of GTFS file name -> extractor. Add new extractors here; the producer +# side (batch_process_dataset/pipeline_tasks.py EXTRACTABLE_FILES) decides which +# files actually get enqueued. +EXTRACTORS: Dict[str, FileDataExtractor] = { + extractor.file_name: extractor for extractor in (FeedInfoExtractor(),) +} + + +def get_extractor(file_name: str) -> Optional[FileDataExtractor]: + """Return the extractor registered for ``file_name``, or None if unsupported.""" + return EXTRACTORS.get(file_name) diff --git a/functions-python/gtfs_file_data_extractor/src/main.py b/functions-python/gtfs_file_data_extractor/src/main.py new file mode 100644 index 000000000..f2b6a9d9e --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/src/main.py @@ -0,0 +1,24 @@ +import logging + +import flask +import functions_framework + +from shared.helpers.logger import init_logger + +init_logger() + + +@functions_framework.http +def gtfs_file_data_extractor(request: flask.Request): + """ + Cloud Function that extracts structured data from a single GTFS file + (e.g. feed_info.txt) and persists it to the database. + + Function trigger: HTTP request by Cloud Tasks, enqueued by + batch_process_dataset once a dataset has been processed. + """ + from processor import process_file_data + + result = process_file_data(request) + logging.info(result) + return result diff --git a/functions-python/gtfs_file_data_extractor/src/processor.py b/functions-python/gtfs_file_data_extractor/src/processor.py new file mode 100644 index 000000000..13ed4ab35 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/src/processor.py @@ -0,0 +1,101 @@ +import io +import logging +from typing import Tuple + +import flask +import pandas as pd +import requests +from sqlalchemy.orm import Session + +from extractors.registry import get_extractor +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Gtfsdataset + +ERROR_STATUS_CODE = 500 +REQUIRED_PARAMETERS = ("stable_id", "dataset_id", "file_name", "file_url") + + +def parse_request_parameters(request: flask.Request) -> Tuple[str, str, str, str]: + """ + Parse and validate the Cloud Task payload. + + Expected JSON body: + { + "stable_id": "", + "dataset_id": "", + "file_name": "feed_info.txt", + "file_url": "" + } + """ + request_json = request.get_json(silent=True) + logging.info("Request JSON: %s", request_json) + if not request_json: + raise ValueError("Missing or invalid JSON body.") + missing = [key for key in REQUIRED_PARAMETERS if not request_json.get(key)] + if missing: + raise ValueError(f"Missing required parameters: {missing}.") + return ( + request_json["stable_id"], + request_json["dataset_id"], + request_json["file_name"], + request_json["file_url"], + ) + + +def load_dataset(dataset_id: str, db_session: Session) -> Gtfsdataset: + dataset = ( + db_session.query(Gtfsdataset) + .filter(Gtfsdataset.stable_id == dataset_id) + .one_or_none() + ) + if not dataset: + raise ValueError( + f"Dataset with ID {dataset_id} does not exist in the database." + ) + return dataset + + +def read_csv_from_url(file_url: str) -> pd.DataFrame: + response = requests.get(file_url) + response.raise_for_status() + return pd.read_csv(io.StringIO(response.content.decode("utf-8"))) + + +@with_db_session +def process_file_data( + request: flask.Request, db_session: Session = None +) -> Tuple[str, int]: + """ + Download a single GTFS file and dispatch it to the registered extractor, + which persists the extracted data to the database. + """ + try: + stable_id, dataset_id, file_name, file_url = parse_request_parameters(request) + except ValueError as error: + logging.error("Invalid request: %s", error) + return str(error), 400 + + extractor = get_extractor(file_name) + if extractor is None: + message = f"No extractor registered for file '{file_name}'. Skipping." + logging.info(message) + return message, 200 + + try: + dataset = load_dataset(dataset_id, db_session) + df = read_csv_from_url(file_url) + extractor.extract(df, dataset, db_session) + db_session.commit() + except Exception as error: + db_session.rollback() + logging.error( + "Error extracting data from %s for dataset %s: %s", + file_name, + dataset_id, + error, + ) + return f"Error extracting data from {file_name}: {error}", ERROR_STATUS_CODE + + message = f"Successfully extracted data from {file_name} for dataset {dataset_id}." + logging.info(message) + return message, 200 diff --git a/functions-python/gtfs_file_data_extractor/tests/test_feed_info_extractor.py b/functions-python/gtfs_file_data_extractor/tests/test_feed_info_extractor.py new file mode 100644 index 000000000..c53cf6189 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/tests/test_feed_info_extractor.py @@ -0,0 +1,157 @@ +import datetime +import unittest + +import pandas as pd +from faker import Faker + +from extractors.feed_info_extractor import ( + FeedInfoExtractor, + clean_str, + parse_gtfs_date, +) +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Feedinfo, Gtfsdataset, Gtfsfeed +from test_shared.test_utils.database_utils import clean_testing_db, default_db_url + +faker = Faker() + + +class TestParsingHelpers(unittest.TestCase): + def test_clean_str(self): + self.assertIsNone(clean_str(None)) + self.assertIsNone(clean_str(float("nan"))) + self.assertIsNone(clean_str(" ")) + self.assertEqual(clean_str(" Agency "), "Agency") + self.assertEqual(clean_str(123), "123") + + def test_parse_gtfs_date_valid(self): + self.assertEqual(parse_gtfs_date("20240115"), datetime.date(2024, 1, 15)) + + def test_parse_gtfs_date_numeric(self): + # pandas may read a purely numeric column as int or float + self.assertEqual(parse_gtfs_date("20240115.0"), datetime.date(2024, 1, 15)) + self.assertEqual(parse_gtfs_date(20240115), datetime.date(2024, 1, 15)) + + def test_parse_gtfs_date_missing_or_invalid(self): + self.assertIsNone(parse_gtfs_date(None)) + self.assertIsNone(parse_gtfs_date("")) + self.assertIsNone(parse_gtfs_date("not-a-date")) + + +class TestFeedInfoExtractor(unittest.TestCase): + def setUp(self): + clean_testing_db() + + def _create_dataset(self, db_session): + feed = Gtfsfeed( + id=faker.uuid4(cast_to=str), + data_type="gtfs", + stable_id=faker.uuid4(cast_to=str), + ) + dataset_id = faker.uuid4(cast_to=str) + dataset = Gtfsdataset(id=dataset_id, stable_id=dataset_id, feed=feed) + db_session.add(dataset) + db_session.commit() + return dataset + + @with_db_session(db_url=default_db_url) + def test_extract_inserts_feedinfo(self, db_session): + dataset = self._create_dataset(db_session) + df = pd.DataFrame( + [ + { + "feed_publisher_name": "Test Agency", + "feed_publisher_url": "https://example.com", + "feed_lang": "en", + "default_lang": "fr", + "feed_start_date": "20240101", + "feed_end_date": "20241231", + "feed_version": "v1", + "feed_contact_email": "info@example.com", + "feed_contact_url": "https://example.com/contact", + } + ] + ) + FeedInfoExtractor().extract(df, dataset, db_session) + db_session.commit() + + feed_info = ( + db_session.query(Feedinfo) + .filter(Feedinfo.gtfs_dataset_id == dataset.id) + .one() + ) + self.assertEqual(feed_info.feed_publisher_name, "Test Agency") + self.assertEqual(feed_info.feed_publisher_url, "https://example.com") + self.assertEqual(feed_info.feed_lang, "en") + self.assertEqual(feed_info.default_lang, "fr") + self.assertEqual(feed_info.feed_start_date, datetime.date(2024, 1, 1)) + self.assertEqual(feed_info.feed_end_date, datetime.date(2024, 12, 31)) + self.assertEqual(feed_info.feed_version, "v1") + self.assertEqual(feed_info.feed_contact_email, "info@example.com") + self.assertEqual(feed_info.feed_contact_url, "https://example.com/contact") + + @with_db_session(db_url=default_db_url) + def test_extract_upserts_existing(self, db_session): + dataset = self._create_dataset(db_session) + FeedInfoExtractor().extract( + pd.DataFrame( + [{"feed_publisher_name": "First", "feed_start_date": "20240101"}] + ), + dataset, + db_session, + ) + db_session.commit() + + FeedInfoExtractor().extract( + pd.DataFrame( + [{"feed_publisher_name": "Second", "feed_start_date": "20250101"}] + ), + dataset, + db_session, + ) + db_session.commit() + + rows = ( + db_session.query(Feedinfo) + .filter(Feedinfo.gtfs_dataset_id == dataset.id) + .all() + ) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0].feed_publisher_name, "Second") + self.assertEqual(rows[0].feed_start_date, datetime.date(2025, 1, 1)) + + @with_db_session(db_url=default_db_url) + def test_extract_missing_optional_columns(self, db_session): + dataset = self._create_dataset(db_session) + FeedInfoExtractor().extract( + pd.DataFrame([{"feed_publisher_name": "Only Name"}]), + dataset, + db_session, + ) + db_session.commit() + + feed_info = ( + db_session.query(Feedinfo) + .filter(Feedinfo.gtfs_dataset_id == dataset.id) + .one() + ) + self.assertEqual(feed_info.feed_publisher_name, "Only Name") + self.assertIsNone(feed_info.feed_start_date) + self.assertIsNone(feed_info.feed_lang) + + @with_db_session(db_url=default_db_url) + def test_extract_empty_dataframe_is_noop(self, db_session): + dataset = self._create_dataset(db_session) + FeedInfoExtractor().extract(pd.DataFrame(), dataset, db_session) + db_session.commit() + + count = ( + db_session.query(Feedinfo) + .filter(Feedinfo.gtfs_dataset_id == dataset.id) + .count() + ) + self.assertEqual(count, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/gtfs_file_data_extractor/tests/test_processor.py b/functions-python/gtfs_file_data_extractor/tests/test_processor.py new file mode 100644 index 000000000..93ec567d6 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/tests/test_processor.py @@ -0,0 +1,137 @@ +import datetime +import unittest +from unittest.mock import MagicMock, patch + +from faker import Faker + +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Feedinfo, Gtfsdataset, Gtfsfeed +from test_shared.test_utils.database_utils import clean_testing_db, default_db_url + +faker = Faker() + + +def make_request(payload): + request = MagicMock() + request.get_json.return_value = payload + return request + + +class TestParseRequestParameters(unittest.TestCase): + def test_valid(self): + from processor import parse_request_parameters + + payload = { + "stable_id": "feed", + "dataset_id": "dataset", + "file_name": "feed_info.txt", + "file_url": "http://example.com/feed_info.txt", + } + self.assertEqual( + parse_request_parameters(make_request(payload)), + ("feed", "dataset", "feed_info.txt", "http://example.com/feed_info.txt"), + ) + + def test_missing_parameters(self): + from processor import parse_request_parameters + + with self.assertRaises(ValueError): + parse_request_parameters(make_request({"stable_id": "feed"})) + + def test_no_body(self): + from processor import parse_request_parameters + + with self.assertRaises(ValueError): + parse_request_parameters(make_request(None)) + + +class TestProcessFileData(unittest.TestCase): + def setUp(self): + clean_testing_db() + + def _create_dataset(self, db_session): + feed = Gtfsfeed( + id=faker.uuid4(cast_to=str), + data_type="gtfs", + stable_id=faker.uuid4(cast_to=str), + ) + dataset_id = faker.uuid4(cast_to=str) + dataset = Gtfsdataset(id=dataset_id, stable_id=dataset_id, feed=feed) + db_session.add(dataset) + db_session.commit() + return dataset_id + + @with_db_session(db_url=default_db_url) + @patch("processor.requests") + def test_process_success_writes_feedinfo(self, requests_mock, db_session): + dataset_id = self._create_dataset(db_session) + csv = ( + "feed_publisher_name,feed_start_date,feed_end_date\n" + "Agency,20240101,20241231\n" + ) + requests_mock.get.return_value.content = csv.encode("utf-8") + + from processor import process_file_data + + payload = { + "stable_id": "feed", + "dataset_id": dataset_id, + "file_name": "feed_info.txt", + "file_url": "http://example.com/feed_info.txt", + } + _, status = process_file_data(make_request(payload), db_session=db_session) + self.assertEqual(status, 200) + + feed_info = ( + db_session.query(Feedinfo) + .filter(Feedinfo.gtfs_dataset_id == dataset_id) + .one() + ) + self.assertEqual(feed_info.feed_publisher_name, "Agency") + self.assertEqual(feed_info.feed_start_date, datetime.date(2024, 1, 1)) + self.assertEqual(feed_info.feed_end_date, datetime.date(2024, 12, 31)) + + @with_db_session(db_url=default_db_url) + @patch("processor.requests") + def test_process_unregistered_file_skips(self, requests_mock, db_session): + from processor import process_file_data + + payload = { + "stable_id": "feed", + "dataset_id": "dataset", + "file_name": "stops.txt", + "file_url": "http://example.com/stops.txt", + } + message, status = process_file_data( + make_request(payload), db_session=db_session + ) + self.assertEqual(status, 200) + self.assertIn("No extractor", message) + requests_mock.get.assert_not_called() + + @with_db_session(db_url=default_db_url) + def test_process_invalid_payload_returns_400(self, db_session): + from processor import process_file_data + + _, status = process_file_data(make_request({}), db_session=db_session) + self.assertEqual(status, 400) + + @with_db_session(db_url=default_db_url) + @patch("processor.requests") + def test_process_missing_dataset_returns_500(self, requests_mock, db_session): + requests_mock.get.return_value.content = b"feed_publisher_name\nAgency\n" + + from processor import process_file_data + + payload = { + "stable_id": "feed", + "dataset_id": "does-not-exist", + "file_name": "feed_info.txt", + "file_url": "http://example.com/feed_info.txt", + } + _, status = process_file_data(make_request(payload), db_session=db_session) + self.assertEqual(status, 500) + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/gtfs_file_data_extractor/tests/test_registry.py b/functions-python/gtfs_file_data_extractor/tests/test_registry.py new file mode 100644 index 000000000..308fb8a83 --- /dev/null +++ b/functions-python/gtfs_file_data_extractor/tests/test_registry.py @@ -0,0 +1,20 @@ +import unittest + +from extractors.feed_info_extractor import FeedInfoExtractor +from extractors.registry import EXTRACTORS, get_extractor + + +class TestRegistry(unittest.TestCase): + def test_known_file_returns_extractor(self): + self.assertIsInstance(get_extractor("feed_info.txt"), FeedInfoExtractor) + + def test_unknown_file_returns_none(self): + self.assertIsNone(get_extractor("stops.txt")) + + def test_registry_keys_match_extractor_file_name(self): + for file_name, extractor in EXTRACTORS.items(): + self.assertEqual(file_name, extractor.file_name) + + +if __name__ == "__main__": + unittest.main() diff --git a/infra/batch/main.tf b/infra/batch/main.tf index 3a1819cf3..ca671d18a 100644 --- a/infra/batch/main.tf +++ b/infra/batch/main.tf @@ -294,6 +294,25 @@ resource "google_cloud_tasks_queue" "gtfs_datasets_comparer_task_queue" { } } +# Task queue to invoke gtfs_file_data_extractor function +# The function itself is deployed by the infra/functions-python stack. +resource "google_cloud_tasks_queue" "gtfs_file_data_extractor_task_queue" { + project = var.project_id + location = var.gcp_region + name = "gtfs-file-data-extractor-queue-${var.environment}-${local.deployment_timestamp}" + + rate_limits { + max_concurrent_dispatches = 10 + max_dispatches_per_second = 1 + } + + retry_config { + max_attempts = 10 + min_backoff = "20s" + max_backoff = "60s" + } +} + # Batch process dataset function resource "google_cloudfunctions2_function" "pubsub_function" { @@ -332,6 +351,7 @@ resource "google_cloudfunctions2_function" "pubsub_function" { MATERIALIZED_VIEW_QUEUE = google_cloud_tasks_queue.refresh_materialized_view_task_queue.name PMTILES_BUILDER_QUEUE = google_cloud_tasks_queue.pmtiles_builder_task_queue.name REVERSE_GEOLOCATION_QUEUE = "reverse-geolocation-processor-task-queue" + GTFS_FILE_DATA_EXTRACTOR_QUEUE = google_cloud_tasks_queue.gtfs_file_data_extractor_task_queue.name GTFS_CHANGE_TRACKER_QUEUE = google_cloud_tasks_queue.gtfs_datasets_comparer_task_queue.name WEB_REVALIDATION_QUEUE = google_cloud_tasks_queue.web_revalidation_task_queue.name } diff --git a/infra/functions-python/main.tf b/infra/functions-python/main.tf index b3f72c4cf..3f7fa008b 100644 --- a/infra/functions-python/main.tf +++ b/infra/functions-python/main.tf @@ -68,6 +68,9 @@ locals { function_gtfs_datasets_comparer_config = jsondecode(file("${path.module}/../../functions-python/gtfs_datasets_comparer/function_config.json")) function_gtfs_datasets_comparer_zip = "${path.module}/../../functions-python/gtfs_datasets_comparer/.dist/gtfs_datasets_comparer.zip" + + function_gtfs_file_data_extractor_config = jsondecode(file("${path.module}/../../functions-python/gtfs_file_data_extractor/function_config.json")) + function_gtfs_file_data_extractor_zip = "${path.module}/../../functions-python/gtfs_file_data_extractor/.dist/gtfs_file_data_extractor.zip" } locals { @@ -80,7 +83,8 @@ locals { local.function_export_csv_config.secret_environment_variables, local.function_tasks_executor_config.secret_environment_variables, local.function_pmtiles_builder_config.secret_environment_variables, - local.function_gtfs_datasets_comparer_config.secret_environment_variables + local.function_gtfs_datasets_comparer_config.secret_environment_variables, + local.function_gtfs_file_data_extractor_config.secret_environment_variables ) # Remove duplicates by key, keeping the first occurrence @@ -233,6 +237,13 @@ resource "google_storage_bucket_object" "gtfs_datasets_comparer_zip" { source = local.function_gtfs_datasets_comparer_zip } +# 17. GTFS File Data Extractor +resource "google_storage_bucket_object" "gtfs_file_data_extractor_zip" { + bucket = google_storage_bucket.functions_bucket.name + name = "gtfs-file-data-extractor-${substr(filebase64sha256(local.function_gtfs_file_data_extractor_zip), 0, 10)}.zip" + source = local.function_gtfs_file_data_extractor_zip +} + # Web app revalidation secret resource "google_secret_manager_secret" "web_app_revalidate_secret" { project = var.project_id @@ -1131,6 +1142,72 @@ resource "google_cloudfunctions2_function_iam_member" "gtfs_datasets_comparer_in member = "serviceAccount:${google_service_account.functions_service_account.email}" } +# 17.1 functions/gtfs_file_data_extractor cloud function +resource "google_cloudfunctions2_function" "gtfs_file_data_extractor" { + name = local.function_gtfs_file_data_extractor_config.name + description = local.function_gtfs_file_data_extractor_config.description + location = var.gcp_region + depends_on = [ + google_project_iam_member.event-receiving, + google_secret_manager_secret_iam_member.secret_iam_member, + ] + + build_config { + runtime = var.python_runtime + entry_point = local.function_gtfs_file_data_extractor_config.entry_point + source { + storage_source { + bucket = google_storage_bucket.functions_bucket.name + object = google_storage_bucket_object.gtfs_file_data_extractor_zip.name + } + } + } + service_config { + environment_variables = { + PYTHONNODEBUGRANGES = 0 + ENVIRONMENT = var.environment + PROJECT_ID = var.project_id + GCP_REGION = var.gcp_region + } + available_memory = local.function_gtfs_file_data_extractor_config.available_memory + timeout_seconds = local.function_gtfs_file_data_extractor_config.timeout + available_cpu = local.function_gtfs_file_data_extractor_config.available_cpu + max_instance_request_concurrency = local.function_gtfs_file_data_extractor_config.max_instance_request_concurrency + max_instance_count = local.function_gtfs_file_data_extractor_config.max_instance_count + min_instance_count = local.function_gtfs_file_data_extractor_config.min_instance_count + service_account_email = google_service_account.functions_service_account.email + ingress_settings = local.function_gtfs_file_data_extractor_config.ingress_settings + vpc_connector = data.google_vpc_access_connector.vpc_connector.id + vpc_connector_egress_settings = "PRIVATE_RANGES_ONLY" + dynamic "secret_environment_variables" { + for_each = local.function_gtfs_file_data_extractor_config.secret_environment_variables + content { + key = secret_environment_variables.value["key"] + project_id = var.project_id + secret = "${upper(var.environment)}_${secret_environment_variables.value["key"]}" + version = "latest" + } + } + } +} + +# Grant execution permission to batchfunctions service account to the gtfs_file_data_extractor function +resource "google_cloudfunctions2_function_iam_member" "gtfs_file_data_extractor_invoker_batch_sa" { + project = var.project_id + location = var.gcp_region + cloud_function = google_cloudfunctions2_function.gtfs_file_data_extractor.name + role = "roles/cloudfunctions.invoker" + member = "serviceAccount:${local.batchfunctions_sa_email}" +} + +resource "google_cloud_run_service_iam_member" "gtfs_file_data_extractor_cloud_run_invoker" { + project = var.project_id + location = var.gcp_region + service = google_cloudfunctions2_function.gtfs_file_data_extractor.name + role = "roles/run.invoker" + member = "serviceAccount:${local.batchfunctions_sa_email}" +} + # 13.3 functions/reverse_geolocation - batch cloud function resource "google_cloudfunctions2_function" "reverse_geolocation_batch" { name = "${local.function_reverse_geolocation_config.name}-batch" diff --git a/liquibase/changelog.xml b/liquibase/changelog.xml index 1408c7220..8ca714f60 100644 --- a/liquibase/changelog.xml +++ b/liquibase/changelog.xml @@ -127,6 +127,8 @@ + + diff --git a/liquibase/changes/feat_1775.sql b/liquibase/changes/feat_1775.sql new file mode 100644 index 000000000..1187ae34d --- /dev/null +++ b/liquibase/changes/feat_1775.sql @@ -0,0 +1,26 @@ +-- Store parsed feed_info.txt fields per GTFS dataset (issue #1775). +-- One row per dataset; populated by the gtfs_file_data_extractor cloud function. +CREATE TABLE IF NOT EXISTS feedinfo ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + gtfs_dataset_id VARCHAR(255) NOT NULL, + feed_publisher_name TEXT, + feed_publisher_url TEXT, + feed_lang VARCHAR(255), + default_lang VARCHAR(255), + feed_start_date DATE, + feed_end_date DATE, + feed_version TEXT, + feed_contact_email TEXT, + feed_contact_url TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + + CONSTRAINT feedinfo_gtfs_dataset_id_fkey + FOREIGN KEY (gtfs_dataset_id) + REFERENCES gtfsdataset(id) + ON DELETE CASCADE, + CONSTRAINT feedinfo_gtfs_dataset_id_key + UNIQUE (gtfs_dataset_id) +); + +CREATE INDEX IF NOT EXISTS idx_feedinfo_gtfs_dataset_id + ON feedinfo (gtfs_dataset_id);