added context manager / decorator for transactions [WIP]

elasticapm/__init__.py

@@ -34,6 +34,7 @@
 from elasticapm.instrumentation.control import instrument, uninstrument  # noqa: F401
 from elasticapm.traces import (  # noqa: F401
     capture_span,
+    transaction,
     get_span_id,
     get_trace_id,
     get_trace_parent_header,

elasticapm/traces.py

@@ -705,7 +705,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
         if transaction and transaction.is_sampled:
             try:
-                outcome = "failure" if exc_val else "success"
+                outcome = constants.OUTCOME.FAILURE if exc_val else constants.OUTCOME.SUCCESS
                 span = transaction.end_span(self.skip_frames, duration=self.duration, outcome=outcome)
                 if exc_val and not isinstance(span, DroppedSpan):
                     try:
@@ -717,6 +717,43 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 logger.debug("ended non-existing span %s of type %s", self.name, self.type)
 
 
+class transaction(object):
+    def __init__(self, name=None, transaction_type=None, trace_parent=None):
+        self.name = name
+        self.transaction_type = transaction_type or "code"
+        self.trace_parent = trace_parent
+
+    def __call__(self, func):
+        self.name = self.name or get_name_from_func(func)
+
+        @functools.wraps(func)
+        def decorated(*args, **kwds):
+            with self:
+                return func(*args, **kwds)
+
+        return decorated
+
+    def __enter__(self):
+        from elasticapm.base import get_client
+
+        client = get_client()
+        if client is None:
+            logger.debug("No active client instance found while tracing transaction %s", str(self.name))
+            return None
+        return client.begin_transaction(self.transaction_type, self.trace_parent)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        from elasticapm.base import get_client
+
+        client = get_client()
+        if client is None:
+            logger.debug("No active client instance found while tracing transaction %s", str(self.name))
+            return None
+        outcome = constants.OUTCOME.FAILURE if exc_val else constants.OUTCOME.SUCCESS
+        set_transaction_outcome(outcome)
+        client.end_transaction(self.name)
+
+
 def label(**labels):
     """
     Labels current transaction. Keys should be strings, values can be strings, booleans,

tests/client/client_tests.py

@@ -48,6 +48,7 @@
 import elasticapm
 from elasticapm.base import Client
 from elasticapm.conf.constants import ERROR, KEYWORD_MAX_LENGTH, SPAN, TRANSACTION
+from elasticapm.traces import transaction
 from elasticapm.utils import compat, encoding
 from elasticapm.utils.disttracing import TraceParent
 from tests.fixtures import DummyTransport, TempStoreClient
@@ -894,3 +895,35 @@ def test_excepthook(elasticapm_client):
         elasticapm_client._excepthook(type_, value, traceback)
 
     assert elasticapm_client.events[ERROR]
+
+
+def test_transaction_as_decorator(elasticapm_client):
+    @elasticapm.transaction()
+    def decorated_function():
+        with elasticapm.capture_span("bla"):
+            pass
+
+    decorated_function()
+    transaction = elasticapm_client.events["transaction"][0]
+    assert transaction["name"] == "tests.client.client_tests.decorated_function"
+    assert transaction["type"] == "code"
+    assert transaction["outcome"] == "success"
+    assert transaction["span_count"]["started"] == 1
+
+
+def test_transaction_as_context_manager(elasticapm_client):
+    with elasticapm.transaction(name="foo"):
+        with elasticapm.capture_span("bla"):
+            pass
+    transaction = elasticapm_client.events["transaction"][0]
+    assert transaction["name"] == "foo"
+    assert transaction["type"] == "code"
+    assert transaction["outcome"] == "success"
+    assert transaction["span_count"]["started"] == 1
+
+
+def test_transaction_context_manager_with_no_client():
+    with elasticapm.transaction(name="foo") as t:
+        with elasticapm.capture_span("bla"):
+            pass
+    assert t is None