make @trace decorator support async functions (#92)
Co-authored-by: Zhenxu Ke <kezhenxu94@apache.org>
diff --git a/skywalking/decorators.py b/skywalking/decorators.py
index a48e057..444c7a5 100644
--- a/skywalking/decorators.py
+++ b/skywalking/decorators.py
@@ -17,6 +17,7 @@
from functools import wraps
from typing import List
+import inspect
from skywalking import Layer, Component
from skywalking.trace.context import get_context
@@ -30,22 +31,37 @@
tags: List[Tag] = None,
):
def decorator(func):
- @wraps(func)
- def wrapper(*args, **kwargs):
- _op = op or func.__name__
- context = get_context()
- with context.new_local_span(op=_op) as span:
- span.layer = layer
- span.component = component
- [span.tag(tag) for tag in tags or []]
- try:
- result = func(*args, **kwargs)
- return result
- except Exception:
- span.raised()
- raise
+ _op = op or func.__name__
+ context = get_context()
+ if inspect.iscoroutinefunction(func):
+ @wraps(func)
+ async def wrapper(*args, **kwargs):
+ with context.new_local_span(op=_op) as span:
+ span.layer = layer
+ span.component = component
+ [span.tag(tag) for tag in tags or []]
+ try:
+ result = func(*args, **kwargs)
+ return await result
+ except Exception:
+ span.raised()
+ raise
+ return wrapper
- return wrapper
+ else:
+ @wraps(func)
+ def wrapper(*args, **kwargs):
+ with context.new_local_span(op=_op) as span:
+ span.layer = layer
+ span.component = component
+ [span.tag(tag) for tag in tags or []]
+ try:
+ result = func(*args, **kwargs)
+ return result
+ except Exception:
+ span.raised()
+ raise
+ return wrapper
return decorator
