diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index ad860873..c539ea05 100644 --- a/datadog_lambda/api.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index a3860567..8db5e352 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ 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 } @@ -38,7 +38,7 @@ flake8 = { version = "^5.0.4", optional = true } [tool.poetry.extras] dev = [ - "boto3", + "botocore", "flake8", "pytest", "pytest-benchmark", diff --git a/tests/test_api.py b/tests/test_api.py index c7facb43..c98d91eb 100644 --- a/tests/test_api.py +++ b/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"}