[PD1-66] - Allow api keys mapping to fetch data across tenants

pyproject.toml

@@ -8,6 +8,7 @@ requires-python = ">=3.10"
 dependencies = [
     "pydantic>=2.7.0,<3.0.0",
     "pyarrow>=18.0.0,<20.0.0",
+    "boto3>=1.26.0",
 ]
 
 [tool.poetry]

src/cvec/cvec.py

@@ -7,6 +7,7 @@
 from urllib.parse import urlencode, urljoin
 from urllib.request import Request, urlopen
 
+from cvec.utils.api_key_service import get_api_key_for_host
 from cvec.models.metric import Metric, MetricDataPoint
 from cvec.models.span import Span
 from cvec.utils.arrow_converter import (
@@ -52,9 +53,23 @@ def __init__(
             raise ValueError(
                 "CVEC_HOST must be set either as an argument or environment variable"
             )
+
+        # If api_key is not provided, try to fetch it from the API key service
+        if not self._api_key:
+            try:
+                self._api_key = get_api_key_for_host(self.host)
+                if self._api_key:
+                    logger.info(
+                        f"API key loaded from mapping service for host: {self.host}"
+                    )
+            except ValueError as e:
+                # Log the error from the service, but we'll raise our own error below
+                logger.debug(f"Failed to load API key from service: {e}")
+
         if not self._api_key:
             raise ValueError(
-                "CVEC_API_KEY must be set either as an argument or environment variable"
+                "CVEC_API_KEY must be set either as an argument, environment variable, "
+                "or available in the API keys mapping (AWS_API_KEYS_SECRET or API_KEYS_MAPPING)"
             )
 
         # Fetch publishable key from host config

src/cvec/utils/api_key_service.py

@@ -0,0 +1,164 @@
+"""
+API Key Service
+
+Loads host-to-API-key mappings from AWS Secrets Manager or local environment variables.
+Supports both AWS Secrets Manager (for production) and local JSON environment variables (for development).
+"""
+
+import json
+import logging
+import os
+from typing import Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+# Global cache for API keys (loaded once per process/cold start)
+_API_KEYS_CACHE: Optional[Dict[str, str]] = None
+
+
+def load_api_keys() -> Optional[Dict[str, str]]:
+    """
+    Load host-to-API-key mapping from AWS Secrets Manager or local environment variables.
+    Caches the mapping in a global variable to avoid repeated API calls.
+
+    The loading priority is:
+    1. AWS Secrets Manager (if AWS_API_KEYS_SECRET environment variable is set)
+    2. Local JSON environment variable API_KEYS_MAPPING (for development)
+
+    Returns:
+        dict: Mapping of host URLs to API keys, or None if no source is available
+    """
+    global _API_KEYS_CACHE
+
+    if _API_KEYS_CACHE is not None:
+        logger.debug("Using cached API keys")
+        return _API_KEYS_CACHE
+
+    # Try AWS Secrets Manager first
+    aws_secret_name = os.environ.get("AWS_API_KEYS_SECRET")
+    if aws_secret_name:
+        logger.info(
+            f"AWS_API_KEYS_SECRET is set, loading from Secrets Manager: {aws_secret_name}"
+        )
+        api_keys = _load_from_aws_secrets_manager(aws_secret_name)
+        if api_keys is not None:
+            _API_KEYS_CACHE = api_keys
+            return _API_KEYS_CACHE
+
+    # Fall back to local environment variable
+    local_mapping = os.environ.get("API_KEYS_MAPPING")
+    if local_mapping:
+        logger.info("Loading API keys from local environment variable API_KEYS_MAPPING")
+        api_keys = _load_from_local_env(local_mapping)
+        if api_keys is not None:
+            _API_KEYS_CACHE = api_keys
+            return _API_KEYS_CACHE
+
+    logger.warning(
+        "No API key mapping found. Set either AWS_API_KEYS_SECRET or API_KEYS_MAPPING environment variable."
+    )
+    return None
+
+
+def _load_from_aws_secrets_manager(secret_name: str) -> Optional[Dict[str, str]]:
+    """
+    Load API keys from AWS Secrets Manager.
+
+    Args:
+        secret_name: The name/ARN of the secret in AWS Secrets Manager
+
+    Returns:
+        dict: Mapping of host URLs to API keys, or None if loading fails
+    """
+    try:
+        import boto3  # type: ignore[import-untyped]
+
+        secretsmanager = boto3.client("secretsmanager")
+
+        response = secretsmanager.get_secret_value(SecretId=secret_name)
+
+        # Check if SecretString exists and is not empty
+        if "SecretString" not in response:
+            logger.warning(
+                f"Secret '{secret_name}' exists but has no value. Please populate it."
+            )
+            return None
+
+        secret_string = response["SecretString"]
+
+        # Check if the secret string is empty or whitespace
+        if not secret_string or not secret_string.strip():
+            logger.warning(f"Secret '{secret_name}' is empty. Please populate it.")
+            return None
+
+        api_keys: Dict[str, str] = json.loads(secret_string)
+        logger.info(
+            f"Loaded API keys for {len(api_keys)} hosts from AWS Secrets Manager"
+        )
+        return api_keys
+
+    except json.JSONDecodeError as e:
+        logger.error(f"Secret '{secret_name}' contains invalid JSON: {str(e)}")
+        return None
+    except Exception as e:
+        # Generic catch-all for any other errors (e.g., ResourceNotFoundException, network errors, etc.)
+        logger.error(
+            f"Error loading API keys from AWS Secrets Manager: {str(e)}",
+            exc_info=True,
+        )
+        return None
+
+
+def _load_from_local_env(json_string: str) -> Optional[Dict[str, str]]:
+    """
+    Load API keys from local environment variable containing JSON.
+
+    Args:
+        json_string: JSON string containing host-to-API-key mapping
+
+    Returns:
+        dict: Mapping of host URLs to API keys, or None if parsing fails
+    """
+    try:
+        api_keys = json.loads(json_string)
+        if not isinstance(api_keys, dict):
+            logger.error(
+                "API_KEYS_MAPPING must be a JSON object/dict mapping hosts to API keys"
+            )
+            return None
+
+        logger.info(f"Loaded API keys for {len(api_keys)} hosts from local environment")
+        return api_keys
+
+    except json.JSONDecodeError as e:
+        logger.error(f"API_KEYS_MAPPING contains invalid JSON: {str(e)}")
+        return None
+
+
+def get_api_key_for_host(host: str) -> Optional[str]:
+    """
+    Get the API key for a specific host.
+
+    Args:
+        host: The host URL
+
+    Returns:
+        str: The API key for the host, or None if not found
+
+    Raises:
+        ValueError: If no API keys mapping is available
+    """
+    api_keys = load_api_keys()
+
+    if api_keys is None:
+        raise ValueError(
+            "No API keys mapping available. Set either AWS_API_KEYS_SECRET or "
+            "API_KEYS_MAPPING environment variable."
+        )
+
+    api_key = api_keys.get(host)
+    if not api_key:
+        logger.warning(f"No API key found for host: {host}")
+        return None
+
+    return api_key