feat: add ldms feature (#8051)

* feat: add ldms service

* feat: add ldms service

* feat: add ldms service

Author: leandrodamascena

.github/workflows/quality_check.yml

@@ -43,6 +43,7 @@ jobs:
   quality_check:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       max-parallel: 5
       matrix:
         python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]

aws_lambda_powertools/shared/constants.py

@@ -67,5 +67,12 @@
 PRETTY_INDENT: int = 4
 COMPACT_INDENT: None = None
 
+# Metadata constants
+LAMBDA_METADATA_API_ENV: str = "AWS_LAMBDA_METADATA_API"
+LAMBDA_METADATA_TOKEN_ENV: str = "AWS_LAMBDA_METADATA_TOKEN"
+METADATA_API_VERSION: str = "2026-01-15"
+METADATA_PATH: str = "/metadata/execution-environment"
+METADATA_DEFAULT_TIMEOUT_SECS: float = 1.0
+
 # Idempotency constants
 IDEMPOTENCY_DISABLED_ENV: str = "POWERTOOLS_IDEMPOTENCY_DISABLED"

aws_lambda_powertools/utilities/metadata/__init__.py

@@ -0,0 +1,17 @@
+"""
+Utility to fetch data from the AWS Lambda Metadata Endpoint
+"""
+
+from aws_lambda_powertools.utilities.metadata.exceptions import LambdaMetadataError
+from aws_lambda_powertools.utilities.metadata.lambda_metadata import (
+    LambdaMetadata,
+    clear_metadata_cache,
+    get_lambda_metadata,
+)
+
+__all__ = [
+    "LambdaMetadata",
+    "LambdaMetadataError",
+    "get_lambda_metadata",
+    "clear_metadata_cache",
+]

aws_lambda_powertools/utilities/metadata/exceptions.py

@@ -0,0 +1,11 @@
+"""
+Lambda Metadata Service exceptions
+"""
+
+
+class LambdaMetadataError(Exception):
+    """Raised when the Lambda Metadata Endpoint is unavailable or returns an error."""
+
+    def __init__(self, message: str, status_code: int = -1):
+        self.status_code = status_code
+        super().__init__(message)

aws_lambda_powertools/utilities/metadata/lambda_metadata.py

@@ -0,0 +1,174 @@
+"""
+Lambda Metadata Service client
+
+Fetches execution environment metadata from the Lambda Metadata Endpoint,
+with caching for the sandbox lifetime.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import urllib.request
+from dataclasses import dataclass, field
+from json import JSONDecodeError
+from json import loads as json_loads
+from typing import Any
+
+from aws_lambda_powertools.shared.constants import (
+    LAMBDA_INITIALIZATION_TYPE,
+    LAMBDA_METADATA_API_ENV,
+    LAMBDA_METADATA_TOKEN_ENV,
+    METADATA_API_VERSION,
+    METADATA_DEFAULT_TIMEOUT_SECS,
+    METADATA_PATH,
+    POWERTOOLS_DEV_ENV,
+)
+from aws_lambda_powertools.utilities.metadata.exceptions import LambdaMetadataError
+
+logger = logging.getLogger(__name__)
+
+_cache: dict[str, Any] = {}
+
+
+@dataclass(frozen=True)
+class LambdaMetadata:
+    """Lambda execution environment metadata returned by the metadata endpoint."""
+
+    availability_zone_id: str | None = None
+    """The Availability Zone ID where the function is executing (e.g. ``use1-az1``)."""
+
+    _raw: dict[str, Any] = field(default_factory=dict, repr=False)
+    """Full raw response for forward-compatibility with future fields."""
+
+
+def _is_lambda_environment() -> bool:
+    """Check whether we are running inside a Lambda execution environment."""
+    return os.environ.get(LAMBDA_INITIALIZATION_TYPE, "") != ""
+
+
+def _is_dev_mode() -> bool:
+    """Check whether POWERTOOLS_DEV is enabled."""
+    return os.environ.get(POWERTOOLS_DEV_ENV, "false").strip().lower() in ("true", "1")
+
+
+def _build_metadata(data: dict[str, Any]) -> LambdaMetadata:
+    """Build a LambdaMetadata dataclass from the raw endpoint response."""
+    return LambdaMetadata(
+        availability_zone_id=data.get("AvailabilityZoneID"),
+        _raw=data,
+    )
+
+
+def _fetch_metadata(timeout: float = METADATA_DEFAULT_TIMEOUT_SECS) -> dict[str, Any]:
+    """
+    Fetch metadata from the Lambda Metadata Endpoint via HTTP.
+
+    Parameters
+    ----------
+    timeout : float
+        Request timeout in seconds.
+
+    Returns
+    -------
+    dict[str, Any]
+        Parsed JSON response from the metadata endpoint.
+
+    Raises
+    ------
+    LambdaMetadataError
+        If required environment variables are missing, the endpoint returns
+        a non-200 status, or the response cannot be parsed.
+    """
+    api = os.environ.get(LAMBDA_METADATA_API_ENV)
+    token = os.environ.get(LAMBDA_METADATA_TOKEN_ENV)
+
+    if not api:
+        raise LambdaMetadataError(
+            f"Environment variable {LAMBDA_METADATA_API_ENV} is not set. Ensure {LAMBDA_METADATA_API_ENV} is set.",
+        )
+    if not token:
+        raise LambdaMetadataError(
+            f"Environment variable {LAMBDA_METADATA_TOKEN_ENV} is not set. Ensure {LAMBDA_METADATA_TOKEN_ENV} is set.",
+        )
+
+    url = f"http://{api}/{METADATA_API_VERSION}{METADATA_PATH}"
+    logger.debug("Fetching Lambda metadata from: %s", url)
+
+    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"})
+
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:  # nosec B310
+            status = resp.status
+            body = resp.read().decode("utf-8")
+    except urllib.error.HTTPError as exc:
+        raise LambdaMetadataError(
+            f"Metadata request failed with status {exc.code}",
+            status_code=exc.code,
+        ) from exc
+    except Exception as exc:
+        raise LambdaMetadataError(f"Failed to fetch Lambda metadata: {exc}") from exc
+
+    if status != 200:
+        raise LambdaMetadataError(
+            f"Metadata request failed with status {status}",
+            status_code=status,
+        )
+
+    try:
+        data: dict[str, Any] = json_loads(body)
+    except (JSONDecodeError, TypeError) as exc:
+        raise LambdaMetadataError(f"Failed to parse metadata response: {exc}") from exc
+
+    logger.debug("Lambda metadata response: %s", data)
+    return data
+
+
+def get_lambda_metadata(*, timeout: float = METADATA_DEFAULT_TIMEOUT_SECS) -> LambdaMetadata:
+    """
+    Retrieve Lambda execution environment metadata.
+
+    Returns cached metadata on subsequent calls. When not running in a Lambda
+    environment (local dev, tests) or when ``POWERTOOLS_DEV`` is enabled,
+    returns an empty ``LambdaMetadata``.
+
+    Parameters
+    ----------
+    timeout : float
+        HTTP request timeout in seconds (default 1.0).
+
+    Returns
+    -------
+    LambdaMetadata
+        Metadata about the current execution environment.
+
+    Raises
+    ------
+    LambdaMetadataError
+        If the metadata endpoint is unavailable or returns an error.
+
+    Example
+    -------
+        >>> from aws_lambda_powertools.utilities.metadata import get_lambda_metadata
+        >>> metadata = get_lambda_metadata()
+        >>> metadata.availability_zone_id  # e.g. "use1-az1"
+    """
+    if _is_dev_mode() or not _is_lambda_environment():
+        return LambdaMetadata()
+
+    if _cache:
+        return _build_metadata(_cache)
+
+    data = _fetch_metadata(timeout=timeout)
+    _cache.update(data)
+    return _build_metadata(_cache)
+
+
+def clear_metadata_cache() -> None:
+    """
+    Clear the cached metadata.
+
+    Useful for testing or when you need to force a fresh fetch
+    (e.g. after SnapStart restore).
+    """
+    _cache.clear()

docs/utilities/metadata.md

@@ -0,0 +1,59 @@
+---
+title: Metadata
+description: Utility
+status: new
+---
+
+<!-- markdownlint-disable MD043 -->
+
+The Metadata utility allows you to fetch data from the [AWS Lambda Metadata Endpoint (LMDS)](https://docs.aws.amazon.com/lambda/latest/dg/lambda-metadata-endpoint.html){target="_blank"}. This can be useful for retrieving information about the Lambda execution environment, such as the Availability Zone ID.
+
+## Key features
+
+* Fetch execution environment metadata from the Lambda Metadata Endpoint
+* Automatic caching for the duration of the Lambda sandbox
+* Graceful fallback to empty metadata outside Lambda (local dev, tests)
+* Forward-compatible dataclass that can be extended as new fields are added
+
+## Getting started
+
+### Usage
+
+You can fetch data from the Lambda Metadata Endpoint using the `get_lambda_metadata` function.
+
+???+ tip
+    Metadata is cached for the duration of the Lambda sandbox, so subsequent calls to `get_lambda_metadata` will return the cached data.
+
+=== "getting_started_metadata.py"
+
+    ```python hl_lines="2 9 10"
+    --8<-- "examples/metadata/src/getting_started_metadata.py"
+    ```
+
+You can also fetch metadata eagerly during cold start, so it's ready for subsequent invocations:
+
+=== "getting_started_metadata_eager.py"
+
+    ```python hl_lines="2 8"
+    --8<-- "examples/metadata/src/getting_started_metadata_eager.py"
+    ```
+
+### Available metadata
+
+| Property               | Type            | Description                                                    |
+| ---------------------- | --------------- | -------------------------------------------------------------- |
+| `availability_zone_id` | `str` or `None` | The AZ where the function is running (e.g., `use1-az1`)        |
+
+## Testing your code
+
+The metadata endpoint is not available during local development or testing. To ease testing, the `get_lambda_metadata` function automatically detects when it's running in a non-Lambda environment and returns an empty `LambdaMetadata` instance. This allows you to write tests without needing to mock the endpoint.
+
+If you want to mock specific metadata values for testing purposes, you can patch the internal `_fetch_metadata` function and set the required environment variables:
+
+=== "testing_metadata.py"
+
+    ```python hl_lines="6-8 13-18 21"
+    --8<-- "examples/metadata/src/testing_metadata.py"
+    ```
+
+We also expose a `clear_metadata_cache` function that can be used to clear the cached metadata, allowing you to test different metadata values within the same execution context.

examples/metadata/src/getting_started_metadata.py

@@ -0,0 +1,15 @@
+from aws_lambda_powertools import Logger
+from aws_lambda_powertools.utilities.metadata import LambdaMetadata, get_lambda_metadata
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+logger = Logger()
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    metadata: LambdaMetadata = get_lambda_metadata()
+    az_id = metadata.availability_zone_id  # e.g., "use1-az1"
+
+    logger.append_keys(az_id=az_id)
+    logger.info("Processing request")
+
+    return {"az_id": az_id}

examples/metadata/src/getting_started_metadata_eager.py

@@ -0,0 +1,15 @@
+from aws_lambda_powertools import Logger
+from aws_lambda_powertools.utilities.metadata import LambdaMetadata, get_lambda_metadata
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+logger = Logger()
+
+# Fetch during cold start — cached for subsequent invocations
+metadata: LambdaMetadata = get_lambda_metadata()
+
+
+def lambda_handler(event: dict, context: LambdaContext) -> dict:
+    logger.append_keys(az_id=metadata.availability_zone_id)
+    logger.info("Processing request")
+
+    return {"az_id": metadata.availability_zone_id}

examples/metadata/src/testing_metadata.py

@@ -0,0 +1,34 @@
+from unittest.mock import patch
+
+from aws_lambda_powertools.utilities.metadata import LambdaMetadata, clear_metadata_cache, get_lambda_metadata
+
+
+def test_handler_uses_metadata(monkeypatch):
+    # GIVEN a Lambda environment with metadata env vars
+    monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand")
+    monkeypatch.setenv("AWS_LAMBDA_METADATA_API", "127.0.0.1:1234")
+    monkeypatch.setenv("AWS_LAMBDA_METADATA_TOKEN", "test-token")
+
+    mock_response = {"AvailabilityZoneID": "use1-az1"}
+
+    with patch(
+        "aws_lambda_powertools.utilities.metadata.lambda_metadata._fetch_metadata",
+        return_value=mock_response,
+    ):
+        # WHEN calling get_lambda_metadata
+        metadata: LambdaMetadata = get_lambda_metadata()
+
+        # THEN it returns the mocked metadata
+        assert metadata.availability_zone_id == "use1-az1"
+
+    # Clean up cache between tests
+    clear_metadata_cache()
+
+
+def test_handler_works_outside_lambda():
+    # GIVEN no Lambda environment variables are set
+    # WHEN calling get_lambda_metadata
+    metadata: LambdaMetadata = get_lambda_metadata()
+
+    # THEN it returns empty metadata without errors
+    assert metadata.availability_zone_id is None

mkdocs.yml

@@ -34,6 +34,7 @@ nav:
           - utilities/idempotency.md
           - utilities/data_masking.md
           - utilities/feature_flags.md
+          - utilities/metadata.md
           - utilities/streaming.md
           - utilities/middleware_factory.md
           - utilities/jmespath_functions.md
@@ -244,6 +245,7 @@ plugins:
           - utilities/idempotency.md
           - utilities/data_masking.md
           - utilities/feature_flags.md
+          - utilities/metadata.md
           - utilities/streaming.md
           - utilities/middleware_factory.md
           - utilities/jmespath_functions.md