From cfe9bb9e1e50f9ca3c61ed2894fa7f6c4068f3ce Mon Sep 17 00:00:00 2001
From: Georgi <georgi.ajaeiya@datadoghq.com>
Date: Tue, 11 Feb 2025 15:20:49 +0100
Subject: [PATCH] Address closed PRs - take 2 (#904)

- Update README.
- Enable retry when validating Datadog's API key.
- Add error logs for cloudfront source handler.
---
 aws/logs_monitoring/README.md                 |  8 +++----
 aws/logs_monitoring/lambda_function.py        | 24 +++++++++++++------
 .../steps/handlers/s3_handler.py              |  8 +++++++
 3 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/aws/logs_monitoring/README.md b/aws/logs_monitoring/README.md
index 417213896..6e6dd92e7 100644
--- a/aws/logs_monitoring/README.md
+++ b/aws/logs_monitoring/README.md
@@ -12,10 +12,10 @@ aliases:
 
 The Datadog Forwarder is an AWS Lambda function that ships logs from AWS to Datadog, specifically:
 
-- Forward CloudWatch, ELB, S3, CloudTrail, VPC, SNS, and CloudFront logs to Datadog.
-- Forward S3 events to Datadog.
-- Forward Kinesis data stream events to Datadog (only CloudWatch logs are supported).
-- Forward metrics, traces, and logs from AWS Lambda functions to Datadog. Datadog recommends you use the [Datadog Lambda Extension][1] to monitor your Lambda functions.
+- Forward CloudWatch and S3 logs.
+- Forward logs from SNS, and Kinesis events to Datadog.
+- Kinesis data stream events support CloudWatch logs only.
+- Forward metrics, traces, and logs from AWS Lambda functions to Datadog. Datadog recommends to use [Datadog Lambda Extension][1] to monitor Lambda functions.
 
 For Serverless customers using the Forwarder to forward metrics, traces, and logs from AWS Lambda logs to Datadog, you should [migrate to the Datadog Lambda Extension][3] to collect telemetry directly from the Lambda execution environments. The Forwarder is still available for use in Serverless Monitoring, but will not be updated to support the latest features.
 
diff --git a/aws/logs_monitoring/lambda_function.py b/aws/logs_monitoring/lambda_function.py
index a2849982f..ccb8f3c9e 100644
--- a/aws/logs_monitoring/lambda_function.py
+++ b/aws/logs_monitoring/lambda_function.py
@@ -9,6 +9,7 @@
 import logging
 import requests
 from hashlib import sha1
+
 from datadog_lambda.wrapper import datadog_lambda_wrapper
 from datadog import api
 from enhanced_lambda_metrics import parse_and_submit_enhanced_metrics
@@ -41,13 +42,22 @@
     )
 # Validate the API key
 logger.debug("Validating the Datadog API key")
-validation_res = requests.get(
-    "{}/api/v1/validate?api_key={}".format(DD_API_URL, DD_API_KEY),
-    verify=(not DD_SKIP_SSL_VALIDATION),
-    timeout=10,
-)
-if not validation_res.ok:
-    raise Exception("The API key is not valid.")
+
+with requests.Session() as s:
+    retries = requests.adapters.Retry(
+        total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504]
+    )
+
+    s.mount("http://", requests.adapters.HTTPAdapter(max_retries=retries))
+    s.mount("https://", requests.adapters.HTTPAdapter(max_retries=retries))
+
+    validation_res = s.get(
+        "{}/api/v1/validate?api_key={}".format(DD_API_URL, DD_API_KEY),
+        verify=(not DD_SKIP_SSL_VALIDATION),
+        timeout=10,
+    )
+    if not validation_res.ok:
+        raise Exception("The API key is not valid.")
 
 # Force the layer to use the exact same API key and host as the forwarder
 api._api_key = DD_API_KEY
diff --git a/aws/logs_monitoring/steps/handlers/s3_handler.py b/aws/logs_monitoring/steps/handlers/s3_handler.py
index f4a4cb8f2..6a369cfd2 100644
--- a/aws/logs_monitoring/steps/handlers/s3_handler.py
+++ b/aws/logs_monitoring/steps/handlers/s3_handler.py
@@ -160,6 +160,9 @@ def _handle_cloudfront_source(self):
         # 2. We extract the AWS account id from the lambda arn
         namesplit = self.data_store.key.split("/")
         if len(namesplit) == 0:
+            self.logger.error(
+                f"Invalid key {self.data_store.key}, handle cloudfront source failed"
+            )
             return None
 
         filename = namesplit[len(namesplit) - 1]
@@ -167,6 +170,9 @@ def _handle_cloudfront_source(self):
         filenamesplit = filename.split(".")
 
         if len(filenamesplit) <= 3:
+            self.logger.error(
+                f"Invalid filename {filename}, handle cloudfront source failed"
+            )
             return None
 
         distributionID = filenamesplit[len(filenamesplit) - 4].lower()
@@ -174,9 +180,11 @@ def _handle_cloudfront_source(self):
         arnsplit = arn.split(":")
 
         if len(arnsplit) != 7:
+            self.logger.error(f"Invalid ARN {arn}, handle cloudfront source failed")
             return None
 
         awsaccountID = arnsplit[4].lower()
+
         return "arn:aws:cloudfront::{}:distribution/{}".format(
             awsaccountID, distributionID
         )