fix: limit and offset behaviour for feeds with multiple locations (#773)

Author: cka-y

.github/workflows/db-update-qa.yml

@@ -3,7 +3,8 @@ name: Database Update - QA
 on:
   workflow_dispatch:
   workflow_call:
-
+  repository_dispatch: # Update on mobility-database-catalog repo dispatch
+    types: [ catalog-sources-updated ]
 jobs:
   update:
     uses: ./.github/workflows/db-update.yml

api/src/feeds/impl/feeds_api_impl.py

@@ -1,8 +1,9 @@
 from datetime import datetime
 from typing import List, Union, TypeVar
-
+from sqlalchemy import select
 from sqlalchemy.orm import joinedload
 from sqlalchemy.orm.query import Query
+
 from database.database import Database
 from database_gen.sqlacodegen_models import (
     Feed,
@@ -11,8 +12,8 @@
     Gtfsrealtimefeed,
     Location,
     Validationreport,
-    Entitytype,
     t_location_with_translations_en,
+    Entitytype,
 )
 from feeds.filters.feed_filter import FeedFilter
 from feeds.filters.gtfs_dataset_filter import GtfsDatasetFilter
@@ -37,7 +38,11 @@
 from feeds_gen.models.gtfs_feed import GtfsFeed
 from feeds_gen.models.gtfs_rt_feed import GtfsRTFeed
 from utils.date_utils import valid_iso_date
-from utils.location_translation import create_location_translation_object, LocationTranslation
+from utils.location_translation import (
+    create_location_translation_object,
+    LocationTranslation,
+    get_feeds_location_translations,
+)
 
 T = TypeVar("T", bound="BasicFeed")
 
@@ -197,25 +202,28 @@ def get_gtfs_feeds(
                 municipality__ilike=municipality,
             ),
         )
-        gtfs_feed_query = gtfs_feed_filter.filter(
-            Database().get_session().query(Gtfsfeed, t_location_with_translations_en)
-        )
-        gtfs_feed_query = (
-            gtfs_feed_query.outerjoin(Location, Feed.locations)
-            .outerjoin(t_location_with_translations_en, Location.id == t_location_with_translations_en.c.location_id)
+
+        subquery = gtfs_feed_filter.filter(select(Gtfsfeed.id).join(Location, Gtfsfeed.locations))
+        subquery = DatasetsApiImpl.apply_bounding_filtering(
+            subquery, dataset_latitudes, dataset_longitudes, bounding_filter_method
+        ).subquery()
+
+        feed_query = (
+            Database()
+            .get_session()
+            .query(Gtfsfeed)
+            .filter(Gtfsfeed.id.in_(subquery))
             .options(
                 joinedload(Gtfsfeed.gtfsdatasets)
                 .joinedload(Gtfsdataset.validation_reports)
                 .joinedload(Validationreport.notices),
                 *BasicFeedImpl.get_joinedload_options(),
             )
             .order_by(Gtfsfeed.provider, Gtfsfeed.stable_id)
+            .limit(limit)
+            .offset(offset)
         )
-        gtfs_feed_query = gtfs_feed_query.order_by(Gtfsfeed.provider, Gtfsfeed.stable_id)
-        gtfs_feed_query = DatasetsApiImpl.apply_bounding_filtering(
-            gtfs_feed_query, dataset_latitudes, dataset_longitudes, bounding_filter_method
-        )
-        return self._get_response(gtfs_feed_query, limit, offset, GtfsFeedImpl)
+        return self._get_response(feed_query, GtfsFeedImpl)
 
     def get_gtfs_rt_feed(
         self,
@@ -283,32 +291,33 @@ def get_gtfs_rt_feeds(
                 municipality__ilike=municipality,
             ),
         )
-        gtfs_rt_feed_query = gtfs_rt_feed_filter.filter(
-            Database().get_session().query(Gtfsrealtimefeed, t_location_with_translations_en)
-        )
-        gtfs_rt_feed_query = (
-            gtfs_rt_feed_query.outerjoin(Location, Gtfsrealtimefeed.locations)
-            .outerjoin(t_location_with_translations_en, Location.id == t_location_with_translations_en.c.location_id)
-            .outerjoin(Entitytype, Gtfsrealtimefeed.entitytypes)
+        subquery = gtfs_rt_feed_filter.filter(
+            select(Gtfsrealtimefeed.id)
+            .join(Location, Gtfsrealtimefeed.locations)
+            .join(Entitytype, Gtfsrealtimefeed.entitytypes)
+        ).subquery()
+        feed_query = (
+            Database()
+            .get_session()
+            .query(Gtfsrealtimefeed)
+            .filter(Gtfsrealtimefeed.id.in_(subquery))
             .options(
                 joinedload(Gtfsrealtimefeed.entitytypes),
                 joinedload(Gtfsrealtimefeed.gtfs_feeds),
                 *BasicFeedImpl.get_joinedload_options(),
             )
             .order_by(Gtfsrealtimefeed.provider, Gtfsrealtimefeed.stable_id)
+            .limit(limit)
+            .offset(offset)
         )
-        return self._get_response(gtfs_rt_feed_query, limit, offset, GtfsRTFeedImpl)
+        return self._get_response(feed_query, GtfsRTFeedImpl)
 
     @staticmethod
-    def _get_response(feed_query: Query, limit: int, offset: int, impl_cls: type[T]) -> List[T]:
+    def _get_response(feed_query: Query, impl_cls: type[T]) -> List[T]:
         """Get the response for the feed query."""
-        if limit is not None:
-            feed_query = feed_query.limit(limit)
-        if offset is not None:
-            feed_query = feed_query.offset(offset)
         results = feed_query.all()
-        location_translations = {row[1]: create_location_translation_object(row) for row in results}
-        response = [impl_cls.from_orm(feed[0], location_translations) for feed in results]
+        location_translations = get_feeds_location_translations(results)
+        response = [impl_cls.from_orm(feed, location_translations) for feed in results]
         return list({feed.id: feed for feed in response}.values())
 
     def get_gtfs_feed_gtfs_rt_feeds(

api/src/scripts/populate_db_test_data.py

@@ -1,11 +1,21 @@
 import json
+from uuid import uuid4
 
 from geoalchemy2 import WKTElement
+from google.cloud.sql.connector.instance import logger
 from sqlalchemy import text
 
 from database.database import Database
-from database_gen.sqlacodegen_models import Gtfsdataset, Validationreport, Gtfsfeed, Notice, Feature, t_feedsearch
-from scripts.populate_db import set_up_configs
+from database_gen.sqlacodegen_models import (
+    Gtfsdataset,
+    Validationreport,
+    Gtfsfeed,
+    Notice,
+    Feature,
+    t_feedsearch,
+    Location,
+)
+from scripts.populate_db import set_up_configs, DatabasePopulateHelper
 from utils.logger import Logger
 
 
@@ -36,70 +46,78 @@ def populate_test_datasets(self, filepath):
         with open(filepath) as f:
             data = json.load(f)
 
+        # GTFS Feeds
+        if "feeds" in data:
+            self.populate_test_feeds(data["feeds"])
+
         # GTFS Datasets
         dataset_dict = {}
-        for dataset in data["datasets"]:
-            # query the db using feed_id to get the feed object
-            gtfsfeed = self.db.session.query(Gtfsfeed).filter(Gtfsfeed.stable_id == dataset["feed_stable_id"]).all()
-            if not gtfsfeed:
-                self.logger.error(f"No feed found with stable_id: {dataset['feed_stable_id']}")
-                continue
-
-            gtfs_dataset = Gtfsdataset(
-                id=dataset["id"],
-                feed_id=gtfsfeed[0].id,
-                stable_id=dataset["id"],
-                latest=dataset["latest"],
-                hosted_url=dataset["hosted_url"],
-                hash=dataset["hash"],
-                downloaded_at=dataset["downloaded_at"],
-                bounding_box=None
-                if dataset.get("bounding_box") is None
-                else WKTElement(dataset["bounding_box"], srid=4326),
-                validation_reports=[],
-            )
-            dataset_dict[dataset["id"]] = gtfs_dataset
-            self.db.session.add(gtfs_dataset)
+        if "datasets" in data:
+            for dataset in data["datasets"]:
+                # query the db using feed_id to get the feed object
+                gtfsfeed = self.db.session.query(Gtfsfeed).filter(Gtfsfeed.stable_id == dataset["feed_stable_id"]).all()
+                if not gtfsfeed:
+                    self.logger.error(f"No feed found with stable_id: {dataset['feed_stable_id']}")
+                    continue
+
+                gtfs_dataset = Gtfsdataset(
+                    id=dataset["id"],
+                    feed_id=gtfsfeed[0].id,
+                    stable_id=dataset["id"],
+                    latest=dataset["latest"],
+                    hosted_url=dataset["hosted_url"],
+                    hash=dataset["hash"],
+                    downloaded_at=dataset["downloaded_at"],
+                    bounding_box=None
+                    if dataset.get("bounding_box") is None
+                    else WKTElement(dataset["bounding_box"], srid=4326),
+                    validation_reports=[],
+                )
+                dataset_dict[dataset["id"]] = gtfs_dataset
+                self.db.session.add(gtfs_dataset)
 
         # Validation reports
-        validation_report_dict = {}
-        for report in data["validation_reports"]:
-            validation_report = Validationreport(
-                id=report["id"],
-                validator_version=report["validator_version"],
-                validated_at=report["validated_at"],
-                html_report=report["html_report"],
-                json_report=report["json_report"],
-                features=[],
-            )
-            dataset_dict[report["dataset_id"]].validation_reports.append(validation_report)
-            validation_report_dict[report["id"]] = validation_report
-            self.db.session.add(validation_report)
+        if "validation_reports" in data:
+            validation_report_dict = {}
+            for report in data["validation_reports"]:
+                validation_report = Validationreport(
+                    id=report["id"],
+                    validator_version=report["validator_version"],
+                    validated_at=report["validated_at"],
+                    html_report=report["html_report"],
+                    json_report=report["json_report"],
+                    features=[],
+                )
+                dataset_dict[report["dataset_id"]].validation_reports.append(validation_report)
+                validation_report_dict[report["id"]] = validation_report
+                self.db.session.add(validation_report)
 
         # Notices
-        for report_notice in data["notices"]:
-            notice = Notice(
-                dataset_id=report_notice["dataset_id"],
-                validation_report_id=report_notice["validation_report_id"],
-                severity=report_notice["severity"],
-                notice_code=report_notice["notice_code"],
-                total_notices=report_notice["total_notices"],
-            )
-            self.db.session.add(notice)
+        if "notices" in data:
+            for report_notice in data["notices"]:
+                notice = Notice(
+                    dataset_id=report_notice["dataset_id"],
+                    validation_report_id=report_notice["validation_report_id"],
+                    severity=report_notice["severity"],
+                    notice_code=report_notice["notice_code"],
+                    total_notices=report_notice["total_notices"],
+                )
+                self.db.session.add(notice)
         # Features
-        for featureName in data["features"]:
-            feature = Feature(name=featureName)
-            self.db.session.add(feature)
+        if "features" in data:
+            for featureName in data["features"]:
+                feature = Feature(name=featureName)
+                self.db.session.add(feature)
 
         # Features in Validation Reports
-        for report_features in data["validation_report_features"]:
-            validation_report_dict[report_features["validation_report_id"]].features.append(
-                self.db.session.query(Feature).filter(Feature.name == report_features["feature_name"]).first()
-            )
-
-        self.db.session.execute(text(f"REFRESH MATERIALIZED VIEW CONCURRENTLY {t_feedsearch.name}"))
+        if "validation_report_features" in data:
+            for report_features in data["validation_report_features"]:
+                validation_report_dict[report_features["validation_report_id"]].features.append(
+                    self.db.session.query(Feature).filter(Feature.name == report_features["feature_name"]).first()
+                )
 
         self.db.session.commit()
+        self.db.session.execute(text(f"REFRESH MATERIALIZED VIEW CONCURRENTLY {t_feedsearch.name}"))
 
     def populate(self):
         """
@@ -116,6 +134,47 @@ def populate(self):
 
         self.logger.info("Database populated with test data")
 
+    def populate_test_feeds(self, feeds_data):
+        for feed_data in feeds_data:
+            feed = Gtfsfeed(
+                id=str(uuid4()),
+                stable_id=feed_data["id"],
+                data_type=feed_data["data_type"],
+                status=feed_data["status"],
+                created_at=feed_data["created_at"],
+                provider=feed_data["provider"],
+                feed_name=feed_data["feed_name"],
+                note=feed_data["note"],
+                authentication_info_url=None,
+                api_key_parameter_name=None,
+                license_url=None,
+                feed_contact_email=feed_data["feed_contact_email"],
+                producer_url=feed_data["source_info"]["producer_url"],
+            )
+            locations = []
+            for location_data in feed_data["locations"]:
+                location_id = DatabasePopulateHelper.get_location_id(
+                    location_data["country_code"],
+                    location_data["subdivision_name"],
+                    location_data["municipality"],
+                )
+                location = self.db.session.get(Location, location_id)
+                location = (
+                    location
+                    if location
+                    else Location(
+                        id=location_id,
+                        country_code=location_data["country_code"],
+                        subdivision_name=location_data["subdivision_name"],
+                        municipality=location_data["municipality"],
+                        country=location_data["country"],
+                    )
+                )
+                locations.append(location)
+            feed.locations = locations
+            self.db.session.add(feed)
+            logger.info(f"Added feed {feed.stable_id}")
+
 
 if __name__ == "__main__":
     db_helper = DatabasePopulateTestDataHelper(set_up_configs())