adds support for modern xray patching

fleece/log.py

@@ -203,11 +203,11 @@ def get_logger(name=None, level=None, stream=DEFAULT_STREAM,
     return log
 
 
-def initial_trace_and_context_binds(logger, trace_id, lambda_context):
+def get_initial_trace_and_context_binds(trace_id, lambda_context):
     """A helper to set up standard trace_id and lambda_context binds"""
-    return logger.new(
-        trace_id=trace_id,
-        lambda_context={
+    return {
+        "trace_id": trace_id,
+        "lambda_context": {
             "function_name": lambda_context.function_name,
             "function_version": lambda_context.function_version,
             "invoked_function_arn": lambda_context.invoked_function_arn,
@@ -216,6 +216,13 @@ def initial_trace_and_context_binds(logger, trace_id, lambda_context):
             "log_group_name": lambda_context.log_group_name,
             "log_stream_name": lambda_context.log_stream_name,
         }
+    }
+
+
+def initial_trace_and_context_binds(logger, trace_id, lambda_context):
+    """A helper to set up standard trace_id and lambda_context binds"""
+    return logger.new(
+        **get_initial_trace_and_context_binds(trace_id, lambda_context)
     )
 
 

fleece/xray.py

@@ -21,6 +21,10 @@
 import time
 import uuid
 
+from aws_xray_sdk.core import patch as aws_xray_patch
+from aws_xray_sdk.core import patch_all
+from aws_xray_sdk.core import xray_recorder
+
 from botocore.exceptions import ClientError
 import wrapt
 
@@ -496,3 +500,25 @@ def to_safe_annotation_key(key):
 def to_safe_annotation_value(value):
     """Xray doesn't like values that are not strings."""
     return str(value)
+
+
+def patch(modules=None):  # type: (t.Optional[t.List[str]]) -> None
+    """
+    Patch known external packages, (requests, pynamo, etc) along with the
+    given set of modules.
+    """
+    if os.environ.get("AWS_XRAY_DAEMON_ADDRESS") is not None:
+        patch_all()
+        if modules:
+            aws_xray_patch(modules)
+
+
+def log_args(lambda_context):  # type: (t.Any) -> t.Mapping
+    """
+    Returns arguments that should be bound to a log (via .new or .bind) when a
+    lambda handler is first called.
+    """
+    return log.get_initial_trace_and_context_binds(
+        trace_id=xray_recorder.get_trace_entity().trace_id,
+        lambda_context=lambda_context
+    )

tests/test_xray.py

@@ -5,6 +5,7 @@
 import mock
 
 from fleece import xray
+from fleece import testing
 
 
 class GetTraceIDTestCase(unittest.TestCase):
@@ -124,3 +125,72 @@ def test_full_segment_document_with_extra_data(self):
         self.assertEqual(segment_document['end_time'], end_time)
         self.assertEqual(segment_document['name'], 'NAME')
         self.assertEqual(segment_document['foo'], 'BAR')
+
+
+class TestPatch(unittest.TestCase):
+    def setUp(self):
+        self.patch = mock.patch('fleece.xray.aws_xray_patch').start()
+        self.patch_all = mock.patch('fleece.xray.patch_all').start()
+
+    def tearDown(self):
+        self.patch.stop()
+        self.patch_all.stop()
+
+    def test_patch_when_not_in_lambda(self):
+        xray.patch()
+        self.patch_all.assert_not_called()
+        self.patch.assert_not_called()
+
+    @mock.patch.dict(
+        os.environ,
+        {
+            'AWS_XRAY_DAEMON_ADDRESS': 'http://localhost',  # noqa
+        }
+    )
+    def test_patch_when_in_lambda(self):
+        xray.patch()
+        self.patch_all.assert_called_once()
+        self.patch.assert_not_called()
+
+    @mock.patch.dict(
+        os.environ,
+        {
+            'AWS_XRAY_DAEMON_ADDRESS': 'http://localhost',  # noqa
+        }
+    )
+    def test_patch_when_in_lambda_2(self):
+        xray.patch(["fleece.utils"])
+        self.patch_all.assert_called_once()
+        self.patch.assert_called_once()
+
+
+class TestLogArgs(unittest.TestCase):
+    def setUp(self):
+        self.maxDiff = None
+        self.get_trace_entity = mock.patch(
+            'fleece.xray.xray_recorder.get_trace_entity').start()
+
+    def tearDown(self):
+        self.get_trace_entity.stop()
+
+    def test_patch_when_not_in_lambda(self):
+        class Trace:
+            def __init__(self):
+                self.trace_id = "trace-id"
+
+        self.get_trace_entity.return_value = Trace()
+        lambda_context = testing.LambdaContext("FakeFunction")
+        actual = xray.log_args(lambda_context)
+        expected = {
+            "trace_id": 'trace-id',
+            "lambda_context": {
+                "function_name": lambda_context.function_name,
+                "function_version": lambda_context.function_version,
+                "invoked_function_arn": lambda_context.invoked_function_arn,
+                "memory_limit_in_mb": lambda_context.memory_limit_in_mb,
+                "aws_request_id": lambda_context.aws_request_id,
+                "log_group_name": lambda_context.log_group_name,
+                "log_stream_name": lambda_context.log_stream_name,
+            }
+        }
+        self.assertEqual(expected, actual)