Add Custom logger

Author: davidgamez

api/src/feeds/impl/metadata_api_impl.py

@@ -4,7 +4,9 @@
 
 from feeds_gen.apis.metadata_api_base import BaseMetadataApi
 from feeds_gen.models.metadata import Metadata
+from shared.common.logging_utils import Logger
 
+Logger.init_logger()
 
 class MetadataApiImpl(BaseMetadataApi):
     """
@@ -37,7 +39,7 @@ def get_metadata(
             # Access the values using the get() method
             long_commit_hash = config.get("DEFAULT", "LONG_COMMIT_HASH")
             version = config.get("DEFAULT", "EXTRACTED_VERSION")
-
+            logging.info(f"Retrieved metadata for {version}")
         except Exception as e:
             logging.error(f"Cannot read {file} file from directory {current_directory}: \n {e}")
 

api/src/shared/common/logging_utils.py

@@ -6,6 +6,10 @@
 import google.cloud.logging
 from google.cloud.logging_v2 import Client
 
+from middleware.request_context import get_request_context
+from utils.config import get_config, PROJECT_ID
+
+
 def get_env_logging_level():
     """
     Get the logging level from the environment via OS variable LOGGING_LEVEL. Returns INFO if not set.
@@ -60,12 +64,12 @@ def init_logger():
             if hasattr(Logger, "initialized"):
                 return
             logging.basicConfig(level=get_env_logging_level())
-            if not is_local_env():
+            if is_local_env():
                 # Use the default logging handler
                 logging.info("Using default logging handler")
                 return
             try:
-                client = google.cloud.logging.Client()
+                client = google.cloud.logging_v2.Client()
                 client.get_default_handler()
                 client.setup_logging()
                 logging.info("GCP logging client initialized")
@@ -78,10 +82,69 @@ def init_logger():
     def get_logger(self):
         return self.logger
 
+
+class TraceLogger:
+    def __init__(self, logger: logging.Logger, project_id: str, trace_id: str, span_id: str):
+        self._logger = logger
+        self._trace_id = None
+        self._span_id = None
+        self._project_id = project_id
+        self._trace_id = trace_id
+        self._span_id = span_id
+
+    def _inject_trace(self, extra):
+        if not self._trace_id or not self._span_id:
+            return extra or {}
+
+        trace_fields = {
+            "logging.googleapis.com/trace": f"projects/{self._project_id}/traces/{self._trace_id}",
+            "logging.googleapis.com/spanId": self._span_id,
+            "logging.googleapis.com/trace_sampled": True,
+        }
+        if extra:
+            trace_fields.update(extra)
+        return trace_fields
+
+    def info(self, msg, *args, extra=None, **kwargs):
+        return self._logger.info(msg, *args, extra=self._inject_trace(extra), **kwargs)
+
+    def error(self, msg, *args, extra=None, **kwargs):
+        return self._logger.error(msg, *args, extra=self._inject_trace(extra), **kwargs)
+
+    def warning(self, msg, *args, extra=None, **kwargs):
+        return self._logger.warning(msg, *args, extra=self._inject_trace(extra), **kwargs)
+
+    def debug(self, msg, *args, extra=None, **kwargs):
+        return self._logger.debug(msg, *args, extra=self._inject_trace(extra), **kwargs)
+
+    def exception(self, msg, *args, extra=None, **kwargs):
+        return self._logger.exception(msg, *args, extra=self._inject_trace(extra), **kwargs)
+
+
 def new_logger(name: str):
     """
     Create a new logger with the given name.
     """
-    Logger.init_logger()
-    return logging.getLogger(name)
+    logger = logging.getLogger(name)
+    logger.setLevel(get_env_logging_level())
+    if not logger.handlers:
+        if is_local_env():
+            handler = logging.StreamHandler()
+        else:
+            from google.cloud.logging import Client
+            from google.cloud.logging.handlers import CloudLoggingHandler
+            client = Client()
+            handler = CloudLoggingHandler(client)
+
+        formatter = logging.Formatter(
+            '%(asctime)s %(levelname)s %(name)s %(message)s'
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+        request_context = get_request_context()
+        if not request_context:
+            return logger
+        return TraceLogger(logger, get_config(PROJECT_ID), request_context.get("trace_id"), request_context.get("span_id"))
+