codekingpro/portable-devtools

1# --------------------------------------------------------------------------2#3# Copyright (c) Microsoft Corporation. All rights reserved.4#5# The MIT License (MIT)6#7# Permission is hereby granted, free of charge, to any person obtaining a copy8# of this software and associated documentation files (the ""Software""), to9# deal in the Software without restriction, including without limitation the10# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or11# sell copies of the Software, and to permit persons to whom the Software is12# furnished to do so, subject to the following conditions:13#14# The above copyright notice and this permission notice shall be included in15# all copies or substantial portions of the Software.16#17# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING22# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS23# IN THE SOFTWARE.24#25# --------------------------------------------------------------------------26import asyncio27from collections.abc import AsyncIterator28import functools29import logging30from typing import Any, Callable, Optional, AsyncIterator as AsyncIteratorType31 32from oauthlib import oauth233import requests34from requests.models import CONTENT_CHUNK_SIZE35 36from ..exceptions import (37    TokenExpiredError,38    ClientRequestError,39    raise_with_traceback40)41from ..universal_http.async_requests import AsyncBasicRequestsHTTPSender42from . import AsyncHTTPSender, AsyncHTTPPolicy, Response, Request43from .requests import RequestsContext44 45 46_LOGGER = logging.getLogger(__name__)47 48 49class AsyncPipelineRequestsHTTPSender(AsyncHTTPSender):50    """Implements a basic Pipeline, that supports universal HTTP lib "requests" driver.51    """52 53    def __init__(self, universal_http_requests_driver: Optional[AsyncBasicRequestsHTTPSender]=None) -> None:54        self.driver = universal_http_requests_driver or AsyncBasicRequestsHTTPSender()55 56    async def __aenter__(self) -> 'AsyncPipelineRequestsHTTPSender':57        await self.driver.__aenter__()58        return self59 60    async def __aexit__(self, *exc_details):  # pylint: disable=arguments-differ61        await self.driver.__aexit__(*exc_details)62 63    async def close(self):64        await self.__aexit__()65 66    def build_context(self):67        # type: () -> RequestsContext68        return RequestsContext(69            session=self.driver.session,70        )71 72    async def send(self, request: Request, **kwargs) -> Response:73        """Send request object according to configuration.74 75        :param Request request: The request object to be sent.76        """77        if request.context is None:  # Should not happen, but make mypy happy and does not hurt78            request.context = self.build_context()79 80        if request.context.session is not self.driver.session:81            kwargs['session'] = request.context.session82 83        return Response(84            request,85            await self.driver.send(request.http_request, **kwargs)86        )87 88 89class AsyncRequestsCredentialsPolicy(AsyncHTTPPolicy):90    """Implementation of request-oauthlib except and retry logic.91    """92    def __init__(self, credentials):93        super(AsyncRequestsCredentialsPolicy, self).__init__()94        self._creds = credentials95 96    async def send(self, request, **kwargs):97        session = request.context.session98        try:99            self._creds.signed_session(session)100        except TypeError: # Credentials does not support session injection101            _LOGGER.warning("Your credentials class does not support session injection. Performance will not be at the maximum.")102            request.context.session = session = self._creds.signed_session()103 104        try:105            try:106                return await self.next.send(request, **kwargs)107            except (oauth2.rfc6749.errors.InvalidGrantError,108                    oauth2.rfc6749.errors.TokenExpiredError) as err:109                error = "Token expired or is invalid. Attempting to refresh."110                _LOGGER.warning(error)111 112            try:113                try:114                    self._creds.refresh_session(session)115                except TypeError: # Credentials does not support session injection116                    _LOGGER.warning("Your credentials class does not support session injection. Performance will not be at the maximum.")117                    request.context.session = session = self._creds.refresh_session()118 119                return await self.next.send(request, **kwargs)120            except (oauth2.rfc6749.errors.InvalidGrantError,121                    oauth2.rfc6749.errors.TokenExpiredError) as err:122                msg = "Token expired or is invalid."123                raise_with_traceback(TokenExpiredError, msg, err)124 125        except (requests.RequestException,126                oauth2.rfc6749.errors.OAuth2Error) as err:127            msg = "Error occurred in request."128            raise_with_traceback(ClientRequestError, msg, err)129 130