Skip to content

Commit bafb232

Browse files
committed
added context manager / decorator for transactions [WIP]
1 parent f6ba092 commit bafb232

File tree

3 files changed

+72
-1
lines changed

3 files changed

+72
-1
lines changed

elasticapm/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
from elasticapm.instrumentation.control import instrument, uninstrument # noqa: F401
3535
from elasticapm.traces import ( # noqa: F401
3636
capture_span,
37+
transaction,
3738
get_span_id,
3839
get_trace_id,
3940
get_trace_parent_header,

elasticapm/traces.py

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -701,7 +701,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
701701

702702
if transaction and transaction.is_sampled:
703703
try:
704-
outcome = "failure" if exc_val else "success"
704+
outcome = constants.OUTCOME.FAILURE if exc_val else constants.OUTCOME.SUCCESS
705705
span = transaction.end_span(self.skip_frames, duration=self.duration, outcome=outcome)
706706
if exc_val and not isinstance(span, DroppedSpan):
707707
try:
@@ -713,6 +713,43 @@ def __exit__(self, exc_type, exc_val, exc_tb):
713713
logger.debug("ended non-existing span %s of type %s", self.name, self.type)
714714

715715

716+
class transaction(object):
717+
def __init__(self, name=None, transaction_type=None, trace_parent=None):
718+
self.name = name
719+
self.transaction_type = transaction_type or "code"
720+
self.trace_parent = trace_parent
721+
722+
def __call__(self, func):
723+
self.name = self.name or get_name_from_func(func)
724+
725+
@functools.wraps(func)
726+
def decorated(*args, **kwds):
727+
with self:
728+
return func(*args, **kwds)
729+
730+
return decorated
731+
732+
def __enter__(self):
733+
from elasticapm.base import get_client
734+
735+
client = get_client()
736+
if client is None:
737+
logger.debug("No active client instance found while tracing transaction %s", str(self.name))
738+
return None
739+
return client.begin_transaction(self.transaction_type, self.trace_parent)
740+
741+
def __exit__(self, exc_type, exc_val, exc_tb):
742+
from elasticapm.base import get_client
743+
744+
client = get_client()
745+
if client is None:
746+
logger.debug("No active client instance found while tracing transaction %s", str(self.name))
747+
return None
748+
outcome = constants.OUTCOME.FAILURE if exc_val else constants.OUTCOME.SUCCESS
749+
set_transaction_outcome(outcome)
750+
client.end_transaction(self.name)
751+
752+
716753
def label(**labels):
717754
"""
718755
Labels current transaction. Keys should be strings, values can be strings, booleans,

tests/client/client_tests.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import elasticapm
4949
from elasticapm.base import Client
5050
from elasticapm.conf.constants import ERROR, KEYWORD_MAX_LENGTH, SPAN, TRANSACTION
51+
from elasticapm.traces import transaction
5152
from elasticapm.utils import compat, encoding
5253
from elasticapm.utils.disttracing import TraceParent
5354
from tests.fixtures import DummyTransport, TempStoreClient
@@ -894,3 +895,35 @@ def test_excepthook(elasticapm_client):
894895
elasticapm_client._excepthook(type_, value, traceback)
895896

896897
assert elasticapm_client.events[ERROR]
898+
899+
900+
def test_transaction_as_decorator(elasticapm_client):
901+
@elasticapm.transaction()
902+
def decorated_function():
903+
with elasticapm.capture_span("bla"):
904+
pass
905+
906+
decorated_function()
907+
transaction = elasticapm_client.events["transaction"][0]
908+
assert transaction["name"] == "tests.client.client_tests.decorated_function"
909+
assert transaction["type"] == "code"
910+
assert transaction["outcome"] == "success"
911+
assert transaction["span_count"]["started"] == 1
912+
913+
914+
def test_transaction_as_context_manager(elasticapm_client):
915+
with elasticapm.transaction(name="foo"):
916+
with elasticapm.capture_span("bla"):
917+
pass
918+
transaction = elasticapm_client.events["transaction"][0]
919+
assert transaction["name"] == "foo"
920+
assert transaction["type"] == "code"
921+
assert transaction["outcome"] == "success"
922+
assert transaction["span_count"]["started"] == 1
923+
924+
925+
def test_transaction_context_manager_with_no_client():
926+
with elasticapm.transaction(name="foo") as t:
927+
with elasticapm.capture_span("bla"):
928+
pass
929+
assert t is None

0 commit comments

Comments
 (0)