|
1 | 1 | import contextlib |
| 2 | +import functools |
2 | 3 | import inspect |
3 | 4 | import os |
4 | 5 | import re |
5 | 6 | import sys |
6 | 7 | from collections.abc import Mapping |
7 | 8 | from datetime import timedelta |
8 | 9 | from decimal import ROUND_DOWN, Decimal, DefaultContext, localcontext |
9 | | -from functools import wraps |
10 | 10 | from random import Random |
11 | 11 | from urllib.parse import quote, unquote |
12 | 12 | import uuid |
@@ -776,70 +776,86 @@ def normalize_incoming_data(incoming_data): |
776 | 776 | return data |
777 | 777 |
|
778 | 778 |
|
779 | | -def start_child_span_decorator(func): |
780 | | - # type: (Any) -> Any |
| 779 | +def create_span_decorator(op=None, name=None, attributes=None): |
| 780 | + # type: (Optional[str], Optional[str], Optional[dict[str, Any]]) -> Any |
781 | 781 | """ |
782 | | - Decorator to add child spans for functions. |
| 782 | + Create a span decorator that can wrap both sync and async functions. |
783 | 783 |
|
784 | | - See also ``sentry_sdk.tracing.trace()``. |
| 784 | + :param op: The operation type for the span. |
| 785 | + :param name: The name of the span. |
| 786 | + :param attributes: Additional attributes to set on the span. |
785 | 787 | """ |
786 | | - # Asynchronous case |
787 | | - if inspect.iscoroutinefunction(func): |
788 | 788 |
|
789 | | - @wraps(func) |
790 | | - async def func_with_tracing(*args, **kwargs): |
791 | | - # type: (*Any, **Any) -> Any |
| 789 | + def span_decorator(f): |
| 790 | + # type: (Any) -> Any |
| 791 | + """ |
| 792 | + Decorator to create a span for the given function. |
| 793 | + """ |
792 | 794 |
|
793 | | - span = get_current_span() |
| 795 | + @functools.wraps(f) |
| 796 | + async def async_wrapper(*args, **kwargs): |
| 797 | + # type: (*Any, **Any) -> Any |
| 798 | + current_span = get_current_span() |
794 | 799 |
|
795 | | - if span is None: |
| 800 | + if current_span is None: |
796 | 801 | logger.debug( |
797 | 802 | "Cannot create a child span for %s. " |
798 | 803 | "Please start a Sentry transaction before calling this function.", |
799 | | - qualname_from_function(func), |
| 804 | + qualname_from_function(f), |
800 | 805 | ) |
801 | | - return await func(*args, **kwargs) |
| 806 | + return await f(*args, **kwargs) |
| 807 | + |
| 808 | + span_op = op or OP.FUNCTION |
| 809 | + span_name = name or qualname_from_function(f) or "" |
802 | 810 |
|
803 | | - with span.start_child( |
804 | | - op=OP.FUNCTION, |
805 | | - name=qualname_from_function(func), |
806 | | - ): |
807 | | - return await func(*args, **kwargs) |
| 811 | + with current_span.start_child( |
| 812 | + op=span_op, |
| 813 | + name=span_name, |
| 814 | + ) as span: |
| 815 | + span.update_data(attributes or {}) |
| 816 | + result = await f(*args, **kwargs) |
| 817 | + return result |
808 | 818 |
|
809 | 819 | try: |
810 | | - func_with_tracing.__signature__ = inspect.signature(func) # type: ignore[attr-defined] |
| 820 | + async_wrapper.__signature__ = inspect.signature(f) # type: ignore[attr-defined] |
811 | 821 | except Exception: |
812 | 822 | pass |
813 | 823 |
|
814 | | - # Synchronous case |
815 | | - else: |
816 | | - |
817 | | - @wraps(func) |
818 | | - def func_with_tracing(*args, **kwargs): |
| 824 | + @functools.wraps(f) |
| 825 | + def sync_wrapper(*args, **kwargs): |
819 | 826 | # type: (*Any, **Any) -> Any |
| 827 | + current_span = get_current_span() |
820 | 828 |
|
821 | | - span = get_current_span() |
822 | | - |
823 | | - if span is None: |
| 829 | + if current_span is None: |
824 | 830 | logger.debug( |
825 | 831 | "Cannot create a child span for %s. " |
826 | 832 | "Please start a Sentry transaction before calling this function.", |
827 | | - qualname_from_function(func), |
| 833 | + qualname_from_function(f), |
828 | 834 | ) |
829 | | - return func(*args, **kwargs) |
| 835 | + return f(*args, **kwargs) |
| 836 | + |
| 837 | + span_op = op or OP.FUNCTION |
| 838 | + span_name = name or qualname_from_function(f) or "" |
830 | 839 |
|
831 | | - with span.start_child( |
832 | | - op=OP.FUNCTION, |
833 | | - name=qualname_from_function(func), |
834 | | - ): |
835 | | - return func(*args, **kwargs) |
| 840 | + with current_span.start_child( |
| 841 | + op=span_op, |
| 842 | + name=span_name, |
| 843 | + ) as span: |
| 844 | + span.update_data(attributes or {}) |
| 845 | + result = f(*args, **kwargs) |
| 846 | + return result |
836 | 847 |
|
837 | 848 | try: |
838 | | - func_with_tracing.__signature__ = inspect.signature(func) # type: ignore[attr-defined] |
| 849 | + sync_wrapper.__signature__ = inspect.signature(f) # type: ignore[attr-defined] |
839 | 850 | except Exception: |
840 | 851 | pass |
841 | 852 |
|
842 | | - return func_with_tracing |
| 853 | + if inspect.iscoroutinefunction(f): |
| 854 | + return async_wrapper |
| 855 | + else: |
| 856 | + return sync_wrapper |
| 857 | + |
| 858 | + return span_decorator |
843 | 859 |
|
844 | 860 |
|
845 | 861 | def get_current_span(scope=None): |
|
0 commit comments