traces_middleware.py 6.82 KB
Newer Older
ethiraj krishnamanaidu's avatar
ethiraj krishnamanaidu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2021 Schlumberger
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

15
16
17
18
19
from typing import Any

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response
Luc Yriarte's avatar
Luc Yriarte committed
20
from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
21
22

from opencensus.trace import tracer as open_tracer
Luc Yriarte's avatar
Luc Yriarte committed
23
24
from opencensus.trace.samplers import AlwaysOnSampler
from opencensus.trace.span import SpanKind
25
26

from app.helper import traces, utils
27
from app.utils import get_or_create_ctx
Luc Yriarte's avatar
Luc Yriarte committed
28
29
from app import conf
from inspect import isfunction as is_function
30
31
32
33
34


class TracingMiddleware(BaseHTTPMiddleware):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
Luc Yriarte's avatar
Luc Yriarte committed
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
        self._trace_propagator = traces.get_trace_propagator()

    @staticmethod
    def _retrieve_raw_path(request):
        """
        Returns the raw path of given request, else default request's url path
        E.g.:
            /ddms/v2/wellbores/{wellboreid} instead of /ddms/v2/wellbores/opendes:doc:blablabla14587

        It retrieves the raw path by finding the APIRoute object by name. By default the name of the route is the name
         of python method where there is the implementation.


        >>> @router.get('/wellbores/{wellboreid}')
        >>> async def get_wellbore(wellboreid: str, ctx: Context):
        >>>     # instructions here
        In this example 'get_wellbore' is called_endpoint_func variable, this function's name is needed to retrieve
        the APIRoute that contains the raw path.
        """
Luc Yriarte's avatar
Luc Yriarte committed
54
        called_endpoint_func = request.scope.get('endpoint')
Luc Yriarte's avatar
Luc Yriarte committed
55
56
57
58
59
60
61
62
63

        if called_endpoint_func and is_function(called_endpoint_func):
            function_name = called_endpoint_func.__name__
            called_routes = [route for route in request.app.routes
                             if route.name == function_name]
            if called_routes:
                return called_routes[0].path

        return request.url.path
64
65
66
67

    @staticmethod
    def _before_request(request: Request, tracer: open_tracer.Tracer):
        tracer.add_attribute_to_current_span(
68
            attribute_key=utils.HTTP_HOST,
69
70
            attribute_value=request.url.hostname)
        tracer.add_attribute_to_current_span(
71
            attribute_key=utils.HTTP_METHOD,
72
            attribute_value=request.method)
Luc Yriarte's avatar
Luc Yriarte committed
73
74

        tracer.add_attribute_to_current_span(
75
            attribute_key=utils.HTTP_ROUTE,
Luc Yriarte's avatar
Luc Yriarte committed
76
            attribute_value=request.url.path)
77
        tracer.add_attribute_to_current_span(
78
            attribute_key=utils.HTTP_PATH,
79
80
            attribute_value=str(request.url.path))
        tracer.add_attribute_to_current_span(
81
            attribute_key=utils.HTTP_URL,
82
83
            attribute_value=str(request.url))

Luc Yriarte's avatar
Luc Yriarte committed
84
85
86
87
88
89
        ctx_correlation_id = get_or_create_ctx().correlation_id
        correlation_id = ctx_correlation_id if ctx_correlation_id is not None \
            else request.headers.get(conf.CORRELATION_ID_HEADER_NAME)
        tracer.add_attribute_to_current_span(
            attribute_key=conf.CORRELATION_ID_HEADER_NAME,
            attribute_value=correlation_id)
90
91
92
93
    
        ctx_partition_id = get_or_create_ctx().partition_id
        partition_id = ctx_partition_id if ctx_partition_id is not None \
            else request.headers.get(conf.PARTITION_ID_HEADER_NAME)
94
95
96
        tracer.add_attribute_to_current_span(
            attribute_key=conf.PARTITION_ID_HEADER_NAME,
            attribute_value=partition_id)
Luc Yriarte's avatar
Luc Yriarte committed
97

98
99
100
        request_content_type = request.headers.get("Content-type")
        tracer.add_attribute_to_current_span(attribute_key="request.header Content-type",
                                             attribute_value=request_content_type)
101

102
103
104
105
        request_content_length = request.headers.get("Content-Length")
        tracer.add_attribute_to_current_span(attribute_key="request.header Content-length",
                                             attribute_value=request_content_length)

Alexandre Vincent's avatar
Alexandre Vincent committed
106
107
108
109
110
        app_id = request.headers.get(conf.APP_ID_HEADER_NAME)
        if app_id is not None:
            tracer.add_attribute_to_current_span(attribute_key=conf.APP_ID_HEADER_NAME,
                                                 attribute_value=app_id)

Luc Yriarte's avatar
Luc Yriarte committed
111
    @staticmethod
112
113
114
115
116
117
118
119
120
    def _after_request(request: Request, response: Response, tracer):

        status = response.status_code if response else HTTP_500_INTERNAL_SERVER_ERROR
        tracer.add_attribute_to_current_span(attribute_key=utils.HTTP_STATUS_CODE,
                                             attribute_value=status)

        tracer.add_attribute_to_current_span(attribute_key=utils.HTTP_ROUTE,
                                             attribute_value=TracingMiddleware._retrieve_raw_path(request))

121
122
123
124
        if response:
            response_content_type = response.headers.get("Content-type")
            tracer.add_attribute_to_current_span(attribute_key="response.header Content-type",
                                                 attribute_value=response_content_type)
125

126
127
128
            response_content_length = response.headers.get("Content-Length")
            tracer.add_attribute_to_current_span(attribute_key="response.header Content-length",
                                                 attribute_value=response_content_length)
Luc Yriarte's avatar
Luc Yriarte committed
129

130
131
    async def dispatch(self, request: Request, call_next: Any) -> Response:

Luc Yriarte's avatar
Luc Yriarte committed
132
133
        # Create tracing context, from headers if exists, else create a new one
        span_context = self._trace_propagator.from_headers(request.headers)
134

Luc Yriarte's avatar
Luc Yriarte committed
135
136
137
138
        tracer = open_tracer.Tracer(span_context=span_context,
                                    sampler=AlwaysOnSampler(),
                                    propagator=self._trace_propagator,
                                    exporter=request.app.trace_exporter)
139
140

        ctx = get_or_create_ctx()
Luc Yriarte's avatar
Luc Yriarte committed
141
142
143
        with tracer.span(request.url.path) as parent_span:
            parent_span.span_kind = SpanKind.SERVER
            ctx.set_current_with_value(tracer=tracer)
144
145

            self._before_request(request, tracer)
Luc Yriarte's avatar
Luc Yriarte committed
146
147
            ctx.logger.debug(f'Request start: {request.method} {request.url}')

148
            response = None
Luc Yriarte's avatar
Luc Yriarte committed
149
150
151
152
153
            try:
                response = await call_next(request)
                return response
            finally:
                status = response.status_code if response else HTTP_500_INTERNAL_SERVER_ERROR
154
                ctx.logger.info(utils.process_message(request, status))
155
                self._after_request(request, response, tracer)