Feature/eli 369 dynamodb x ray tracing (#256)

* handled none headers from request

* x-ray tracing setup for dynamo, s3, firehose

* enable_xray_patching env variable for lambda

* sonar fixes

---------

Co-authored-by: karthikeyannhs <174426205+Karthikeyannhs@users.noreply.github.com>

Author: ayeshalshukri1-nhs

infrastructure/modules/lambda/lambda.tf

@@ -22,6 +22,7 @@ resource "aws_lambda_function" "eligibility_signposting_lambda" {
       KINESIS_AUDIT_STREAM_TO_S3 = var.kinesis_audit_stream_to_s3_name
       ENV                        = var.environment
       LOG_LEVEL                  = var.log_level
+      ENABLE_XRAY_PATCHING       = var.enable_xray_patching
     }
   }
 

infrastructure/modules/lambda/variables.tf

@@ -47,3 +47,8 @@ variable "log_level" {
   description = "log level"
   type        = string
 }
+
+variable "enable_xray_patching"{
+  description = "flag to enable xray tracing, which puts an entry for dynamodb, s3 and firehose in trace map"
+  type        = string
+}

infrastructure/stacks/api-layer/lambda.tf

@@ -24,5 +24,6 @@ module "eligibility_signposting_lambda_function" {
   eligibility_status_table_name   = module.eligibility_status_table.table_name
   kinesis_audit_stream_to_s3_name = module.eligibility_audit_firehose_delivery_stream.firehose_stream_name
   log_level                       = "INFO"
+  enable_xray_patching            = "true"
   stack_name                      = local.stack_name
 }

poetry.lock

@@ -35,6 +35,7 @@ python-json-logger = "^3.3.0"
 fhir-resources = "^8.0.0"
 python-dateutil = "^2.9.0"
 pyhamcrest = "^2.1.0"
+aws-xray-sdk = "2.14.0"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.11.13"

pyproject.toml

@@ -1,8 +1,10 @@
 import logging
+import os
 from typing import Any
 
 import wireup.integration.flask
 from asgiref.wsgi import WsgiToAsgi
+from aws_xray_sdk.core import patch_all
 from flask import Flask
 from mangum import Mangum
 from mangum.types import LambdaContext, LambdaEvent
@@ -11,10 +13,14 @@
 from eligibility_signposting_api.common.error_handler import handle_exception
 from eligibility_signposting_api.common.request_validator import validate_request_params
 from eligibility_signposting_api.config.config import config
-from eligibility_signposting_api.logging.logs_helper import log_request_ids
-from eligibility_signposting_api.logging.logs_manager import add_request_id_to_logs, init_logging
+from eligibility_signposting_api.logging.logs_helper import log_request_ids_from_headers
+from eligibility_signposting_api.logging.logs_manager import add_lambda_request_id_to_logger, init_logging
+from eligibility_signposting_api.logging.tracing_helper import tracing_setup
 from eligibility_signposting_api.views import eligibility_blueprint
 
+if os.getenv("ENABLE_XRAY_PATCHING"):
+    patch_all()
+
 init_logging()
 logger = logging.getLogger(__name__)
 
@@ -25,8 +31,9 @@ def main() -> None:  # pragma: no cover
     app.run(debug=config()["log_level"] == logging.DEBUG)
 
 
-@add_request_id_to_logs()
-@log_request_ids()
+@add_lambda_request_id_to_logger()
+@tracing_setup()
+@log_request_ids_from_headers()
 @validate_request_params()
 def lambda_handler(event: LambdaEvent, context: LambdaContext) -> dict[str, Any]:  # pragma: no cover
     """Run the Flask app as an AWS Lambda."""

src/eligibility_signposting_api/app.py

@@ -22,6 +22,7 @@ def config() -> dict[str, Any]:
     rules_bucket_name = BucketName(os.getenv("RULES_BUCKET_NAME", "test-rules-bucket"))
     audit_bucket_name = BucketName(os.getenv("AUDIT_BUCKET_NAME", "test-audit-bucket"))
     aws_default_region = AwsRegion(os.getenv("AWS_DEFAULT_REGION", "eu-west-1"))
+    enable_xray_patching = bool(os.getenv("ENABLE_XRAY_PATCHING", "false"))
     kinesis_audit_stream_to_s3 = AwsKinesisFirehoseStreamName(
         os.getenv("KINESIS_AUDIT_STREAM_TO_S3", "test_kinesis_audit_stream_to_s3")
     )
@@ -39,19 +40,22 @@ def config() -> dict[str, Any]:
             "audit_bucket_name": audit_bucket_name,
             "firehose_endpoint": None,
             "kinesis_audit_stream_to_s3": kinesis_audit_stream_to_s3,
+            "enable_xray_patching": enable_xray_patching,
             "log_level": log_level,
         }
 
+    local_stack_endpoint = "http://localhost:4566"
     return {
         "aws_access_key_id": AwsAccessKey(os.getenv("AWS_ACCESS_KEY_ID", "dummy_key")),
         "aws_default_region": aws_default_region,
         "aws_secret_access_key": AwsSecretAccessKey(os.getenv("AWS_SECRET_ACCESS_KEY", "dummy_secret")),
-        "dynamodb_endpoint": URL(os.getenv("DYNAMODB_ENDPOINT", "http://localhost:4566")),
+        "dynamodb_endpoint": URL(os.getenv("DYNAMODB_ENDPOINT", local_stack_endpoint)),
         "person_table_name": person_table_name,
-        "s3_endpoint": URL(os.getenv("S3_ENDPOINT", "http://localhost:4566")),
+        "s3_endpoint": URL(os.getenv("S3_ENDPOINT", local_stack_endpoint)),
         "rules_bucket_name": rules_bucket_name,
         "audit_bucket_name": audit_bucket_name,
-        "firehose_endpoint": URL(os.getenv("FIREHOSE_ENDPOINT", "http://localhost:4566")),
+        "firehose_endpoint": URL(os.getenv("FIREHOSE_ENDPOINT", local_stack_endpoint)),
         "kinesis_audit_stream_to_s3": kinesis_audit_stream_to_s3,
+        "enable_xray_patching": enable_xray_patching,
         "log_level": log_level,
     }

src/eligibility_signposting_api/config/config.py

@@ -8,12 +8,12 @@
 logger = logging.getLogger(__name__)
 
 
-def log_request_ids() -> Callable:
+def log_request_ids_from_headers() -> Callable:
     def decorator(func: Callable) -> Callable:
         @wraps(func)
         def wrapper(event: LambdaEvent, context: LambdaContext) -> dict[str, Any] | None:
-            gateway_request_id = event.get("requestContext", {}).get("requestId")
-            headers = event.get("headers", {})
+            gateway_request_id = (event.get("requestContext") or {}).get("requestId")
+            headers = event.get("headers") or {}
             logger.info(
                 "request trace metadata",
                 extra={

src/eligibility_signposting_api/logging/logs_helper.py

@@ -14,7 +14,7 @@
 LOG_FORMAT = "%(asctime)s %(levelname)-8s %(name)s %(module)s.py:%(funcName)s():%(lineno)d %(message)s"
 
 
-def add_request_id_to_logs() -> Callable:
+def add_lambda_request_id_to_logger() -> Callable:
     def decorator(func: Callable) -> Callable:
         @wraps(func)
         def wrapper(event: LambdaEvent, context: LambdaContext) -> dict[str, Any] | None:

src/eligibility_signposting_api/logging/logs_manager.py

@@ -0,0 +1,21 @@
+from collections.abc import Callable
+from functools import wraps
+from typing import Any
+
+from aws_xray_sdk.core import xray_recorder
+from mangum.types import LambdaContext, LambdaEvent
+
+
+def tracing_setup() -> Callable:
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(event: LambdaEvent, context: LambdaContext) -> dict[str, Any] | None:
+            xray_recorder.begin_subsegment("Lambda")
+            try:
+                return func(event, context)
+            finally:
+                xray_recorder.end_subsegment()
+
+        return wrapper
+
+    return decorator