Merge pull request #598 from NHSDigital/feature/ELI-686-adding-subsegments

[ELI-686] - Adding custom subsegments for xray

Author: TOEL2

src/eligibility_signposting_api/audit/audit_service.py

@@ -2,6 +2,7 @@
 import logging
 from typing import Annotated
 
+from aws_xray_sdk.core import xray_recorder
 from botocore.client import BaseClient
 from wireup import Inject, service
 
@@ -21,6 +22,7 @@ def __init__(
         self.firehose = firehose
         self.audit_delivery_stream = audit_delivery_stream
 
+    @xray_recorder.capture("AuditService.audit")  # pyright: ignore[reportCallIssue]
     def audit(self, audit_record: dict) -> None:
         """
         Sends an audit record to the configured Firehose delivery stream.

src/eligibility_signposting_api/repos/campaign_repo.py

@@ -2,6 +2,7 @@
 from collections.abc import Generator
 from typing import Annotated, NewType
 
+from aws_xray_sdk.core import xray_recorder
 from botocore.client import BaseClient
 from wireup import Inject, service
 
@@ -25,9 +26,12 @@ def __init__(
         self.s3_client = s3_client
         self.bucket_name = bucket_name
 
+    @xray_recorder.capture("CampaignRepo.get_campaign_configs")  # pyright: ignore[reportCallIssue]
     def get_campaign_configs(self) -> Generator[CampaignConfig]:
-        campaign_objects = self.s3_client.list_objects(Bucket=self.bucket_name)
-        for campaign_object in campaign_objects["Contents"]:
-            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=f"{campaign_object['Key']}")
-            body = response["Body"].read()
-            yield Rules.model_validate(json.loads(body)).campaign_config
+        with xray_recorder.in_subsegment("CampaignRepo.get_campaign_configs:list_objects"):
+            campaign_objects = self.s3_client.list_objects(Bucket=self.bucket_name)
+        with xray_recorder.in_subsegment("CampaignRepo.get_campaign_configs:get_objects"):
+            for campaign_object in campaign_objects["Contents"]:
+                response = self.s3_client.get_object(Bucket=self.bucket_name, Key=f"{campaign_object['Key']}")
+                body = response["Body"].read()
+                yield Rules.model_validate(json.loads(body)).campaign_config

src/eligibility_signposting_api/repos/consumer_mapping_repo.py

@@ -2,6 +2,7 @@
 import logging
 from typing import Annotated, NewType
 
+from aws_xray_sdk.core import xray_recorder
 from botocore.client import BaseClient
 from botocore.exceptions import ClientError
 from wireup import Inject, service
@@ -28,6 +29,7 @@ def __init__(
         self.s3_client = s3_client
         self.bucket_name = bucket_name
 
+    @xray_recorder.capture("ConsumerMappingRepo.get_permitted_campaign_ids")  # pyright: ignore[reportCallIssue]
     def get_permitted_campaign_ids(self, consumer_id: ConsumerId) -> list[CampaignID] | None:
         try:
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=CONSUMER_MAPPING_FILE_NAME)

src/eligibility_signposting_api/repos/person_repo.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Annotated, Any, NewType
 
+from aws_xray_sdk.core import xray_recorder
 from boto3.dynamodb.conditions import Key
 from boto3.resources.base import ServiceResource
 from wireup import Inject, service
@@ -54,6 +55,7 @@ def get_person_record(self, nhs_hash: str | None) -> Any:
 
         return None
 
+    @xray_recorder.capture("PersonRepo.get_eligibility_data")  # pyright: ignore[reportCallIssue]
     def get_eligibility_data(self, nhs_number: NHSNumber) -> Person:
         # Hash using AWSCURRENT secret and fetch items
         items = None

src/eligibility_signposting_api/repos/secret_repo.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Annotated, NewType
 
+from aws_xray_sdk.core import xray_recorder
 from botocore.client import BaseClient
 from botocore.exceptions import ClientError
 from wireup import Inject, service
@@ -16,6 +17,7 @@ def __init__(self, secret_manager: Annotated[BaseClient, Inject(qualifier="secre
         super().__init__()
         self.secret_manager = secret_manager
 
+    @xray_recorder.capture("SecretRepo._get_secret_by_stage")  # pyright: ignore[reportCallIssue]
     def _get_secret_by_stage(self, secret_name: str, stage: str) -> dict[str, str]:
         """Internal helper to fetch a secret by version stage."""
         try: