Lazy load boto client when using datadogpy for metrics. (#558)

* Lazy load boto client when using datadogpy for metrics.

* Update test mocking.

Author: purple4reina

datadog_lambda/api.py

@@ -57,8 +57,6 @@ def get_api_key() -> str:
     if api_key:
         return api_key
 
-    import boto3
-
     DD_API_KEY_SECRET_ARN = os.environ.get("DD_API_KEY_SECRET_ARN", "")
     DD_API_KEY_SSM_NAME = os.environ.get("DD_API_KEY_SSM_NAME", "")
     DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "")
@@ -85,7 +83,7 @@ def get_api_key() -> str:
             if is_gov_region
             else None
         )
-        secrets_manager_client = boto3.client(
+        secrets_manager_client = _boto3_client(
             "secretsmanager", endpoint_url=endpoint_url, region_name=secrets_region
         )
         api_key = secrets_manager_client.get_secret_value(
@@ -96,7 +94,7 @@ def get_api_key() -> str:
         fips_endpoint = (
             f"https://ssm-fips.{LAMBDA_REGION}.amazonaws.com" if is_gov_region else None
         )
-        ssm_client = boto3.client("ssm", endpoint_url=fips_endpoint)
+        ssm_client = _boto3_client("ssm", endpoint_url=fips_endpoint)
         api_key = ssm_client.get_parameter(
             Name=DD_API_KEY_SSM_NAME, WithDecryption=True
         )["Parameter"]["Value"]
@@ -105,7 +103,7 @@ def get_api_key() -> str:
         fips_endpoint = (
             f"https://kms-fips.{LAMBDA_REGION}.amazonaws.com" if is_gov_region else None
         )
-        kms_client = boto3.client("kms", endpoint_url=fips_endpoint)
+        kms_client = _boto3_client("kms", endpoint_url=fips_endpoint)
         api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY)
     else:
         api_key = DD_API_KEY
@@ -133,3 +131,9 @@ def init_api():
 
         # Unmute exceptions from datadog api client, so we can catch and handle them
         api._mute = False
+
+
+def _boto3_client(*args, **kwargs):
+    import botocore.session
+
+    return botocore.session.get_session().create_client(*args, **kwargs)

pyproject.toml

@@ -30,15 +30,15 @@ datadog = ">=0.51.0,<1.0.0"
 wrapt = "^1.11.2"
 ddtrace = ">=2.20.0,<4"
 ujson = ">=5.9.0"
-boto3 = { version = "^1.34.0", optional = true }
+botocore = { version = "^1.34.0", optional = true }
 requests = { version ="^2.22.0", optional = true }
 pytest = { version= "^8.0.0", optional = true }
 pytest-benchmark = { version = "^4.0", optional = true }
 flake8 = { version = "^5.0.4", optional = true }
 
 [tool.poetry.extras]
 dev = [
-    "boto3",
+    "botocore",
     "flake8",
     "pytest",
     "pytest-benchmark",

tests/test_api.py

@@ -22,7 +22,7 @@ def setUp(self):
         )
         self.env_patcher.start()
 
-    @patch("boto3.client")
+    @patch("botocore.session.Session.create_client")
     def test_secrets_manager_fips_endpoint(self, mock_boto3_client):
         mock_client = MagicMock()
         mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"}
@@ -42,7 +42,7 @@ def test_secrets_manager_fips_endpoint(self, mock_boto3_client):
         )
         self.assertEqual(api_key, "test-api-key")
 
-    @patch("boto3.client")
+    @patch("botocore.session.Session.create_client")
     def test_secrets_manager_different_region(self, mock_boto3_client):
         mock_client = MagicMock()
         mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"}
@@ -62,7 +62,7 @@ def test_secrets_manager_different_region(self, mock_boto3_client):
         )
         self.assertEqual(api_key, "test-api-key")
 
-    @patch("boto3.client")
+    @patch("botocore.session.Session.create_client")
     def test_ssm_fips_endpoint(self, mock_boto3_client):
         mock_client = MagicMock()
         mock_client.get_parameter.return_value = {
@@ -80,7 +80,7 @@ def test_ssm_fips_endpoint(self, mock_boto3_client):
         )
         self.assertEqual(api_key, "test-api-key")
 
-    @patch("boto3.client")
+    @patch("botocore.session.Session.create_client")
     @patch("datadog_lambda.api.decrypt_kms_api_key")
     def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client):
         mock_client = MagicMock()
@@ -97,7 +97,7 @@ def test_kms_fips_endpoint(self, mock_decrypt_kms, mock_boto3_client):
         )
         self.assertEqual(api_key, "test-api-key")
 
-    @patch("boto3.client")
+    @patch("botocore.session.Session.create_client")
     def test_no_fips_for_standard_regions(self, mock_boto3_client):
         mock_client = MagicMock()
         mock_client.get_secret_value.return_value = {"SecretString": "test-api-key"}