Check LambdaRuntimeClient before logging tracebacks in _ensure_flush_after_aws_lambda (#388)

Author: alexmojaki

logfire/_internal/config.py

@@ -8,6 +8,7 @@
 import sys
 import time
 import warnings
+from contextlib import suppress
 from dataclasses import dataclass, field
 from functools import cached_property
 from pathlib import Path
@@ -862,15 +863,26 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
                 client = getattr(mod, 'LambdaRuntimeClient', None)
             except Exception:  # pragma: no cover
                 continue
-            if not client:  # pragma: no branch
+            if not client:
                 continue
             try:  # pragma: no cover
                 client.post_invocation_error = wrap_client_post_invocation_method(client.post_invocation_error)
                 client.post_invocation_result = wrap_client_post_invocation_method(client.post_invocation_result)
-            except Exception:  # pragma: no cover
+            except Exception as e:  # pragma: no cover
+                with suppress(Exception):
+                    # client is likely some random object from a dynamic module unrelated to AWS lambda.
+                    # If it doesn't look like the LambdaRuntimeClient class, ignore this error.
+                    # We don't check this beforehand so that if the lambda runtime library changes
+                    # LambdaRuntimeClient to some object other than a class,
+                    # or something else patches it with some kind of wrapper,
+                    # our patching still has some chance of working.
+                    # But we also don't want to log spurious noisy tracebacks.
+                    if not (isinstance(client, type) and client.__name__ == 'LambdaRuntimeClient'):
+                        continue
+
                 import traceback
 
-                traceback.print_exc()
+                traceback.print_exception(e)
 
 
 def _get_default_span_processor(exporter: SpanExporter) -> SpanProcessor:

tests/import_used_for_tests/module_with_getattr.py

@@ -0,0 +1,5 @@
+# See test_dynamic_module_ignored_in_ensure_flush_after_aws_lambda
+
+
+def __getattr__(name: str) -> str:
+    return name

tests/test_auto_trace.py

@@ -138,14 +138,23 @@ def test_check_already_imported() -> None:
     # because we don't want anything to get auto-traced here.
     imported_modules = set(sys.modules.items())
 
+    meta_path = sys.meta_path.copy()
+
     with pytest.raises(AutoTraceModuleAlreadyImportedException, match=r"The module 'tests.*' matches modules to trace"):
         install_auto_tracing(['tests'])
 
+    with pytest.raises(ValueError):
+        install_auto_tracing(['tests'], check_imported_modules='other')  # type: ignore
+
+    # No tracing installed.
+    assert sys.meta_path == meta_path
+
     with pytest.warns(AutoTraceModuleAlreadyImportedWarning, match=r"The module 'tests.*' matches modules to trace"):
         install_auto_tracing(['tests'], check_imported_modules='warn')
 
-    with pytest.raises(ValueError):
-        install_auto_tracing(['tests'], check_imported_modules='other')  # type: ignore
+    # The tracing was installed, undo it.
+    assert sys.meta_path[1:] == meta_path
+    sys.meta_path = meta_path
 
     assert set(sys.modules.items()) == imported_modules
 

tests/test_configure.py

@@ -1452,6 +1452,18 @@ def get_metric_readers() -> Iterable[SpanProcessor]:
     return get_meter_provider().provider._sdk_config.metric_readers  # type: ignore
 
 
+def test_dynamic_module_ignored_in_ensure_flush_after_aws_lambda(
+    config_kwargs: dict[str, Any], capsys: pytest.CaptureFixture[str]
+):
+    from tests.import_used_for_tests.module_with_getattr import module_with_getattr_value
+
+    assert module_with_getattr_value == 'module_with_getattr_value'
+
+    logfire.configure(**config_kwargs)
+
+    assert capsys.readouterr().err == ''
+
+
 def test_collect_system_metrics_false():
     with inline_snapshot.extra.raises(
         snapshot(