updated populate script

Author: qcdyx

api/src/scripts/populate_db_gtfs.py

@@ -4,6 +4,7 @@
 
 import pycountry
 import pytz
+from sqlalchemy import func
 
 from scripts.load_dataset_on_create import publish_all
 from scripts.populate_db import DatabasePopulateHelper, set_up_configs
@@ -14,6 +15,7 @@
     Gtfsrealtimefeed,
     Location,
     Redirectingid,
+    Gtfsfeed,
 )
 from utils.data_utils import set_up_defaults
 
@@ -116,32 +118,95 @@ def process_entity_types(self, session: "Session", feed: Gtfsrealtimefeed, row,
             self.logger.warning(f"Entity types array is empty for feed {stable_id}")
             feed.entitytypes.clear()
 
+    # def process_feed_references(self, session: "Session"):
+    #     """
+    #     Process the feed references
+    #     """
+    #     self.logger.info("Processing feed references")
+    #     for index, row in self.df.iterrows():
+    #         stable_id = self.get_stable_id(row)
+    #         data_type = self.get_data_type(row)
+    #         if data_type != "gtfs_rt":
+    #             continue
+    #         gtfs_rt_feed = self.query_feed_by_stable_id(session, stable_id, "gtfs_rt")
+    #         static_reference = self.get_safe_value(row, "static_reference", "")
+    #         if static_reference:
+    #             try:
+    #                 gtfs_stable_id = f"mdb-{int(float(static_reference))}"
+    #             except ValueError:
+    #                 gtfs_stable_id = static_reference
+    #             gtfs_feed = self.query_feed_by_stable_id(session, gtfs_stable_id, "gtfs")
+    #             if not gtfs_feed:
+    #                 self.logger.warning(f"Could not find static reference feed {gtfs_stable_id} for feed {stable_id}")
+    #                 continue
+    #             already_referenced_ids = {ref.id for ref in gtfs_feed.gtfs_rt_feeds}
+    #             if gtfs_feed and gtfs_rt_feed.id not in already_referenced_ids:
+    #                 gtfs_feed.gtfs_rt_feeds.append(gtfs_rt_feed)
+    #                 # Flush to avoid FK violation
+    #                 session.flush()
+
     def process_feed_references(self, session: "Session"):
         """
-        Process the feed references
+        Process the feed references for GTFS-RT feeds.
+
+        1. Uses 'static_reference' column if present.
+        2. Falls back to matching static feeds by provider name.
         """
         self.logger.info("Processing feed references")
+
         for index, row in self.df.iterrows():
             stable_id = self.get_stable_id(row)
             data_type = self.get_data_type(row)
+
+            # Only process GTFS-RT feeds
             if data_type != "gtfs_rt":
                 continue
+
             gtfs_rt_feed = self.query_feed_by_stable_id(session, stable_id, "gtfs_rt")
-            static_reference = self.get_safe_value(row, "static_reference", "")
+            if not gtfs_rt_feed:
+                self.logger.warning(f"Could not find GTFS-RT feed {stable_id}")
+                continue
+
+            # Try static_reference column first
+            static_reference = self.get_safe_value(row, "static_reference", "").strip()
+            gtfs_feed = None
+
             if static_reference:
+                # Normalize stable_id
                 try:
                     gtfs_stable_id = f"mdb-{int(float(static_reference))}"
                 except ValueError:
                     gtfs_stable_id = static_reference
+
                 gtfs_feed = self.query_feed_by_stable_id(session, gtfs_stable_id, "gtfs")
                 if not gtfs_feed:
                     self.logger.warning(f"Could not find static reference feed {gtfs_stable_id} for feed {stable_id}")
-                    continue
+
+            # Fallback: match by provider if no static_reference or not found
+            if not gtfs_feed:
+                provider_value = (self.get_safe_value(row, "provider", "") or "").strip().lower()
+                if provider_value:
+                    gtfs_feed = (
+                        session.query(Gtfsfeed)
+                        .filter(
+                            Gtfsfeed.data_type == "gtfs",
+                            func.lower(func.trim(Gtfsfeed.provider)) == provider_value,
+                            Gtfsfeed.stable_id != stable_id,
+                        )
+                        .first()
+                    )
+                    if not gtfs_feed:
+                        self.logger.warning(
+                            f"No static GTFS feed found for provider '{provider_value}' for feed {stable_id}"
+                        )
+
+            # Link the feeds if we have a valid static GTFS feed
+            if gtfs_feed:
                 already_referenced_ids = {ref.id for ref in gtfs_feed.gtfs_rt_feeds}
-                if gtfs_feed and gtfs_rt_feed.id not in already_referenced_ids:
+                if gtfs_rt_feed.id not in already_referenced_ids:
                     gtfs_feed.gtfs_rt_feeds.append(gtfs_rt_feed)
-                    # Flush to avoid FK violation
-                    session.flush()
+                    session.flush()  # Avoid FK violations
+                    self.logger.info(f"Linked GTFS-RT feed {stable_id} to static feed {gtfs_feed.stable_id}")
 
     def process_redirects(self, session: "Session"):
         """

api/src/shared/db_models/gtfs_rt_feed_impl.py

@@ -20,23 +20,13 @@ class Config:
         from_attributes = True
 
     @classmethod
-    @with_db_session
-    def from_orm(cls, feed: Gtfsrealtimefeed | None, db_session: Session) -> GtfsRTFeed | None:
+    def from_orm(cls, feed: Gtfsrealtimefeed | None) -> GtfsRTFeed | None:
         gtfs_rt_feed: GtfsRTFeed = super().from_orm(feed)
         if not gtfs_rt_feed:
             return None
         gtfs_rt_feed.locations = [LocationImpl.from_orm(item) for item in feed.locations] if feed.locations else []
         gtfs_rt_feed.entity_types = [item.name for item in feed.entitytypes] if feed.entitytypes else []
-
-        provider_value = (feed.provider or "").strip().lower()
-
-        query = db_session.query(FeedOrm).filter(
-            func.lower(func.trim(FeedOrm.provider)) == provider_value,
-            FeedOrm.stable_id != feed.stable_id,
-        )
-
-        gtfs_rt_feed.feed_references = [gtfs_feed.stable_id for gtfs_feed in query.all()]
-
+        gtfs_rt_feed.feed_references = [item.stable_id for item in feed.gtfs_feeds] if feed.gtfs_feeds else []
         return gtfs_rt_feed
 
     @classmethod