diff --git a/datadog_lambda/api.py b/datadog_lambda/api.py index b5414fd9..a114fe8f 100644 --- a/datadog_lambda/api.py +++ b/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,40 @@ 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 +89,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)) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 2632d22e..6afa9a07 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -56,10 +56,14 @@ if profiling_env_var: from ddtrace.profiling import profiler +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 get_api_key from ddtrace.llmobs import LLMObs + llmobs_api_key = get_api_key() + logger = logging.getLogger(__name__) DD_FLUSH_TO_LOG = "DD_FLUSH_TO_LOG" @@ -229,7 +233,10 @@ def __init__(self, func): # Enable LLM Observability if llmobs_env_var: - LLMObs.enable() + LLMObs.enable( + agentless_enabled=True, + api_key=llmobs_api_key, + ) logger.debug("datadog_lambda_wrapper initialized") except Exception as e: