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,59 @@ 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")
+
+        test_cases = [
+            (
+                logging.ERROR,
+                "TEST 1",
+                {
+                    "level": "ERROR",
+                    "logger": "a.b",
+                    "message": "TEST 1",
+                    "requestId": "",
+                    "tenantId": "test-tenant-id",
+                },
+            ),
+            (
+                logging.ERROR,
+                "test \nwith \nnew \nlines",
+                {
+                    "level": "ERROR",
+                    "logger": "a.b",
+                    "message": "test \nwith \nnew \nlines",
+                    "requestId": "",
+                    "tenantId": "test-tenant-id",
+                },
+            ),
+            (
+                logging.CRITICAL,
+                "TEST CRITICAL",
+                {
+                    "level": "CRITICAL",
+                    "logger": "a.b",
+                    "message": "TEST CRITICAL",
+                    "requestId": "",
+                    "tenantId": "test-tenant-id",
+                },
+            ),
+        ]
+        for level, msg, expected in test_cases:
+            with self.subTest(msg):
+                with patch("sys.stdout", new_callable=StringIO) as mock_stdout:
+                    logger.log(level, msg)
+
+                    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):