From bafb23251930a439c61d6347e5bef8a33b0f46b7 Mon Sep 17 00:00:00 2001
From: Benjamin Wohlwend <beni@elastic.co>
Date: Wed, 7 Apr 2021 15:27:52 +0200
Subject: [PATCH] added context manager / decorator for transactions [WIP]

---
 elasticapm/__init__.py       |  1 +
 elasticapm/traces.py         | 39 +++++++++++++++++++++++++++++++++++-
 tests/client/client_tests.py | 33 ++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/elasticapm/__init__.py b/elasticapm/__init__.py
index 42bbca423..b8cee9992 100644
--- a/elasticapm/__init__.py
+++ b/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,
diff --git a/elasticapm/traces.py b/elasticapm/traces.py
index f78dae996..8d666c6a2 100644
--- a/elasticapm/traces.py
+++ b/elasticapm/traces.py
@@ -701,7 +701,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:
@@ -713,6 +713,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,
diff --git a/tests/client/client_tests.py b/tests/client/client_tests.py
index 46024e711..5c5b103cf 100644
--- a/tests/client/client_tests.py
+++ b/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