fix: removed redundant location extraction and refactored code (#985)

Author: cka-y

api/src/feeds/impl/feeds_api_impl.py

@@ -15,7 +15,6 @@
     Gtfsrealtimefeed,
     Location,
     Validationreport,
-    t_location_with_translations_en,
     Entitytype,
 )
 from shared.feed_filters.feed_filter import FeedFilter
@@ -42,10 +41,6 @@
 from feeds.impl.error_handling import raise_http_error, raise_http_validation_error, convert_exception
 from middleware.request_context import is_user_email_restricted
 from utils.date_utils import valid_iso_date
-from utils.location_translation import (
-    create_location_translation_object,
-    get_feeds_location_translations,
-)
 from utils.logger import Logger
 
 T = TypeVar("T", bound="BasicFeed")
@@ -253,7 +248,7 @@ def get_gtfs_feeds(
             # that needs to be converted to HTTPException before being thrown.
             raise convert_exception(e)
 
-        return self._get_response(feed_query, GtfsFeedImpl, db_session)
+        return self._get_response(feed_query, GtfsFeedImpl)
 
     @with_db_session
     def get_gtfs_rt_feed(self, id: str, db_session: Session) -> GtfsRTFeed:
@@ -266,25 +261,23 @@ def get_gtfs_rt_feed(self, id: str, db_session: Session) -> GtfsRTFeed:
             location=None,
         )
         results = gtfs_rt_feed_filter.filter(
-            db_session.query(Gtfsrealtimefeed, t_location_with_translations_en)
+            db_session.query(Gtfsrealtimefeed)
             .filter(
                 or_(
                     Gtfsrealtimefeed.operational_status == "published",
                     not is_user_email_restricted(),  # Allow all feeds to be returned if the user is not restricted
                 )
             )
             .outerjoin(Location, Gtfsrealtimefeed.locations)
-            .outerjoin(t_location_with_translations_en, Location.id == t_location_with_translations_en.c.location_id)
             .options(
                 joinedload(Gtfsrealtimefeed.entitytypes),
                 joinedload(Gtfsrealtimefeed.gtfs_feeds),
                 *get_joinedload_options(),
             )
         ).all()
 
-        if len(results) > 0 and results[0].Gtfsrealtimefeed:
-            translations = {result[1]: create_location_translation_object(result) for result in results}
-            return GtfsRTFeedImpl.from_orm(results[0].Gtfsrealtimefeed, translations)
+        if len(results) > 0 and results[0]:
+            return GtfsRTFeedImpl.from_orm(results[0])
         else:
             raise_http_error(404, gtfs_rt_feed_not_found.format(id))
 
@@ -321,7 +314,7 @@ def get_gtfs_rt_feeds(
         except InternalHTTPException as e:
             raise convert_exception(e)
 
-        return self._get_response(feed_query, GtfsRTFeedImpl, db_session)
+        return self._get_response(feed_query, GtfsRTFeedImpl)
 
         entity_types_list = entity_types.split(",") if entity_types else None
 
@@ -370,14 +363,13 @@ def get_gtfs_rt_feeds(
         if is_official:
             feed_query = feed_query.filter(Feed.official)
         feed_query = feed_query.limit(limit).offset(offset)
-        return self._get_response(feed_query, GtfsRTFeedImpl, db_session)
+        return self._get_response(feed_query, GtfsRTFeedImpl)
 
     @staticmethod
-    def _get_response(feed_query: Query, impl_cls: type[T], db_session: "Session") -> List[T]:
+    def _get_response(feed_query: Query, impl_cls: type[T]) -> List[T]:
         """Get the response for the feed query."""
         results = feed_query.all()
-        location_translations = get_feeds_location_translations(results, db_session)
-        response = [impl_cls.from_orm(feed, location_translations) for feed in results]
+        response = [impl_cls.from_orm(feed) for feed in results]
         return list({feed.id: feed for feed in response}.values())
 
     @with_db_session

api/src/feeds/impl/models/basic_feed_impl.py

@@ -17,7 +17,7 @@ class Config:
         from_attributes = True
 
     @classmethod
-    def from_orm(cls, feed: Feed | None, _=None) -> BasicFeed | None:
+    def from_orm(cls, feed: Feed | None) -> BasicFeed | None:
         if not feed:
             return None
         return cls(

api/src/feeds/impl/models/gtfs_feed_impl.py

@@ -1,11 +1,8 @@
-from typing import Dict
-
 from shared.database_gen.sqlacodegen_models import Gtfsfeed as GtfsfeedOrm
 from feeds.impl.models.basic_feed_impl import BaseFeedImpl
 from feeds.impl.models.latest_dataset_impl import LatestDatasetImpl
 from feeds.impl.models.location_impl import LocationImpl
 from feeds_gen.models.gtfs_feed import GtfsFeed
-from utils.location_translation import LocationTranslation, translate_feed_locations
 
 
 class GtfsFeedImpl(BaseFeedImpl, GtfsFeed):
@@ -20,11 +17,7 @@ class Config:
         from_attributes = True
 
     @classmethod
-    def from_orm(
-        cls, feed: GtfsfeedOrm | None, location_translations: Dict[str, LocationTranslation] = None
-    ) -> GtfsFeed | None:
-        if location_translations is not None:
-            translate_feed_locations(feed, location_translations)
+    def from_orm(cls, feed: GtfsfeedOrm | None) -> GtfsFeed | None:
         gtfs_feed: GtfsFeed = super().from_orm(feed)
         if not gtfs_feed:
             return None

api/src/feeds/impl/models/gtfs_rt_feed_impl.py

@@ -1,10 +1,7 @@
-from typing import Dict
-
 from shared.database_gen.sqlacodegen_models import Gtfsrealtimefeed as GtfsRTFeedOrm
 from feeds.impl.models.basic_feed_impl import BaseFeedImpl
 from feeds.impl.models.location_impl import LocationImpl
 from feeds_gen.models.gtfs_rt_feed import GtfsRTFeed
-from utils.location_translation import LocationTranslation, translate_feed_locations
 
 
 class GtfsRTFeedImpl(BaseFeedImpl, GtfsRTFeed):
@@ -17,11 +14,7 @@ class Config:
         from_attributes = True
 
     @classmethod
-    def from_orm(
-        cls, feed: GtfsRTFeedOrm | None, location_translations: Dict[str, LocationTranslation] = None
-    ) -> GtfsRTFeed | None:
-        if location_translations is not None:
-            translate_feed_locations(feed, location_translations)
+    def from_orm(cls, feed: GtfsRTFeedOrm | None) -> GtfsRTFeed | None:
         gtfs_rt_feed: GtfsRTFeed = super().from_orm(feed)
         if not gtfs_rt_feed:
             return None

api/src/feeds/impl/models/location_impl.py

@@ -1,4 +1,5 @@
 from feeds_gen.models.location import Location
+import pycountry
 from shared.database_gen.sqlacodegen_models import Location as LocationOrm
 
 
@@ -14,9 +15,15 @@ def from_orm(cls, location: LocationOrm | None) -> Location | None:
         """Create a model instance from a SQLAlchemy a Location row object."""
         if not location:
             return None
+        country_name = location.country
+        if not country_name:
+            try:
+                country_name = pycountry.countries.get(alpha_2=location.country_code).name
+            except AttributeError:
+                pass
         return cls(
             country_code=location.country_code,
-            country=location.country,
+            country=country_name,
             subdivision_name=location.subdivision_name,
             municipality=location.municipality,
         )

api/src/feeds/impl/models/search_feed_item_result_impl.py

@@ -38,7 +38,7 @@ def from_orm_gtfs(cls, feed_search_row):
                 license_url=feed_search_row.license_url,
             ),
             redirects=feed_search_row.redirect_ids,
-            locations=feed_search_row.locations,
+            locations=cls.resolve_locations(feed_search_row.locations),
             latest_dataset=LatestDataset(
                 id=feed_search_row.latest_dataset_id,
                 hosted_url=feed_search_row.latest_dataset_hosted_url,
@@ -74,55 +74,39 @@ def from_orm_gtfs_rt(cls, feed_search_row):
                 license_url=feed_search_row.license_url,
             ),
             redirects=feed_search_row.redirect_ids,
-            locations=feed_search_row.locations,
+            locations=cls.resolve_locations(feed_search_row.locations),
             entity_types=feed_search_row.entities,
             feed_references=feed_search_row.feed_reference_ids,
         )
 
     @classmethod
-    def _translate_locations(cls, feed_search_row):
-        """Translate location information in the feed search row.
-        This method modifies the locations in the feed search row in place."""
-        if feed_search_row.locations is None:
-            return
-        country_translations = cls._create_translation_dict(feed_search_row.country_translations)
-        subdivision_translations = cls._create_translation_dict(feed_search_row.subdivision_name_translations)
-        municipality_translations = cls._create_translation_dict(feed_search_row.municipality_translations)
-
-        for location in feed_search_row.locations:
-            location["country"] = country_translations.get(location["country"], location["country"])
-            if location["country"] is None or len(location["country"]) == 0:
-                location["country"] = SearchFeedItemResultImpl.resolve_country_by_code(location)
-            location["subdivision_name"] = subdivision_translations.get(
-                location["subdivision_name"], location["subdivision_name"]
-            )
-            location["municipality"] = municipality_translations.get(location["municipality"], location["municipality"])
+    def resolve_locations(cls, locations):
+        """Resolve locations by country code."""
+        return [
+            {
+                **location,
+                "country": location.get("country")
+                if location.get("country")
+                else cls.resolve_country_by_code(location),
+            }
+            for location in locations
+        ]
 
     @classmethod
     def resolve_country_by_code(cls, location):
         """Resolve country name by country code.
         If the country code is not found, return the original country name."""
-        country = pycountry.countries.get(alpha_2=location["country_code"])
-        return country.name if country else location["country"]
-
-    @staticmethod
-    def _create_translation_dict(translations):
-        """Helper method to create a translation dictionary."""
-        if translations:
-            return {
-                elem.get("key"): elem.get("value") for elem in translations if elem.get("key") and elem.get("value")
-            }
-        return {}
+        try:
+            country = pycountry.countries.get(alpha_2=location.get("country_code"))
+            return country.name if country else location.get("country")
+        except AttributeError:
+            return location.get("country")
 
     @classmethod
     def from_orm(cls, feed_search_row):
         """Create a model instance from a SQLAlchemy row object."""
         if feed_search_row is None:
             return None
-
-        # Translate location data
-        cls._translate_locations(feed_search_row)
-
         match feed_search_row.data_type:
             case "gtfs":
                 return cls.from_orm_gtfs(feed_search_row)

api/src/utils/location_translation.py

@@ -182,7 +182,7 @@ class TestGtfsFeedImpl(unittest.TestCase):
 
     def test_from_orm_all_fields(self):
         """Test the `from_orm` method with all fields."""
-        result = GtfsFeedImpl.from_orm(gtfs_feed_orm, {})
+        result = GtfsFeedImpl.from_orm(gtfs_feed_orm)
         assert result == expected_gtfs_feed_result
 
     def test_from_orm_empty_fields(self):