fix: sqlalchemy sync bugs + batch fill (#1134)

Author: cka-y

functions-python/reverse_geolocation/README.md

@@ -16,6 +16,10 @@ This HTTP function initiates reverse geolocation for multiple feeds. It accepts
 
 - **`country_codes`** (optional): A comma-separated list of country codes specifying which feeds should be processed.  
   - If not provided, the function processes feeds from all available countries.
+- **`include_only_unprocessed`** (optional): A boolean flag indicating whether to include only feeds that have not been processed yet.  
+  - If set to `true`, only unprocessed feeds will be considered for reverse geolocation.
+  - If set to `false`, all feeds will be processed, regardless of their processing status.
+  - Default is `true`.
 
 **Behavior:**  
 The function publishes a message to the `reverse-geolocation` Pub/Sub topic for each non deprecated feed that matches the specified country codes.  

functions-python/reverse_geolocation/src/reverse_geolocation_batch.py

@@ -16,7 +16,9 @@
 
 
 @with_db_session
-def get_feeds_data(country_codes: List[str], db_session: Session) -> List[Dict]:
+def get_feeds_data(
+    country_codes: List[str], include_only_unprocessed: bool, db_session: Session
+) -> List[Dict]:
     """Get the feeds data for the given country codes. In case no country codes are provided, fetch feeds for all
     countries."""
     query = (
@@ -35,6 +37,10 @@ def get_feeds_data(country_codes: List[str], db_session: Session) -> List[Dict]:
     else:
         logging.warning("No country codes provided. Fetching feeds for all countries.")
 
+    if include_only_unprocessed:
+        logging.info("Filtering for unprocessed feeds.")
+        query = query.filter(~Gtfsfeed.feedlocationgrouppoints.any())
+
     results = query.populate_existing().all()
     logging.info(f"Found {len(results)} feeds.")
 
@@ -49,24 +55,28 @@ def get_feeds_data(country_codes: List[str], db_session: Session) -> List[Dict]:
     return data
 
 
-def parse_request_parameters(request: flask.Request) -> List[str]:
-    """Parse the request parameters to get the country codes."""
-    country_codes = request.args.get("country_codes", "").split(",")
+def parse_request_parameters(request: flask.Request) -> Tuple[List[str], bool]:
+    """Parse the request parameters to get the country codes and whether to include only unprocessed feeds."""
+    json_request = request.get_json()
+    country_codes = json_request.get("country_codes", "").split(",")
     country_codes = [code.strip().upper() for code in country_codes if code]
 
     # Validate country codes
     for country_code in country_codes:
         if not pycountry.countries.get(alpha_2=country_code):
             raise ValueError(f"Invalid country code: {country_code}")
-    return country_codes
+    include_only_unprocessed = (
+        json_request.get("include_only_unprocessed", True) is True
+    )
+    return country_codes, include_only_unprocessed
 
 
 def reverse_geolocation_batch(request: flask.Request) -> Tuple[str, int]:
     """Batch function to trigger reverse geolocation for feeds."""
     try:
         Logger.init_logger()
-        country_codes = parse_request_parameters(request)
-        feeds_data = get_feeds_data(country_codes)
+        country_codes, include_only_unprocessed = parse_request_parameters(request)
+        feeds_data = get_feeds_data(country_codes, include_only_unprocessed)
         logging.info(f"Valid feeds with latest dataset: {len(feeds_data)}")
 
         pubsub_topic_name = os.getenv("PUBSUB_TOPIC_NAME", None)

functions-python/reverse_geolocation/src/reverse_geolocation_processor.py

@@ -215,6 +215,7 @@ def extract_location_aggregate(
     stop = Feedlocationgrouppoint(feed_id=feed_id, geometry=stop_point)
     stop.group = group
     db_session.add(stop)
+    db_session.flush()
     logging.info(
         f"Point {stop_point} matched to {', '.join([g.name for g in geopolygons])}"
     )
@@ -359,11 +360,13 @@ def extract_location_aggregates(
         )
         for location_group in location_aggregates.values()
     ]
-    gtfs_feed.feedosmlocationgroups = osm_location_groups
+    gtfs_feed.feedosmlocationgroups.clear()
+    gtfs_feed.feedosmlocationgroups.extend(osm_location_groups)
+
     for gtfs_rt_feed in gtfs_feed.gtfs_rt_feeds:
         logging.info(f"Updating GTFS-RT feed with stable ID {gtfs_rt_feed.stable_id}")
         gtfs_rt_feed.feedosmlocationgroups.clear()
-        gtfs_rt_feed.feedosmlocationgroups = osm_location_groups
+        gtfs_rt_feed.feedosmlocationgroups.extend(osm_location_groups)
 
     feed_locations = []
     for location_aggregate in location_aggregates.values():

functions-python/reverse_geolocation/tests/test_reverse_geolocation_batch.py

@@ -56,16 +56,16 @@ def test_get_feed_data(self, db_session):
         db_session.add(feed_2)
         db_session.commit()
 
-        results = get_feeds_data(["CA"])
+        results = get_feeds_data(["CA"], True)
         self.assertEqual(len(results), 1)
         self.assertEqual(results[0]["stable_id"], "test_feed")
         self.assertEqual(results[0]["dataset_id"], "test_dataset_latest")
         self.assertEqual(results[0]["url"], "test_url")
 
-        results_2 = get_feeds_data([])
+        results_2 = get_feeds_data([], False)
         self.assertEqual(len(results_2), 2)
 
-        results_3 = get_feeds_data(["US"])
+        results_3 = get_feeds_data(["US"], True)
         self.assertEqual(len(results_3), 1)
         self.assertEqual(results_3[0]["stable_id"], "test_feed_2")
         self.assertEqual(results_3[0]["dataset_id"], "test_dataset_3")
@@ -75,16 +75,17 @@ def test_get_feed_data(self, db_session):
 
     def test_parse_request_parameters(self):
         request = MagicMock()
-        request.args.get = lambda value, default: {"country_codes": "Ca , uS"}.get(
-            value, default
-        )
+        request.get_json.return_value.get = lambda value, default: {
+            "country_codes": "Ca , uS"
+        }.get(value, default)
         from reverse_geolocation_batch import parse_request_parameters
 
-        country_codes = parse_request_parameters(request)
-        self.assertEqual(country_codes, ["CA", "US"])
+        country_codes, include_only_unprocessed = parse_request_parameters(request)
+        self.assertEqual(["CA", "US"], country_codes)
+        self.assertTrue(include_only_unprocessed)
 
         with pytest.raises(ValueError):
-            request.args.get = lambda value, default: {
+            request.get_json.return_value.get = lambda value, default: {
                 "country_codes": "CA , US, XX"
             }.get(value, default)
             parse_request_parameters(request)