Add logfire.instrument_aws_lambda

docs/integrations/aws-lambda.md

@@ -0,0 +1,34 @@
+# AWS Lambda
+
+The [`logfire.instrument_aws_lambda`][logfire.Logfire.instrument_aws_lambda] function can be used to
+instrument AWS Lambda functions to automatically send traces to **Logfire**.
+
+## Installation
+
+Install `logfire` with the `aws-lambda` extra:
+
+{{ install_logfire(extras=['aws-lambda']) }}
+
+## Usage
+
+To instrument an AWS Lambda function, call the `logfire.instrument_aws_lambda` function after defining
+the handler function:
+
+```python
+import logfire
+
+logfire.configure()  # (1)!
+
+
+def handler(event, context):
+    return 'Hello from Lambda'
+
+logfire.instrument_aws_lambda(handler)
+```
+
+1. Remember to set the `LOGFIRE_TOKEN` environment variable on your Lambda function configuration.
+
+[`logfire.instrument_aws_lambda`][logfire.Logfire.instrument_aws_lambda] uses the **OpenTelemetry AWS Lambda Instrumentation** package,
+which you can find more information about [here][opentelemetry-aws-lambda].
+
+[opentelemetry-aws-lambda]: https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/aws_lambda/aws_lambda.html

logfire-api/logfire_api/__init__.py

@@ -113,6 +113,8 @@ def instrument_sqlalchemy(self, *args, **kwargs) -> None: ...
 
             def instrument_sqlite3(self, *args, **kwargs) -> None: ...
 
+            def instrument_aws_lambda(self, *args, **kwargs) -> None: ...
+
             def instrument_redis(self, *args, **kwargs) -> None: ...
 
             def instrument_flask(self, *args, **kwargs) -> None: ...
@@ -175,6 +177,7 @@ def shutdown(self, *args, **kwargs) -> None: ...
         instrument_aiohttp_client = DEFAULT_LOGFIRE_INSTANCE.instrument_aiohttp_client
         instrument_sqlalchemy = DEFAULT_LOGFIRE_INSTANCE.instrument_sqlalchemy
         instrument_sqlite3 = DEFAULT_LOGFIRE_INSTANCE.instrument_sqlite3
+        instrument_aws_lambda = DEFAULT_LOGFIRE_INSTANCE.instrument_aws_lambda
         instrument_redis = DEFAULT_LOGFIRE_INSTANCE.instrument_redis
         instrument_pymongo = DEFAULT_LOGFIRE_INSTANCE.instrument_pymongo
         instrument_mysql = DEFAULT_LOGFIRE_INSTANCE.instrument_mysql

logfire-api/logfire_api/__init__.pyi

@@ -12,7 +12,7 @@ from .version import VERSION as VERSION
 from logfire.sampling import SamplingOptions as SamplingOptions
 from typing import Any
 
-__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_asgi', 'instrument_wsgi', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_sqlite3', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'suppress_scopes', 'shutdown', 'load_spans_from_file', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
+__all__ = ['Logfire', 'LogfireSpan', 'LevelName', 'AdvancedOptions', 'ConsoleOptions', 'CodeSource', 'PydanticPlugin', 'configure', 'span', 'instrument', 'log', 'trace', 'debug', 'notice', 'info', 'warn', 'error', 'exception', 'fatal', 'force_flush', 'log_slow_async_callbacks', 'install_auto_tracing', 'instrument_asgi', 'instrument_wsgi', 'instrument_pydantic', 'instrument_fastapi', 'instrument_openai', 'instrument_anthropic', 'instrument_asyncpg', 'instrument_httpx', 'instrument_celery', 'instrument_requests', 'instrument_psycopg', 'instrument_django', 'instrument_flask', 'instrument_starlette', 'instrument_aiohttp_client', 'instrument_sqlalchemy', 'instrument_sqlite3', 'instrument_aws_lambda', 'instrument_redis', 'instrument_pymongo', 'instrument_mysql', 'instrument_system_metrics', 'AutoTraceModule', 'with_tags', 'with_settings', 'suppress_scopes', 'shutdown', 'load_spans_from_file', 'no_auto_trace', 'ScrubMatch', 'ScrubbingOptions', 'VERSION', 'suppress_instrumentation', 'StructlogProcessor', 'LogfireLoggingHandler', 'loguru_handler', 'SamplingOptions', 'MetricsOptions']
 
 DEFAULT_LOGFIRE_INSTANCE = Logfire()
 span = DEFAULT_LOGFIRE_INSTANCE.span
@@ -37,6 +37,7 @@ instrument_starlette = DEFAULT_LOGFIRE_INSTANCE.instrument_starlette
 instrument_aiohttp_client = DEFAULT_LOGFIRE_INSTANCE.instrument_aiohttp_client
 instrument_sqlalchemy = DEFAULT_LOGFIRE_INSTANCE.instrument_sqlalchemy
 instrument_sqlite3 = DEFAULT_LOGFIRE_INSTANCE.instrument_sqlite3
+instrument_aws_lambda = DEFAULT_LOGFIRE_INSTANCE.instrument_aws_lambda
 instrument_redis = DEFAULT_LOGFIRE_INSTANCE.instrument_redis
 instrument_pymongo = DEFAULT_LOGFIRE_INSTANCE.instrument_pymongo
 instrument_mysql = DEFAULT_LOGFIRE_INSTANCE.instrument_mysql
@@ -55,7 +56,7 @@ error = DEFAULT_LOGFIRE_INSTANCE.error
 fatal = DEFAULT_LOGFIRE_INSTANCE.fatal
 exception = DEFAULT_LOGFIRE_INSTANCE.exception
 
-def loguru_handler() -> dict[str, Any]:
+def loguru_handler() -> Any:
     """Create a **Logfire** handler for Loguru.
 
     Returns:

logfire-api/logfire_api/_internal/integrations/aws_lambda.pyi

@@ -0,0 +1,17 @@
+from opentelemetry.context import Context as Context
+from opentelemetry.metrics import MeterProvider
+from opentelemetry.trace import TracerProvider
+from typing import Any, Callable, TypedDict, Unpack
+
+LambdaEvent = Any
+LambdaFunction = Callable[[LambdaEvent, Any], Any]
+
+class AwsLambdaInstrumentKwargs(TypedDict, total=False):
+    skip_dep_check: bool
+    event_context_extractor: Callable[[LambdaEvent], Context]
+
+def instrument_aws_lambda(lambda_function: LambdaFunction, *, tracer_provider: TracerProvider, meter_provider: MeterProvider, **kwargs: Unpack[AwsLambdaInstrumentKwargs]) -> None:
+    """Instrument the AWS Lambda runtime so that spans are automatically created for each invocation.
+
+    See the `Logfire.instrument_aws_lambda` method for details.
+    """

logfire-api/logfire_api/_internal/main.pyi

@@ -12,6 +12,7 @@ from .formatter import logfire_format as logfire_format, logfire_format_with_mag
 from .instrument import instrument as instrument
 from .integrations.asgi import ASGIApp as ASGIApp, ASGIInstrumentKwargs as ASGIInstrumentKwargs
 from .integrations.asyncpg import AsyncPGInstrumentKwargs as AsyncPGInstrumentKwargs
+from .integrations.aws_lambda import AwsLambdaInstrumentKwargs as AwsLambdaInstrumentKwargs, LambdaFunction as LambdaFunction
 from .integrations.celery import CeleryInstrumentKwargs as CeleryInstrumentKwargs
 from .integrations.flask import FlaskInstrumentKwargs as FlaskInstrumentKwargs
 from .integrations.httpx import HTTPXInstrumentKwargs as HTTPXInstrumentKwargs
@@ -719,12 +720,19 @@ class Logfire:
         Returns:
             If a connection is provided, returns the instrumented connection. If no connection is provided, returns `None`.
         """
+    def instrument_aws_lambda(self, lambda_function: LambdaFunction, **kwargs: Unpack[AwsLambdaInstrumentKwargs]) -> None:
+        """Instrument AWS Lambda so that spans are automatically created for each invocation.
+
+        Uses the
+        [OpenTelemetry AWS Lambda Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/aws_lambda/aws_lambda.html)
+        library, specifically `AwsLambdaInstrumentor().instrument()`, to which it passes `**kwargs`.
+        """
     def instrument_pymongo(self, **kwargs: Unpack[PymongoInstrumentKwargs]) -> None:
         """Instrument the `pymongo` module so that spans are automatically created for each operation.
 
         Uses the
         [OpenTelemetry pymongo Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/pymongo/pymongo.html)
-            library, specifically `PymongoInstrumentor().instrument()`, to which it passes `**kwargs`.
+        library, specifically `PymongoInstrumentor().instrument()`, to which it passes `**kwargs`.
         """
     def instrument_redis(self, capture_statement: bool = False, **kwargs: Unpack[RedisInstrumentKwargs]) -> None:
         """Instrument the `redis` module so that spans are automatically created for each operation.

logfire/__init__.py

@@ -48,6 +48,7 @@
 instrument_aiohttp_client = DEFAULT_LOGFIRE_INSTANCE.instrument_aiohttp_client
 instrument_sqlalchemy = DEFAULT_LOGFIRE_INSTANCE.instrument_sqlalchemy
 instrument_sqlite3 = DEFAULT_LOGFIRE_INSTANCE.instrument_sqlite3
+instrument_aws_lambda = DEFAULT_LOGFIRE_INSTANCE.instrument_aws_lambda
 instrument_redis = DEFAULT_LOGFIRE_INSTANCE.instrument_redis
 instrument_pymongo = DEFAULT_LOGFIRE_INSTANCE.instrument_pymongo
 instrument_mysql = DEFAULT_LOGFIRE_INSTANCE.instrument_mysql
@@ -132,6 +133,7 @@ def loguru_handler() -> Any:
     'instrument_aiohttp_client',
     'instrument_sqlalchemy',
     'instrument_sqlite3',
+    'instrument_aws_lambda',
     'instrument_redis',
     'instrument_pymongo',
     'instrument_mysql',

logfire/_internal/integrations/aws_lambda.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+try:
+    from opentelemetry.context import Context
+    from opentelemetry.instrumentation.aws_lambda import AwsLambdaInstrumentor
+    from opentelemetry.metrics import MeterProvider
+    from opentelemetry.trace import TracerProvider
+except ImportError:
+    raise RuntimeError(
+        '`logfire.instrument_aws_lambda()` requires the `opentelemetry-instrumentation-aws-lambda` package.\n'
+        'You can install this with:\n'
+        "    pip install 'logfire[aws-lambda]'"
+    )
+
+if TYPE_CHECKING:
+    from typing import Any, Callable, TypedDict, Unpack
+
+    LambdaEvent = Any
+    LambdaFunction = Callable[[LambdaEvent, Any], Any]
+
+    class AwsLambdaInstrumentKwargs(TypedDict, total=False):
+        skip_dep_check: bool
+        event_context_extractor: Callable[[LambdaEvent], Context]
+
+
+def instrument_aws_lambda(
+    lambda_function: LambdaFunction,
+    *,
+    tracer_provider: TracerProvider,
+    meter_provider: MeterProvider,
+    **kwargs: Unpack[AwsLambdaInstrumentKwargs],
+) -> None:
+    """Instrument the AWS Lambda runtime so that spans are automatically created for each invocation.
+
+    See the `Logfire.instrument_aws_lambda` method for details.
+    """
+    return AwsLambdaInstrumentor().instrument(  # type: ignore[no-any-return]
+        tracer_provider=tracer_provider, meter_provider=meter_provider, **kwargs
+    )

logfire/_internal/main.py

@@ -83,6 +83,7 @@
 
     from .integrations.asgi import ASGIApp, ASGIInstrumentKwargs
     from .integrations.asyncpg import AsyncPGInstrumentKwargs
+    from .integrations.aws_lambda import AwsLambdaInstrumentKwargs, LambdaFunction
     from .integrations.celery import CeleryInstrumentKwargs
     from .integrations.flask import FlaskInstrumentKwargs
     from .integrations.httpx import HTTPXInstrumentKwargs
@@ -1451,12 +1452,33 @@ def instrument_sqlite3(
             },
         )
 
+    def instrument_aws_lambda(
+        self, lambda_function: LambdaFunction, **kwargs: Unpack[AwsLambdaInstrumentKwargs]
+    ) -> None:
+        """Instrument AWS Lambda so that spans are automatically created for each invocation.
+
+        Uses the
+        [OpenTelemetry AWS Lambda Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/aws_lambda/aws_lambda.html)
+        library, specifically `AwsLambdaInstrumentor().instrument()`, to which it passes `**kwargs`.
+        """
+        from .integrations.aws_lambda import instrument_aws_lambda
+
+        self._warn_if_not_initialized_for_instrumentation()
+        return instrument_aws_lambda(
+            lambda_function=lambda_function,
+            **{  # type: ignore
+                'tracer_provider': self._config.get_tracer_provider(),
+                'meter_provider': self._config.get_meter_provider(),
+                **kwargs,
+            },
+        )
+
     def instrument_pymongo(self, **kwargs: Unpack[PymongoInstrumentKwargs]) -> None:
         """Instrument the `pymongo` module so that spans are automatically created for each operation.
 
         Uses the
         [OpenTelemetry pymongo Instrumentation](https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/pymongo/pymongo.html)
-            library, specifically `PymongoInstrumentor().instrument()`, to which it passes `**kwargs`.
+        library, specifically `PymongoInstrumentor().instrument()`, to which it passes `**kwargs`.
         """
         from .integrations.pymongo import instrument_pymongo
 

mkdocs.yml

@@ -126,6 +126,7 @@ nav:
           - Celery: integrations/celery.md
           - Stripe: integrations/stripe.md
           - System Metrics: integrations/system-metrics.md
+          - AWS Lambda: integrations/aws-lambda.md
       - Custom:
           - Pydantic: integrations/pydantic.md
           - OpenAI: integrations/openai.md

pyproject.toml

@@ -75,6 +75,7 @@ redis = ["opentelemetry-instrumentation-redis >= 0.42b0"]
 requests = ["opentelemetry-instrumentation-requests >= 0.42b0"]
 mysql = ["opentelemetry-instrumentation-mysql >= 0.42b0"]
 sqlite3 = ["opentelemetry-instrumentation-sqlite3 >= 0.42b0"]
+aws-lambda = ["opentelemetry-instrumentation-aws-lambda >= 0.42b0"]
 
 [project.scripts]
 logfire = "logfire.cli:main"
@@ -106,25 +107,26 @@ dev = [
     "pandas; python_version < '3.9'",
     "attrs",
     "openai",
-    "opentelemetry-instrumentation-aiohttp-client",
-    "opentelemetry-instrumentation-asgi",
-    "opentelemetry-instrumentation-wsgi",
-    "opentelemetry-instrumentation-fastapi",
-    "opentelemetry-instrumentation-starlette",
-    "opentelemetry-instrumentation-flask",
-    "opentelemetry-instrumentation-django",
-    "opentelemetry-instrumentation-httpx",
-    "opentelemetry-instrumentation-requests",
-    "opentelemetry-instrumentation-sqlalchemy",
-    "opentelemetry-instrumentation-system-metrics",
-    "opentelemetry-instrumentation-asyncpg",
-    "opentelemetry-instrumentation-psycopg",
-    "opentelemetry-instrumentation-psycopg2",
-    "opentelemetry-instrumentation-redis",
-    "opentelemetry-instrumentation-pymongo",
-    "opentelemetry-instrumentation-celery",
-    "opentelemetry-instrumentation-mysql",
-    "opentelemetry-instrumentation-sqlite3",
+    "opentelemetry-instrumentation-aiohttp-client>=0.42b0",
+    "opentelemetry-instrumentation-asgi>=0.42b0",
+    "opentelemetry-instrumentation-wsgi>=0.42b0",
+    "opentelemetry-instrumentation-fastapi>=0.42b0",
+    "opentelemetry-instrumentation-starlette>=0.42b0",
+    "opentelemetry-instrumentation-flask>=0.42b0",
+    "opentelemetry-instrumentation-django>=0.42b0",
+    "opentelemetry-instrumentation-httpx>=0.42b0",
+    "opentelemetry-instrumentation-requests>=0.42b0",
+    "opentelemetry-instrumentation-sqlalchemy>=0.42b0",
+    "opentelemetry-instrumentation-system-metrics>=0.42b0",
+    "opentelemetry-instrumentation-asyncpg>=0.42b0",
+    "opentelemetry-instrumentation-psycopg>=0.42b0",
+    "opentelemetry-instrumentation-psycopg2>=0.42b0",
+    "opentelemetry-instrumentation-redis>=0.42b0",
+    "opentelemetry-instrumentation-pymongo>=0.42b0",
+    "opentelemetry-instrumentation-celery>=0.42b0",
+    "opentelemetry-instrumentation-mysql>=0.42b0",
+    "opentelemetry-instrumentation-sqlite3>=0.42b0",
+    "opentelemetry-instrumentation-aws-lambda>=0.42b0",
     "eval-type-backport",
     "requests-mock",
     "inline-snapshot",
@@ -170,6 +172,9 @@ docs = [
 [tool.uv.sources]
 logfire-api = { workspace = true }
 
+[tool.uv]
+default-groups = ["dev", "docs"]
+
 [tool.uv.workspace]
 members = ["logfire-api"]
 

tests/otel_integrations/test_aws_lambda.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import importlib
+from dataclasses import dataclass
+from typing import Any
+from unittest import mock
+
+import pytest
+from inline_snapshot import snapshot
+from opentelemetry.instrumentation.aws_lambda import _HANDLER  # type: ignore[import]
+
+import logfire
+import logfire._internal.integrations.aws_lambda
+import logfire._internal.integrations.pymongo
+from logfire.testing import TestExporter
+
+
+def lambda_handler(event: Any, context: MockLambdaContext):
+    pass
+
+
+# The below mock is based on the following code:
+# https://github.com/open-telemetry/opentelemetry-python-contrib/blob/ecf5529f99222e7d945eddcaa83acb8a47c9ba42/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py#L57-L66
+@dataclass
+class MockLambdaContext:
+    aws_request_id: str
+    invoked_function_arn: str
+
+
+def test_instrument_aws_lambda(exporter: TestExporter) -> None:
+    with mock.patch.dict('os.environ', {_HANDLER: 'tests.otel_integrations.test_aws_lambda.lambda_handler'}):
+        logfire.instrument_aws_lambda(lambda_handler)
+
+        context = MockLambdaContext(
+            aws_request_id='mock_aws_request_id',
+            invoked_function_arn='arn:aws:lambda:us-east-1:123456:function:myfunction:myalias',
+        )
+        lambda_handler({'key': 'value'}, context)
+
+    assert exporter.exported_spans_as_dict() == snapshot(
+        [
+            {
+                'name': 'tests.otel_integrations.test_aws_lambda.lambda_handler',
+                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
+                'parent': None,
+                'start_time': 1000000000,
+                'end_time': 2000000000,
+                'attributes': {
+                    'logfire.span_type': 'span',
+                    'logfire.msg': 'tests.otel_integrations.test_aws_lambda.lambda_handler',
+                    'cloud.resource_id': 'arn:aws:lambda:us-east-1:123456:function:myfunction:myalias',
+                    'faas.invocation_id': 'mock_aws_request_id',
+                    'cloud.account.id': '123456',
+                },
+            }
+        ]
+    )
+
+
+def test_missing_opentelemetry_dependency() -> None:
+    with mock.patch.dict('sys.modules', {'opentelemetry.instrumentation.aws_lambda': None}):
+        with pytest.raises(RuntimeError) as exc_info:
+            importlib.reload(logfire._internal.integrations.aws_lambda)
+        assert str(exc_info.value) == snapshot("""\
+`logfire.instrument_aws_lambda()` requires the `opentelemetry-instrumentation-aws-lambda` package.
+You can install this with:
+    pip install 'logfire[aws-lambda]'\
+""")

tests/test_logfire_api.py

@@ -131,6 +131,10 @@ def func() -> None: ...
     func()
     logfire__all__.remove('instrument')
 
+    assert hasattr(logfire_api, 'instrument_aws_lambda'), 'instrument_aws_lambda'
+    logfire_api.instrument_aws_lambda(lambda_function=MagicMock())
+    logfire__all__.remove('instrument_aws_lambda')
+
     assert hasattr(logfire_api, 'instrument_asgi'), 'instrument_asgi'
     assert getattr(logfire_api, 'instrument_asgi')(app=MagicMock()) is not None
     logfire__all__.remove('instrument_asgi')