from ylog.context import log_context
try:
    from werkzeug.wrappers import Request
except ImportError:
    # no werkzeug installed
    pass


__all__ = [
    'LogReqIdMiddleware',
]


class LogReqIdMiddleware(object):
    """ Adds request id from headers to log context """

    def __init__(self, app):
        super(LogReqIdMiddleware, self).__init__()
        self._app = app

    @staticmethod
    def get_req_id(environ):
        """
        Extracts request id from headers
        see https://doc.qloud.yandex-team.ru/doc/balancers
        """
        req = Request(environ, shallow=True)
        req_id = req.headers.get('X-Req-Id', None)
        req_id = req_id or req.headers.get('X-Request-Id', None)
        return req_id

    def __call__(self, environ, start_response):
        with log_context(request_id=self.get_req_id(environ)):
            return self._app(environ, start_response)
