enhance batch_+process_datasets logging

Author: davidgamez

functions-python/batch_datasets/src/main.py

@@ -152,5 +152,5 @@ def batch_datasets(request, db_session: Session):
         )
     )
     message = f"Publish completed. Published {len(feeds)} feeds to {pubsub_topic_name}."
-    logging.info(message, extra={"message_extra": message})
+    logging.info(message)
     return message

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,27 @@ 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 +179,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 +194,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 +225,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 +248,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 +265,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 +277,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 +315,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 +344,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 +380,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 +403,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/helpers/logger.py

@@ -13,12 +13,11 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-import os
-
-import google.cloud.logging
-from google.cloud.logging_v2 import Client
+import threading
 import logging
 
+from shared.common.logging_utils import get_env_logging_level
+
 
 class StableIdFilter(logging.Filter):
     """Add a stable_id to the log record"""
@@ -33,37 +32,26 @@ def filter(self, record):
         return True
 
 
-class Logger:
-    """
-    Util class for logging information, errors or warnings.
-    This class uses the Google Cloud Logging API enhancing the logs with extra request information.
-    """
+_logger_initialized = False
+lock = threading.Lock()
 
-    def __init__(self, name):
-        self.init_logger()
-        self.logger = self.init_logger().logger(name)
 
-    @staticmethod
-    def init_logger() -> Client | None:
-        """
-        Initializes the logger
-        """
-        if os.getenv("DEBUG", "False") == "True":
-            return None
-        try:
-            client = google.cloud.logging.Client()
-            client.get_default_handler()
-            client.setup_logging()
-            return client
-        except Exception as error:
-            # This might happen when the GCP authorization credentials are not available.
-            # Example, when running the tests locally
-            logging.error(f"Error initializing the logger: {error}")
-        return None
-
-    def get_logger(self) -> Client:
-        """
-        Get the GCP logger instance
-        :return: the logger instance
-        """
-        return self.logger
+def init_logger():
+    """
+    Initializes the logger
+    """
+    with lock:
+        global _logger_initialized
+        if _logger_initialized:
+            return
+        logging.basicConfig(level=get_env_logging_level())
+        _logger_initialized = True
+
+
+def get_logger(name: str, stable_id: str = None):
+    logger = logging.getLogger(name)
+    if stable_id and not any(
+        isinstance(handler, StableIdFilter) for handler in logger.handlers
+    ):
+        logger.addFilter(StableIdFilter(stable_id))
+    return logger