extract getting api key to its own function

Author: sabrenner

datadog_lambda/api.py

@@ -4,6 +4,7 @@
 
 logger = logging.getLogger(__name__)
 KMS_ENCRYPTION_CONTEXT_KEY = "LambdaFunctionName"
+api_key = None
 
 
 def decrypt_kms_api_key(kms_client, ciphertext):
@@ -46,6 +47,41 @@ def decrypt_kms_api_key(kms_client, ciphertext):
     return plaintext
 
 
+def get_api_key() -> str:
+    """
+    Gets the Datadog API key from the environment variables or secrets manager.
+    Extracts the result to a global value to avoid repeated calls to the secrets manager from different products.
+    """
+    global api_key
+    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", "")
+    DD_API_KEY = os.environ.get(
+        "DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")
+    )
+
+    if DD_API_KEY_SECRET_ARN:
+        api_key = boto3.client("secretsmanager").get_secret_value(
+            SecretId=DD_API_KEY_SECRET_ARN
+        )["SecretString"]
+    elif DD_API_KEY_SSM_NAME:
+        api_key = boto3.client("ssm").get_parameter(
+            Name=DD_API_KEY_SSM_NAME, WithDecryption=True
+        )["Parameter"]["Value"]
+    elif DD_KMS_API_KEY:
+        kms_client = boto3.client("kms")
+        api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY)
+    else:
+        api_key = DD_API_KEY
+
+    return api_key
+
+
 def init_api():
     if not os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
         # Make sure that this package would always be lazy-loaded/outside from the critical path
@@ -54,28 +90,7 @@ def init_api():
         from datadog import api
 
         if not api._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", "")
-            DD_API_KEY = os.environ.get(
-                "DD_API_KEY", os.environ.get("DATADOG_API_KEY", "")
-            )
-
-            if DD_API_KEY_SECRET_ARN:
-                api._api_key = boto3.client("secretsmanager").get_secret_value(
-                    SecretId=DD_API_KEY_SECRET_ARN
-                )["SecretString"]
-            elif DD_API_KEY_SSM_NAME:
-                api._api_key = boto3.client("ssm").get_parameter(
-                    Name=DD_API_KEY_SSM_NAME, WithDecryption=True
-                )["Parameter"]["Value"]
-            elif DD_KMS_API_KEY:
-                kms_client = boto3.client("kms")
-                api._api_key = decrypt_kms_api_key(kms_client, DD_KMS_API_KEY)
-            else:
-                api._api_key = DD_API_KEY
+            api._api_key = get_api_key()
 
         logger.debug("Setting DATADOG_API_KEY of length %d", len(api._api_key))
 

datadog_lambda/wrapper.py

@@ -59,12 +59,10 @@
 llmobs_api_key = None
 llmobs_env_var = os.environ.get("DD_LLMOBS_ENABLED", "false").lower() in ("true", "1")
 if llmobs_env_var:
-    from datadog_lambda.api import init_api
-    from datadog import api
+    from datadog_lambda.api import get_api_key
     from ddtrace.llmobs import LLMObs
 
-    init_api()
-    llmobs_api_key = api._api_key
+    llmobs_api_key = get_api_key()
 
 logger = logging.getLogger(__name__)