diff --git a/src/lambda_obs/handler.py b/src/lambda_obs/handler.py index 7586ce8..0511168 100644 --- a/src/lambda_obs/handler.py +++ b/src/lambda_obs/handler.py @@ -2,6 +2,7 @@ from __future__ import annotations +import time from typing import Any from .logging import Logger @@ -11,21 +12,38 @@ metrics = Metrics(namespace="LambdaObsStarter", service="lambda-obs-starter") +def _request_id(event: dict[str, Any], context: Any) -> str | None: + if context is not None and getattr(context, "aws_request_id", None): + return str(context.aws_request_id) + headers = (event or {}).get("headers") or {} + for key in ("x-correlation-id", "X-Correlation-Id", "x-request-id"): + if key in headers and headers[key]: + return str(headers[key]) + return (event or {}).get("correlation_id") + + def handler(event: dict[str, Any], context: Any = None) -> dict[str, Any]: """Handle a simple greeting event. Expected event shape: ``{"name": "world"}``. + Optional correlation via ``correlation_id`` or ``headers.x-correlation-id``. """ + started = time.perf_counter() name = (event or {}).get("name") or "world" - bound = logger.bind(function="handler") + cid = _request_id(event or {}, context) + bound = logger.with_correlation_id(cid).bind(function="handler") + if cid: + metrics.add_dimension("correlation_present", "true") bound.info("invocation started", name=name) metrics.add_metric("Invocations", 1, unit="Count") try: greeting = f"hello {name}" metrics.add_metric("Success", 1, unit="Count") - bound.info("invocation succeeded", greeting=greeting) - return {"ok": True, "message": greeting} + elapsed_ms = (time.perf_counter() - started) * 1000.0 + metrics.timing("HandlerDuration", elapsed_ms) + bound.info("invocation succeeded", greeting=greeting, duration_ms=round(elapsed_ms, 3)) + return {"ok": True, "message": greeting, "correlation_id": bound._base.get("correlation_id")} except Exception as exc: # pragma: no cover - defensive metrics.add_metric("Errors", 1, unit="Count") bound.error("invocation failed", error=str(exc)) diff --git a/tests/test_handler.py b/tests/test_handler.py index b4ece24..176511f 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + from lambda_obs.handler import handler @@ -5,8 +7,28 @@ def test_handler_greets_name(): result = handler({"name": "lei"}, None) assert result["ok"] is True assert result["message"] == "hello lei" + assert "correlation_id" in result def test_handler_default_name(): result = handler({}, None) assert result["message"] == "hello world" + + +def test_handler_uses_event_correlation_id(): + result = handler({"name": "lei", "correlation_id": "abc-999"}, None) + assert result["correlation_id"] == "abc-999" + + +def test_handler_uses_header_correlation_id(): + result = handler( + {"name": "lei", "headers": {"x-correlation-id": "hdr-1"}}, + None, + ) + assert result["correlation_id"] == "hdr-1" + + +def test_handler_uses_context_request_id(): + ctx = SimpleNamespace(aws_request_id="req-42") + result = handler({"name": "lei"}, ctx) + assert result["correlation_id"] == "req-42"