chore: improve logs in batch functions (#1201)

Author: davidgamez

api/src/shared/common/logging_utils.py

@@ -6,7 +6,7 @@ def get_env_logging_level():
     """
     Get the logging level from the environment via OS variable LOGGING_LEVEL. Returns INFO if not set.
     """
-    return os.getenv("LOGGING_LEVEL", "INFO")
+    return logging.getLevelName(os.getenv("LOGGING_LEVEL", "INFO"))
 
 
 class Logger:

functions-python/batch_datasets/src/main.py

@@ -26,12 +26,13 @@
 from google.cloud.pubsub_v1.futures import Future
 from sqlalchemy import or_
 from sqlalchemy.orm import Session
+
 from shared.database_gen.sqlacodegen_models import Gtfsfeed, Gtfsdataset
 from shared.dataset_service.main import BatchExecutionService, BatchExecution
 from shared.database.database import with_db_session
-from shared.helpers.logger import Logger
+from shared.helpers.logger import init_logger
 
-logging.basicConfig(level=logging.INFO)
+init_logger()
 pubsub_topic_name = os.getenv("PUBSUB_TOPIC_NAME")
 project_id = os.getenv("PROJECT_ID")
 
@@ -91,8 +92,6 @@ def get_non_deprecated_feeds(session: Session):
         limit = os.getenv("FEEDS_LIMIT")
         query = query.limit(10 if limit is None else int(limit))
     results = query.all()
-    logging.info(f"Retrieved {len(results)} feeds.")
-
     return results
 
 
@@ -109,7 +108,6 @@ def batch_datasets(request, db_session: Session):
     :param db_session: database session object
     :return: HTTP response object
     """
-    Logger.init_logger()
     try:
         feeds = get_non_deprecated_feeds(db_session)
     except Exception as error:
@@ -139,7 +137,7 @@ def batch_datasets(request, db_session: Session):
             "api_key_parameter_name": feed.api_key_parameter_name,
         }
         data_str = json.dumps(payload)
-        logging.info(f"Publishing {data_str} to {topic_path}.")
+        logging.debug(f"Publishing {data_str} to {topic_path}.")
         future = publish(publisher, topic_path, data_str.encode("utf-8"))
         future.add_done_callback(
             lambda _: publish_callback(future, feed.stable_id, topic_path)
@@ -151,4 +149,6 @@ def batch_datasets(request, db_session: Session):
             timestamp=timestamp,
         )
     )
-    return f"Publish completed. Published {len(feeds)} feeds to {pubsub_topic_name}."
+    message = f"Publish completed. Published {len(feeds)} feeds to {pubsub_topic_name}."
+    logging.info(message)
+    return message

functions-python/batch_process_dataset/main_local_debug.py

@@ -18,7 +18,7 @@
 # - python batch_process_dataset/main_local_debug.py
 
 from dotenv import load_dotenv
-from batch_process_dataset.src.main import process_dataset
+from main import process_dataset
 
 # Load environment variables from .env.local
 load_dotenv(dotenv_path=".env.local")

functions-python/batch_process_dataset/src/main.py

@@ -34,10 +34,12 @@
 from shared.database.database import with_db_session, refresh_materialized_view
 import logging
 
-from shared.helpers.logger import Logger
+from shared.helpers.logger import init_logger, get_logger
 from shared.helpers.utils import download_and_get_hash
 from sqlalchemy.orm import Session
 
+init_logger()
+
 
 @dataclass
 class DatasetFile:
@@ -63,6 +65,7 @@ def __init__(
         api_key_parameter_name,
         public_hosted_datasets_url,
     ):
+        self.logger = get_logger(DatasetProcessor.__name__, feed_stable_id)
         self.producer_url = producer_url
         self.bucket_name = bucket_name
         self.latest_hash = latest_hash
@@ -73,7 +76,7 @@ def __init__(
         self.api_key_parameter_name = api_key_parameter_name
         self.date = datetime.now().strftime("%Y%m%d%H%M")
         if self.authentication_type != 0:
-            logging.info(f"Getting feed credentials for feed {self.feed_stable_id}")
+            self.logger.info(f"Getting feed credentials for feed {self.feed_stable_id}")
             self.feed_credentials = self.get_feed_credentials(self.feed_stable_id)
             if self.feed_credentials is None:
                 raise Exception(
@@ -95,7 +98,9 @@ def get_feed_credentials(feed_stable_id) -> str | None:
             feeds_credentials = json.loads(os.getenv("FEEDS_CREDENTIALS", "{}"))
             return feeds_credentials.get(feed_stable_id, None)
         except Exception as e:
-            logging.error(f"Error getting feed credentials: {e}")
+            get_logger(DatasetProcessor.__name__, feed_stable_id).error(
+                f"Error getting feed credentials: {e}"
+            )
             return None
 
     @staticmethod
@@ -141,23 +146,25 @@ def upload_dataset(self) -> DatasetFile or None:
         :return: the file hash and the hosted url as a tuple or None if no upload is required
         """
         try:
-            logging.info(f"[{self.feed_stable_id}] - Accessing URL {self.producer_url}")
+            self.logger.info("Accessing URL %s", self.producer_url)
             temp_file_path = self.generate_temp_filename()
             file_sha256_hash, is_zip = self.download_content(temp_file_path)
             if not is_zip:
-                logging.error(
+                self.logger.error(
                     f"[{self.feed_stable_id}] The downloaded file from {self.producer_url} is not a valid ZIP file."
                 )
                 return None
 
-            logging.info(f"[{self.feed_stable_id}] File hash is {file_sha256_hash}.")
+            self.logger.info(
+                f"[{self.feed_stable_id}] File hash is {file_sha256_hash}."
+            )
 
             if self.latest_hash != file_sha256_hash:
-                logging.info(
+                self.logger.info(
                     f"[{self.feed_stable_id}] Dataset has changed (hash {self.latest_hash}"
                     f"-> {file_sha256_hash}). Uploading new version."
                 )
-                logging.info(
+                self.logger.info(
                     f"Creating file {self.feed_stable_id}/latest.zip in bucket {self.bucket_name}"
                 )
                 self.upload_file_to_storage(
@@ -170,7 +177,7 @@ def upload_dataset(self) -> DatasetFile or None:
                 dataset_full_path = (
                     f"{self.feed_stable_id}/{dataset_stable_id}/{dataset_stable_id}.zip"
                 )
-                logging.info(
+                self.logger.info(
                     f"Creating file: {dataset_full_path}"
                     f" in bucket {self.bucket_name}"
                 )
@@ -185,7 +192,7 @@ def upload_dataset(self) -> DatasetFile or None:
                     hosted_url=f"{self.public_hosted_datasets_url}/{dataset_full_path}",
                 )
 
-            logging.info(
+            self.logger.info(
                 f"[{self.feed_stable_id}] Datasets hash has not changed (hash {self.latest_hash} "
                 f"-> {file_sha256_hash}). Not uploading it."
             )
@@ -216,11 +223,11 @@ def create_dataset(self, dataset_file: DatasetFile, db_session: Session):
                 .one_or_none()
             )
             if not latest_dataset:
-                logging.info(
+                self.logger.info(
                     f"[{self.feed_stable_id}] No latest dataset found for feed."
                 )
 
-            logging.info(
+            self.logger.info(
                 f"[{self.feed_stable_id}] Creating new dataset for feed with stable id {dataset_file.stable_id}."
             )
             new_dataset = Gtfsdataset(
@@ -239,10 +246,10 @@ def create_dataset(self, dataset_file: DatasetFile, db_session: Session):
                 db_session.add(latest_dataset)
             db_session.add(new_dataset)
             db_session.commit()
-            logging.info(f"[{self.feed_stable_id}] Dataset created successfully.")
+            self.logger.info(f"[{self.feed_stable_id}] Dataset created successfully.")
 
             refresh_materialized_view(db_session, t_feedsearch.name)
-            logging.info(
+            self.logger.info(
                 f"[{self.feed_stable_id}] Materialized view refresh event triggered successfully."
             )
         except Exception as e:
@@ -256,7 +263,7 @@ def process(self) -> DatasetFile or None:
         dataset_file = self.upload_dataset()
 
         if dataset_file is None:
-            logging.info(f"[{self.feed_stable_id}] No database update required.")
+            self.logger.info(f"[{self.feed_stable_id}] No database update required.")
             return None
         self.create_dataset(dataset_file)
         return dataset_file
@@ -268,8 +275,8 @@ def record_trace(
     """
     Record the trace in the datastore
     """
-    logging.info(
-        f"[{stable_id}] Recording trace in execution: [{execution_id}] with status: [{status}]"
+    get_logger("record_trace", stable_id if stable_id else "UNKNOWN").info(
+        f"Recording trace in execution: [{execution_id}] with status: [{status}]"
     )
     trace = DatasetTrace(
         trace_id=None,
@@ -306,7 +313,6 @@ def process_dataset(cloud_event: CloudEvent):
         }
     }
     """
-    Logger.init_logger()
     logging.info("Function Started")
     stable_id = "UNKNOWN"
     bucket_name = os.getenv("DATASETS_BUCKET_NAME")
@@ -336,27 +342,27 @@ def process_dataset(cloud_event: CloudEvent):
         #  Extract  data from message
         data = base64.b64decode(cloud_event.data["message"]["data"]).decode()
         json_payload = json.loads(data)
-        logging.info(
-            f"[{json_payload['feed_stable_id']}] JSON Payload: {json.dumps(json_payload)}"
-        )
         stable_id = json_payload["feed_stable_id"]
+        logger = get_logger("process_dataset", stable_id)
+        logger.info(f"JSON Payload: {json.dumps(json_payload)}")
+
         execution_id = json_payload["execution_id"]
         trace_service = DatasetTraceService()
         trace = trace_service.get_by_execution_and_stable_ids(execution_id, stable_id)
-        logging.info(f"[{stable_id}] Dataset trace: {trace}")
+        logger.info(f"Dataset trace: {trace}")
         executions = len(trace) if trace else 0
-        logging.info(
-            f"[{stable_id}] Dataset executed times={executions}/{maximum_executions} "
+        logger.info(
+            f"Dataset executed times={executions}/{maximum_executions} "
             f"in execution=[{execution_id}] "
         )
 
         if executions > 0:
             if executions >= maximum_executions:
                 error_message = (
-                    f"[{stable_id}] Function already executed maximum times "
+                    f"Function already executed maximum times "
                     f"in execution: [{execution_id}]"
                 )
-                logging.error(error_message)
+                logger.error(error_message)
                 return error_message
 
         processor = DatasetProcessor(
@@ -372,11 +378,14 @@ def process_dataset(cloud_event: CloudEvent):
         )
         dataset_file = processor.process()
     except Exception as e:
-        logging.error(e)
-        error_message = f"[{stable_id}] Error execution: [{execution_id}] error: [{e}]"
-        logging.error(error_message)
-        logging.error(f"Function completed with error:{error_message}")
+        # This makes sure the logger is initialized
+        logger = get_logger("process_dataset", stable_id if stable_id else "UNKNOWN")
+        logger.error(e)
+        error_message = f"Error execution: [{execution_id}] error: [{e}]"
+        logger.error(error_message)
+        logger.error(f"Function completed with error:{error_message}")
     finally:
+        logger = get_logger("process_dataset", stable_id if stable_id else "UNKNOWN")
         if stable_id and execution_id:
             status = (
                 Status.PUBLISHED if dataset_file is not None else Status.NOT_PUBLISHED
@@ -392,12 +401,12 @@ def process_dataset(cloud_event: CloudEvent):
                 trace_service,
             )
         else:
-            logging.error(
+            logger.error(
                 f"Function completed with errors, missing stable={stable_id} or execution_id={execution_id}"
             )
             return f"Function completed with errors, missing stable={stable_id} or execution_id={execution_id}"
-    logging.info(
-        f"[{stable_id}] Function %s in execution: [{execution_id}]",
+    logger.info(
+        f"Function %s in execution: [{execution_id}]",
         "successfully completed" if not error_message else "Failed",
     )
     return "Completed." if error_message is None else error_message

functions-python/batch_process_dataset/tests/test_batch_process_dataset_main.py

@@ -356,11 +356,10 @@ def test_process_no_change(self):
         self.assertIsNone(result)
         processor.create_dataset.assert_not_called()
 
-    @patch("main.Logger")
     @patch("main.DatasetTraceService")
     @patch("main.DatasetProcessor")
     def test_process_dataset_normal_execution(
-        self, mock_dataset_processor, mock_dataset_trace, _
+        self, mock_dataset_processor, mock_dataset_trace
     ):
         db_url = os.getenv("TEST_FEEDS_DATABASE_URL", default=default_db_url)
         os.environ["FEEDS_DATABASE_URL"] = db_url
@@ -392,11 +391,12 @@ def test_process_dataset_normal_execution(
         mock_dataset_processor.assert_called_once()
         mock_dataset_processor_instance.process.assert_called_once()
 
-    @patch("main.Logger")
     @patch("main.DatasetTraceService")
     @patch("main.DatasetProcessor")
     def test_process_dataset_exception_caught(
-        self, mock_dataset_processor, mock_dataset_trace, _
+        self,
+        mock_dataset_processor,
+        mock_dataset_trace,
     ):
         db_url = os.getenv("TEST_FEEDS_DATABASE_URL", default=default_db_url)
         os.environ["FEEDS_DATABASE_URL"] = db_url
@@ -416,9 +416,8 @@ def test_process_dataset_exception_caught(
         # Call the function
         process_dataset(cloud_event)
 
-    @patch("main.Logger")
     @patch("main.DatasetTraceService")
-    def test_process_dataset_missing_stable_id(self, mock_dataset_trace, _):
+    def test_process_dataset_missing_stable_id(self, mock_dataset_trace):
         db_url = os.getenv("TEST_FEEDS_DATABASE_URL", default=default_db_url)
         os.environ["FEEDS_DATABASE_URL"] = db_url
 

functions-python/big_query_ingestion/function_config.json

@@ -6,7 +6,7 @@
   "memory": "8Gi",
   "trigger_http": false,
   "include_folders": ["helpers"],
-  "include_api_folders": ["database_gen"],
+  "include_api_folders": ["database_gen", "common"],
   "environment_variables": [],
   "secret_environment_variables": [
     {