fix: failing metrics cloud functions (#942)

Author: cka-y

.gitignore

@@ -74,3 +74,6 @@ __pycache__
 .coverage
 coverage_reports
 tf.plan
+
+# CSV generation output files
+functions-python/**/*.csv

functions-python/big_query_ingestion/src/gbfs/gbfs_big_query_ingest.py

@@ -10,5 +10,5 @@ def __init__(self):
         super().__init__()
         current_dir = os.path.dirname(os.path.abspath(__file__))
         self.schema_path = os.path.join(
-            current_dir, "../helpers/bq_schema/gbfs_schema.json"
+            current_dir, "../shared/helpers/bq_schema/gbfs_schema.json"
         )

functions-python/big_query_ingestion/src/gtfs/gtfs_big_query_ingest.py

@@ -10,5 +10,5 @@ def __init__(self):
         super().__init__()
         current_dir = os.path.dirname(os.path.abspath(__file__))
         self.schema_path = os.path.join(
-            current_dir, "../helpers/bq_schema/gtfs_schema.json"
+            current_dir, "../shared/helpers/bq_schema/gtfs_schema.json"
         )

functions-python/preprocessed_analytics/function_config.json

@@ -3,7 +3,7 @@
   "description": "Preprocess analytics",
   "entry_point": "preprocess_analytics",
   "timeout": 540,
-  "memory": "2Gi",
+  "memory": "4Gi",
   "trigger_http": false,
   "include_folders": ["helpers"],
   "include_api_folders": ["database_gen"],

functions-python/preprocessed_analytics/src/processors/base_analytics_processor.py

@@ -6,14 +6,14 @@
 import pandas as pd
 from google.cloud import storage
 from sqlalchemy.orm import Query
-
+from sqlalchemy.orm.session import Session
 from shared.database_gen.sqlacodegen_models import (
     Gbfsfeed,
     Gbfssnapshot,
     Gtfsfeed,
     Gtfsdataset,
 )
-from shared.helpers.database import Database
+from shared.helpers.database import with_db_session
 
 
 class NoFeedDataException(Exception):
@@ -23,7 +23,6 @@ class NoFeedDataException(Exception):
 class BaseAnalyticsProcessor:
     def __init__(self, run_date):
         self.run_date = run_date
-        self.session = Database().start_db_session(echo=False)
         self.processed_feeds = set()
         self.data = []
         self.feed_metrics_data = []
@@ -33,14 +32,11 @@ def __init__(self, run_date):
             os.getenv("ANALYTICS_BUCKET")
         )
 
-    def get_latest_data(self) -> Query:
+    def get_latest_data(self, db_session: Session) -> Query:
         raise NotImplementedError("Subclasses should implement this method.")
 
     def process_feed_data(
-        self,
-        feed: Gtfsfeed | Gbfsfeed,
-        dataset_or_snapshot: Gtfsdataset | Gbfssnapshot,
-        translations: Dict,
+        self, feed: Gtfsfeed | Gbfsfeed, dataset_or_snapshot: Gtfsdataset | Gbfssnapshot
     ) -> None:
         raise NotImplementedError("Subclasses should implement this method.")
 
@@ -135,65 +131,25 @@ def save_analytics(self) -> None:
         self.save()
         logging.info(f"Analytics saved to bucket as {file_name}")
 
-    def run(self) -> None:
-        for (
-            feed,
-            dataset_or_snapshot,
-            translation_fields,
-        ) in self._get_data_with_translations():
-            self.process_feed_data(feed, dataset_or_snapshot, translation_fields)
+    @with_db_session
+    def run(self, db_session: Session) -> None:
+        for feed, dataset_or_snapshot in self._get_data(db_session):
+            self.process_feed_data(feed, dataset_or_snapshot)
 
-        self.session.close()
         self.save_summary()
         self.save_analytics()
         self.update_analytics_files()
         logging.info(f"Finished running analytics for date: {self.run_date}")
 
-    def _get_data_with_translations(self):
-        query = self.get_latest_data()
+    def _get_data(self, db_session: Session):
+        query = self.get_latest_data(db_session)
         all_results = query.all()
         if len(all_results) == 0:
             raise NoFeedDataException("No feed data found")
         logging.info(f"Loaded {len(all_results)} feeds to process")
-        try:
-            location_translations = [
-                self._extract_translation_fields(result[2:]) for result in all_results
-            ]
-            logging.info("Location translations loaded")
-            location_translations_dict = {
-                translation["location_id"]: translation
-                for translation in location_translations
-                if translation["location_id"] is not None
-            }
-        except Exception as e:
-            logging.warning(
-                f"Error loading location translations: {e}\n Continuing without translations"
-            )
-            location_translations_dict = {}
         unique_feeds = {result[0].stable_id: result for result in all_results}
         logging.info(f"Nb of unique feeds loaded: {len(unique_feeds)}")
-        return [
-            (result[0], result[1], location_translations_dict)
-            for result in unique_feeds.values()
-        ]
-
-    @staticmethod
-    def _extract_translation_fields(translation_data):
-        keys = [
-            "location_id",
-            "country_code",
-            "country",
-            "subdivision_name",
-            "municipality",
-            "country_translation",
-            "subdivision_name_translation",
-            "municipality_translation",
-        ]
-        try:
-            return dict(zip(keys, translation_data))
-        except Exception as e:
-            logging.error(f"Error extracting translation fields: {e}")
-            return dict(zip(keys, [None] * len(keys)))
+        return [(result[0], result[1]) for result in unique_feeds.values()]
 
     def update_analytics_files(self) -> None:
         try:

functions-python/preprocessed_analytics/src/processors/gbfs_analytics_processor.py

@@ -1,6 +1,7 @@
 from typing import List
 
 import sqlalchemy
+from sqlalchemy.orm.session import Session
 from sqlalchemy.orm import joinedload
 from sqlalchemy.sql import func, and_
 
@@ -18,9 +19,9 @@ def __init__(self, run_date):
         super().__init__(run_date)
         self.versions_metrics_data = []
 
-    def get_latest_data(self) -> sqlalchemy.orm.Query:
+    def get_latest_data(self, db_session: Session) -> sqlalchemy.orm.Query:
         subquery = (
-            self.session.query(
+            db_session.query(
                 Gbfssnapshot.feed_id,
                 func.max(Gbfssnapshot.downloaded_at).label("max_downloaded_at"),
             )
@@ -30,7 +31,7 @@ def get_latest_data(self) -> sqlalchemy.orm.Query:
         )
 
         query = (
-            self.session.query(Gbfsfeed, Gbfssnapshot)
+            db_session.query(Gbfsfeed, Gbfssnapshot)
             .join(Gbfssnapshot, Gbfsfeed.id == Gbfssnapshot.feed_id)
             .join(
                 subquery,
@@ -50,7 +51,7 @@ def get_latest_data(self) -> sqlalchemy.orm.Query:
         )
         return query
 
-    def process_feed_data(self, feed: Gbfsfeed, snapshot: Gbfssnapshot, _) -> None:
+    def process_feed_data(self, feed: Gbfsfeed, snapshot: Gbfssnapshot) -> None:
         if feed.stable_id in self.processed_feeds:
             return
         self.processed_feeds.add(feed.stable_id)

functions-python/preprocessed_analytics/src/processors/gtfs_analytics_processor.py

@@ -1,19 +1,18 @@
-from typing import List, Dict
+from typing import List
 
 import sqlalchemy
 from sqlalchemy.orm import joinedload
+from sqlalchemy.orm.session import Session
 from sqlalchemy.sql import func, and_
+
 from shared.database_gen.sqlacodegen_models import (
     Gtfsdataset,
     Gtfsfeed,
     Validationreport,
     Notice,
     Feature,
     Feed,
-    t_location_with_translations_en,
-    Location,
 )
-from shared.helpers.locations import translate_feed_locations
 from .base_analytics_processor import BaseAnalyticsProcessor
 
 
@@ -22,9 +21,9 @@ def __init__(self, run_date):
         super().__init__(run_date)
         self.features_metrics_data = []
 
-    def get_latest_data(self) -> sqlalchemy.orm.Query:
+    def get_latest_data(self, db_session: Session) -> sqlalchemy.orm.Query:
         subquery = (
-            self.session.query(
+            db_session.query(
                 Gtfsdataset.feed_id,
                 func.max(Gtfsdataset.downloaded_at).label("max_downloaded_at"),
             )
@@ -34,7 +33,7 @@ def get_latest_data(self) -> sqlalchemy.orm.Query:
         )
 
         query = (
-            self.session.query(Gtfsfeed, Gtfsdataset, t_location_with_translations_en)
+            db_session.query(Gtfsfeed, Gtfsdataset)
             .join(Gtfsdataset, Gtfsfeed.id == Gtfsdataset.feed_id)
             .join(
                 subquery,
@@ -43,11 +42,6 @@ def get_latest_data(self) -> sqlalchemy.orm.Query:
                     Gtfsdataset.downloaded_at == subquery.c.max_downloaded_at,
                 ),
             )
-            .outerjoin(Location, Feed.locations)
-            .outerjoin(
-                t_location_with_translations_en,
-                Location.id == t_location_with_translations_en.c.location_id,
-            )
             .where(Gtfsfeed.status != "deprecated")
             .options(
                 joinedload(Gtfsfeed.locations),
@@ -72,9 +66,7 @@ def save_summary(self) -> None:
         }
         self._save_json(summary_file_name, summary_data)
 
-    def process_feed_data(
-        self, feed: Feed, dataset: Gtfsdataset, translations: Dict
-    ) -> None:
+    def process_feed_data(self, feed: Feed, dataset: Gtfsdataset) -> None:
         if feed.stable_id in self.processed_feeds:
             return
         self.processed_feeds.add(feed.stable_id)
@@ -83,8 +75,6 @@ def process_feed_data(
         if not validation_reports:
             return
 
-        translate_feed_locations(feed, translations)
-
         latest_validation_report = max(validation_reports, key=lambda x: x.validated_at)
         notices = latest_validation_report.notices
         errors = [notice for notice in notices if notice.severity == "ERROR"]

functions-python/preprocessed_analytics/tests/test_base_processor.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 from unittest.mock import patch, MagicMock
 from datetime import datetime
@@ -6,12 +7,12 @@
 from processors.base_analytics_processor import (
     BaseAnalyticsProcessor,
 )
+from test_shared.test_utils.database_utils import default_db_url
 
 
 class TestBaseAnalyticsProcessor(unittest.TestCase):
-    @patch("processors.base_analytics_processor.Database")
     @patch("processors.base_analytics_processor.storage.Client")
-    def setUp(self, mock_storage_client, _):
+    def setUp(self, mock_storage_client):
         self.mock_storage_client = mock_storage_client
         self.mock_bucket = MagicMock()
         self.mock_storage_client().bucket.return_value = self.mock_bucket
@@ -81,6 +82,7 @@ def test_update_analytics_files(self, mock_blob, mock_save_json):
         "processors.base_analytics_processor.BaseAnalyticsProcessor.update_analytics_files"
     )
     @patch("processors.base_analytics_processor.BaseAnalyticsProcessor.save_summary")
+    @patch.dict(os.environ, {"FEEDS_DATABASE_URL": default_db_url})
     def test_run(
         self,
         mock_save_summary,

functions-python/preprocessed_analytics/tests/test_gbfs_processor.py

@@ -1,15 +1,16 @@
+import os
 import unittest
 from unittest.mock import patch, MagicMock
 from datetime import datetime
 from processors.gbfs_analytics_processor import (
     GBFSAnalyticsProcessor,
 )
+from test_shared.test_utils.database_utils import default_db_url
 
 
 class TestGBFSAnalyticsProcessor(unittest.TestCase):
-    @patch("processors.base_analytics_processor.Database")
     @patch("processors.base_analytics_processor.storage.Client")
-    def setUp(self, mock_storage_client, _):
+    def setUp(self, mock_storage_client):
         self.mock_storage_client = mock_storage_client
         self.mock_bucket = MagicMock()
         self.mock_storage_client().bucket.return_value = self.mock_bucket
@@ -26,6 +27,7 @@ def setUp(self, mock_storage_client, _):
         "processors.gbfs_analytics_processor."
         "GBFSAnalyticsProcessor.update_analytics_files"
     )
+    @patch.dict(os.environ, {"FEEDS_DATABASE_URL": default_db_url})
     def test_run(
         self,
         mock_update_analytics_files,
@@ -60,8 +62,8 @@ def test_run(
 
         # Assert that process_feed_data was called twice (once for each feed-snapshot pair)
         self.assertEqual(mock_process_feed_data.call_count, 2)
-        mock_process_feed_data.assert_any_call(mock_feed1, mock_snapshot1, {})
-        mock_process_feed_data.assert_any_call(mock_feed2, mock_snapshot2, {})
+        mock_process_feed_data.assert_any_call(mock_feed1, mock_snapshot1)
+        mock_process_feed_data.assert_any_call(mock_feed2, mock_snapshot2)
 
         # Assert that save was called once
         mock_save.assert_called_once()
@@ -96,7 +98,7 @@ def test_process_feed_data(self):
         ]
 
         # Run process_feed_data
-        self.processor.process_feed_data(mock_feed, mock_snapshot, None)
+        self.processor.process_feed_data(mock_feed, mock_snapshot)
 
         # Assert the data was appended correctly
         self.assertEqual(len(self.processor.data), 1)

functions-python/preprocessed_analytics/tests/test_gtfs_processor.py

@@ -1,15 +1,16 @@
+import os
 import unittest
 from unittest.mock import patch, MagicMock
 from datetime import datetime
 from processors.gtfs_analytics_processor import (
     GTFSAnalyticsProcessor,
 )
+from test_shared.test_utils.database_utils import default_db_url
 
 
 class TestGTFSAnalyticsProcessor(unittest.TestCase):
-    @patch("processors.base_analytics_processor.Database")
     @patch("processors.base_analytics_processor.storage.Client")
-    def setUp(self, mock_storage_client, _):
+    def setUp(self, mock_storage_client):
         self.mock_storage_client = mock_storage_client
         self.mock_bucket = MagicMock()
         self.mock_storage_client().bucket.return_value = self.mock_bucket
@@ -26,6 +27,7 @@ def setUp(self, mock_storage_client, _):
         "processors.gtfs_analytics_processor.GTFSAnalyticsProcessor"
         ".update_analytics_files"
     )
+    @patch.dict(os.environ, {"FEEDS_DATABASE_URL": default_db_url})
     def test_run(
         self,
         mock_update_analytics_files,
@@ -59,8 +61,8 @@ def test_run(
 
         # Assert that process_feed_data was called twice (once for each feed-dataset pair)
         self.assertEqual(mock_process_feed_data.call_count, 2)
-        mock_process_feed_data.assert_any_call(mock_feed1, mock_dataset1, {})
-        mock_process_feed_data.assert_any_call(mock_feed2, mock_dataset2, {})
+        mock_process_feed_data.assert_any_call(mock_feed1, mock_dataset1)
+        mock_process_feed_data.assert_any_call(mock_feed2, mock_dataset2)
 
         # Assert that save was called once
         mock_save.assert_called_once()
@@ -98,7 +100,7 @@ def test_process_feed_data(self):
         mock_dataset.validation_reports[0].features = [MagicMock(name="feature1")]
 
         # Run process_feed_data
-        self.processor.process_feed_data(mock_feed, mock_dataset, {})
+        self.processor.process_feed_data(mock_feed, mock_dataset)
 
         # Assert the data was appended correctly
         self.assertEqual(len(self.processor.data), 1)