feat: add tenant id to structed log messages

Author: Ramisa Alam

awslambdaric/bootstrap.py

@@ -341,6 +341,7 @@ def emit(self, record):
 class LambdaLoggerFilter(logging.Filter):
     def filter(self, record):
         record.aws_request_id = _GLOBAL_AWS_REQUEST_ID or ""
+        record.tenant_id = _GLOBAL_TENANT_ID
         return True
 
 
@@ -449,6 +450,7 @@ def create_log_sink():
 
 
 _GLOBAL_AWS_REQUEST_ID = None
+_GLOBAL_TENANT_ID = None
 
 
 def _setup_logging(log_format, log_level, log_sink):
@@ -494,7 +496,7 @@ def run(app_root, handler, lambda_runtime_api_addr):
 
         try:
             _setup_logging(_AWS_LAMBDA_LOG_FORMAT, _AWS_LAMBDA_LOG_LEVEL, log_sink)
-            global _GLOBAL_AWS_REQUEST_ID
+            global _GLOBAL_AWS_REQUEST_ID, _GLOBAL_TENANT_ID
 
             request_handler = _get_handler(handler)
         except FaultException as e:
@@ -519,6 +521,7 @@ def run(app_root, handler, lambda_runtime_api_addr):
             event_request = lambda_runtime_client.wait_next_invocation()
 
             _GLOBAL_AWS_REQUEST_ID = event_request.invoke_id
+            _GLOBAL_TENANT_ID = event_request.tenant_id
 
             update_xray_env_variable(event_request.x_amzn_trace_id)
 

awslambdaric/lambda_runtime_log_utils.py

@@ -30,6 +30,7 @@
     "processName",
     "process",
     "aws_request_id",
+    "tenant_id",
     "_frame_type",
 }
 
@@ -124,6 +125,9 @@ def format(self, record: logging.LogRecord) -> str:
             "requestId": getattr(record, "aws_request_id", None),
             "location": self.__format_location(record),
         }
+        if hasattr(record, "tenant_id") and record.tenant_id is not None:
+            result["tenantId"] = record.tenant_id
+
         result.update(
             (key, value)
             for key, value in record.__dict__.items()

tests/test_bootstrap.py

@@ -1376,6 +1376,31 @@ def test_json_formatter(self, mock_stderr):
                     )
         self.assertEqual(mock_stderr.getvalue(), "")
 
+    @patch("awslambdaric.bootstrap._GLOBAL_TENANT_ID", "test-tenant-id")
+    @patch("sys.stderr", new_callable=StringIO)
+    def test_json_formatter_with_tenant_id(self, mock_stderr):
+        logger = logging.getLogger("a.b")
+        level = logging.INFO
+        message = "Test json formatting with tenant id"
+        expected = {
+            "level": "INFO",
+            "logger": "a.b",
+            "message": message,
+            "requestId": "",
+            "tenantId": "test-tenant-id",
+        }
+
+        with patch("sys.stdout", new_callable=StringIO) as mock_stdout:
+            logger.log(level, message)
+
+            data = json.loads(mock_stdout.getvalue())
+            data.pop("timestamp")
+            self.assertEqual(
+                data,
+                expected,
+            )
+        self.assertEqual(mock_stderr.getvalue(), "")
+
     @patch("sys.stdout", new_callable=StringIO)
     @patch("sys.stderr", new_callable=StringIO)
     def test_exception(self, mock_stderr, mock_stdout):