|
| 1 | +import logging |
| 2 | + |
| 3 | +import httpx |
| 4 | +import wrapt |
| 5 | + |
| 6 | +from opencensus.trace import ( |
| 7 | + attributes_helper, |
| 8 | + exceptions_status, |
| 9 | + execution_context, |
| 10 | + integrations, |
| 11 | +) |
| 12 | +from opencensus.trace import span as span_module |
| 13 | +from opencensus.trace import utils |
| 14 | + |
| 15 | +try: |
| 16 | + from urllib.parse import urlparse |
| 17 | +except ImportError: |
| 18 | + from urlparse import urlparse |
| 19 | + |
| 20 | + |
| 21 | +log = logging.getLogger(__name__) |
| 22 | + |
| 23 | +MODULE_NAME = "httpx" |
| 24 | + |
| 25 | +HTTP_HOST = attributes_helper.COMMON_ATTRIBUTES["HTTP_HOST"] |
| 26 | +HTTP_METHOD = attributes_helper.COMMON_ATTRIBUTES["HTTP_METHOD"] |
| 27 | +HTTP_PATH = attributes_helper.COMMON_ATTRIBUTES["HTTP_PATH"] |
| 28 | +HTTP_ROUTE = attributes_helper.COMMON_ATTRIBUTES["HTTP_ROUTE"] |
| 29 | +HTTP_STATUS_CODE = attributes_helper.COMMON_ATTRIBUTES["HTTP_STATUS_CODE"] |
| 30 | +HTTP_URL = attributes_helper.COMMON_ATTRIBUTES["HTTP_URL"] |
| 31 | + |
| 32 | + |
| 33 | +def trace_integration(tracer=None): |
| 34 | + """Wrap the requests library to trace it.""" |
| 35 | + log.info("Integrated module: {}".format(MODULE_NAME)) |
| 36 | + |
| 37 | + if tracer is not None: |
| 38 | + # The execution_context tracer should never be None - if it has not |
| 39 | + # been set it returns a no-op tracer. Most code in this library does |
| 40 | + # not handle None being used in the execution context. |
| 41 | + execution_context.set_opencensus_tracer(tracer) |
| 42 | + |
| 43 | + wrapt.wrap_function_wrapper(MODULE_NAME, "Client.request", wrap_client_request) |
| 44 | + # pylint: disable=protected-access |
| 45 | + integrations.add_integration(integrations._Integrations.HTTPX) |
| 46 | + |
| 47 | + |
| 48 | +def wrap_client_request(wrapped, instance, args, kwargs): |
| 49 | + """Wrap the session function to trace it.""" |
| 50 | + # Check if request was sent from an exporter. If so, do not wrap. |
| 51 | + if execution_context.is_exporter(): |
| 52 | + return wrapped(*args, **kwargs) |
| 53 | + |
| 54 | + method = kwargs.get("method") or args[0] |
| 55 | + url = kwargs.get("url") or args[1] |
| 56 | + |
| 57 | + excludelist_hostnames = execution_context.get_opencensus_attr( |
| 58 | + "excludelist_hostnames" |
| 59 | + ) |
| 60 | + parsed_url = urlparse(url) |
| 61 | + if parsed_url.port is None: |
| 62 | + dest_url = parsed_url.hostname |
| 63 | + else: |
| 64 | + dest_url = "{}:{}".format(parsed_url.hostname, parsed_url.port) |
| 65 | + if utils.disable_tracing_hostname(dest_url, excludelist_hostnames): |
| 66 | + return wrapped(*args, **kwargs) |
| 67 | + |
| 68 | + path = parsed_url.path if parsed_url.path else "/" |
| 69 | + |
| 70 | + _tracer = execution_context.get_opencensus_tracer() |
| 71 | + _span = _tracer.start_span() |
| 72 | + |
| 73 | + _span.name = "{}".format(path) |
| 74 | + _span.span_kind = span_module.SpanKind.CLIENT |
| 75 | + |
| 76 | + try: |
| 77 | + tracer_headers = _tracer.propagator.to_headers(_tracer.span_context) |
| 78 | + kwargs.setdefault("headers", {}).update(tracer_headers) |
| 79 | + except Exception: # pragma: NO COVER |
| 80 | + pass |
| 81 | + |
| 82 | + # Add the component type to attributes |
| 83 | + _tracer.add_attribute_to_current_span("component", "HTTP") |
| 84 | + |
| 85 | + # Add the requests host to attributes |
| 86 | + _tracer.add_attribute_to_current_span(HTTP_HOST, dest_url) |
| 87 | + |
| 88 | + # Add the requests method to attributes |
| 89 | + _tracer.add_attribute_to_current_span(HTTP_METHOD, method.upper()) |
| 90 | + |
| 91 | + # Add the requests path to attributes |
| 92 | + _tracer.add_attribute_to_current_span(HTTP_PATH, path) |
| 93 | + |
| 94 | + # Add the requests url to attributes |
| 95 | + _tracer.add_attribute_to_current_span(HTTP_URL, url) |
| 96 | + |
| 97 | + try: |
| 98 | + result = wrapped(*args, **kwargs) |
| 99 | + except httpx.TimeoutException: |
| 100 | + _span.set_status(exceptions_status.TIMEOUT) |
| 101 | + raise |
| 102 | + except httpx.InvalidURL: |
| 103 | + _span.set_status(exceptions_status.INVALID_URL) |
| 104 | + raise |
| 105 | + except Exception as e: |
| 106 | + _span.set_status(exceptions_status.unknown(e)) |
| 107 | + raise |
| 108 | + else: |
| 109 | + # Add the status code to attributes |
| 110 | + _tracer.add_attribute_to_current_span(HTTP_STATUS_CODE, result.status_code) |
| 111 | + _span.set_status(utils.status_from_http_code(result.status_code)) |
| 112 | + return result |
| 113 | + finally: |
| 114 | + _tracer.end_span() |
0 commit comments