feat: adding AWS_LAMBDA_LOG_TIMESTAMP_PRECISION env

This is done to support millisecond precision timestamp in the
structured logs in cloudwatch as requested in #156.
Changing directly to millisecond is a potentially breaking change, so we
agreed on supporting this environmental variable on this version of the
RIC. We'll remove the support in RIC v5, defaulting to milisecond.

Author: darklight3it

awslambdaric/bootstrap.py

@@ -14,10 +14,10 @@
 from .lambda_runtime_client import LambdaRuntimeClient
 from .lambda_runtime_exception import FaultException
 from .lambda_runtime_log_utils import (
-    _DATETIME_FORMAT,
     _DEFAULT_FRAME_TYPE,
     _JSON_FRAME_TYPES,
     _TEXT_FRAME_TYPES,
+    _format_timestamp,
     JsonFormatter,
     LogFormat,
     _format_log_level,
@@ -106,9 +106,7 @@ def replace_line_indentation(line, indent_char, new_indent_char):
 
     def log_error(error_result, log_sink):
         error_result = {
-            "timestamp": time.strftime(
-                _DATETIME_FORMAT, logging.Formatter.converter(time.time())
-            ),
+            "timestamp": _format_timestamp(time.time()),
             "log_level": "ERROR",
             **error_result,
         }

awslambdaric/lambda_runtime_log_utils.py

@@ -5,10 +5,34 @@
 
 import json
 import logging
+import os
 import traceback
+from datetime import datetime, timezone
 from enum import IntEnum
 
-_DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
+# TODO(v5.0): Remove this env var check and make milliseconds the default.
+# Once removed, _format_timestamp simplifies to just the millis branch.
+_TIMESTAMP_PRECISION_MILLIS = (
+    os.environ.get("AWS_LAMBDA_LOG_TIMESTAMP_PRECISION", "").lower() == "milliseconds"
+)
+
+
+def _format_timestamp(epoch_secs):
+    """Format a UTC timestamp in ISO 8601 format.
+
+    Returns second precision by default (e.g. 2024-06-19T23:13:05Z), to avoid breaking changes.
+    When AWS_LAMBDA_LOG_TIMESTAMP_PRECISION=milliseconds, includes
+    milliseconds (e.g. 2024-06-19T23:13:05.068Z).
+    """
+    dt = datetime.fromtimestamp(epoch_secs, tz=timezone.utc)
+
+    if _TIMESTAMP_PRECISION_MILLIS:
+        millis = dt.microsecond // 1000
+        return f"{dt:%Y-%m-%dT%H:%M:%S}.{millis:03d}Z"
+
+    return f"{dt:%Y-%m-%dT%H:%M:%S}Z"
+
+
 _RESERVED_FIELDS = {
     "name",
     "msg",
@@ -79,7 +103,10 @@ def _format_log_level(record: logging.LogRecord) -> int:
 
 class JsonFormatter(logging.Formatter):
     def __init__(self):
-        super().__init__(datefmt=_DATETIME_FORMAT)
+        super().__init__()
+
+    def formatTime(self, record, datefmt=None):
+        return _format_timestamp(record.created)
 
     @staticmethod
     def __format_stacktrace(exc_info):

tests/test_lambda_runtime_log_utils.py

@@ -3,20 +3,24 @@
 SPDX-License-Identifier: Apache-2.0
 """
 
+import importlib
 import logging
 import json
-import re
+import os
 import time
 import unittest
+from unittest.mock import patch
 
-from awslambdaric.lambda_runtime_log_utils import JsonFormatter
 
-
-class TestJsonFormatterTimestamp(unittest.TestCase):
+class TestJsonFormatterTimestampDefaultPrecision(unittest.TestCase):
     def setUp(self):
-        self.formatter = JsonFormatter()
-        self.logger = logging.getLogger("test")
-        self.logger.setLevel(logging.INFO)
+        env = os.environ.copy()
+        env.pop("AWS_LAMBDA_LOG_TIMESTAMP_PRECISION", None)
+        with patch.dict(os.environ, env, clear=True):
+            import awslambdaric.lambda_runtime_log_utils as mod
+
+            importlib.reload(mod)
+            self.formatter = mod.JsonFormatter()
 
     def test_timestamp_format_is_second_precision_with_z(self):
         record = logging.LogRecord(
@@ -53,9 +57,7 @@ def test_timestamp_value_is_accurate(self):
         output = self.formatter.format(record)
         log_entry = json.loads(output)
 
-        expected = time.strftime(
-            "%Y-%m-%dT%H:%M:%SZ", self.formatter.converter(record.created)
-        )
+        expected = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(record.created))
         self.assertEqual(log_entry["timestamp"], expected)
 
     def test_timestamp_does_not_include_milliseconds(self):
@@ -102,3 +104,114 @@ def test_timestamps_same_within_same_second(self):
         output2 = json.loads(self.formatter.format(record2))
 
         self.assertEqual(output1["timestamp"], output2["timestamp"])
+
+
+class TestJsonFormatterTimestampMillisecondPrecision(unittest.TestCase):
+    def setUp(self):
+        env = os.environ.copy()
+        env["AWS_LAMBDA_LOG_TIMESTAMP_PRECISION"] = "milliseconds"
+        with patch.dict(os.environ, env, clear=True):
+            import awslambdaric.lambda_runtime_log_utils as mod
+
+            importlib.reload(mod)
+            self.formatter = mod.JsonFormatter()
+
+    def test_timestamp_includes_milliseconds(self):
+        record = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="hello",
+            args=None,
+            exc_info=None,
+        )
+        output = self.formatter.format(record)
+        log_entry = json.loads(output)
+        timestamp = log_entry["timestamp"]
+
+        pattern = r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{3}Z$"
+        self.assertRegex(
+            timestamp,
+            pattern,
+            f"Timestamp '{timestamp}' does not match expected format YYYY-MM-DDTHH:MM:SS.mmmZ",
+        )
+
+    def test_timestamp_milliseconds_are_accurate(self):
+        record = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="hello",
+            args=None,
+            exc_info=None,
+        )
+        record.created = 1718838785.068
+        output = self.formatter.format(record)
+        log_entry = json.loads(output)
+
+        self.assertEqual(log_entry["timestamp"], "2024-06-19T23:13:05.068Z")
+
+    def test_timestamp_zero_milliseconds(self):
+        record = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="hello",
+            args=None,
+            exc_info=None,
+        )
+        record.created = 1718838785.0
+        output = self.formatter.format(record)
+        log_entry = json.loads(output)
+
+        self.assertEqual(log_entry["timestamp"], "2024-06-19T23:13:05.000Z")
+
+    def test_timestamps_differ_within_same_second(self):
+        record1 = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="first",
+            args=None,
+            exc_info=None,
+        )
+        record1.created = 1718838785.100
+
+        record2 = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="second",
+            args=None,
+            exc_info=None,
+        )
+        record2.created = 1718838785.200
+
+        output1 = json.loads(self.formatter.format(record1))
+        output2 = json.loads(self.formatter.format(record2))
+
+        self.assertNotEqual(output1["timestamp"], output2["timestamp"])
+        self.assertEqual(output1["timestamp"], "2024-06-19T23:13:05.100Z")
+        self.assertEqual(output2["timestamp"], "2024-06-19T23:13:05.200Z")
+
+
+class TestTimestampPrecisionVersionGate(unittest.TestCase):
+    """Fails if we bump to v5+ without removing the seconds-precision path."""
+
+    def test_v5_must_remove_timestamp_precision_env_var(self):
+        from awslambdaric import __version__
+
+        major = int(__version__.split(".")[0])
+        if major >= 5:
+            import awslambdaric.lambda_runtime_log_utils as mod
+
+            self.assertFalse(
+                hasattr(mod, "_TIMESTAMP_PRECISION_MILLIS"),
+                "v5+: remove _TIMESTAMP_PRECISION_MILLIS and make milliseconds "
+                "the default. See TODO(v5.0) in lambda_runtime_log_utils.py",
+            )