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 copy
8# of this software and associated documentation files (the ""Software""), to
9# deal in the Software without restriction, including without limitation the
10# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
11# sell copies of the Software, and to permit persons to whom the Software is
12# furnished to do so, subject to the following conditions:
13#
14# The above copyright notice and this permission notice shall be included in
15# all copies or substantial portions of the Software.
16#
17# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
22# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
23# IN THE SOFTWARE.
24#
25# --------------------------------------------------------------------------
26from typing import TypeVar, Any
27from azure.core.pipeline import PipelineRequest, PipelineResponse
28from azure.core.pipeline.transport import (
29 HttpResponse as LegacyHttpResponse,
30 HttpRequest as LegacyHttpRequest,
31)
32from azure.core.rest import HttpResponse, HttpRequest
33from ._base import SansIOHTTPPolicy
34
35HTTPResponseType = TypeVar("HTTPResponseType", HttpResponse, LegacyHttpResponse)
36HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest)
37
38
39class CustomHookPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]):
40 """A simple policy that enable the given callback
41 with the response.
42
43 :keyword callback raw_request_hook: Callback function. Will be invoked on request.
44 :keyword callback raw_response_hook: Callback function. Will be invoked on response.
45 """
46
47 def __init__(self, **kwargs: Any):
48 self._request_callback = kwargs.get("raw_request_hook")
49 self._response_callback = kwargs.get("raw_response_hook")
50
51 def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
52 """This is executed before sending the request to the next policy.
53
54 :param request: The PipelineRequest object.
55 :type request: ~azure.core.pipeline.PipelineRequest
56 """
57 request_callback = request.context.options.pop("raw_request_hook", None)
58 if request_callback:
59 request.context["raw_request_hook"] = request_callback
60 request_callback(request)
61 elif self._request_callback:
62 self._request_callback(request)
63
64 response_callback = request.context.options.pop("raw_response_hook", None)
65 if response_callback:
66 request.context["raw_response_hook"] = response_callback
67
68 def on_response(
69 self,
70 request: PipelineRequest[HTTPRequestType],
71 response: PipelineResponse[HTTPRequestType, HTTPResponseType],
72 ) -> None:
73 """This is executed after the request comes back from the policy.
74
75 :param request: The PipelineRequest object.
76 :type request: ~azure.core.pipeline.PipelineRequest
77 :param response: The PipelineResponse object.
78 :type response: ~azure.core.pipeline.PipelineResponse
79 """
80 response_callback = response.context.get("raw_response_hook")
81 if response_callback:
82 response_callback(response)
83 elif self._response_callback:
84 self._response_callback(response)