|
9 | 9 | import time |
10 | 10 | from concurrent.futures import ThreadPoolExecutor, as_completed |
11 | 11 | from functools import partial, reduce, wraps |
12 | | -from http.client import HTTPConnection, HTTPResponse |
13 | 12 | from typing import ( |
14 | 13 | IO, |
15 | 14 | Any, |
|
37 | 36 | from boto3.dynamodb.types import TypeDeserializer |
38 | 37 | from boto3.s3.transfer import TransferConfig |
39 | 38 | from boto3.session import Session |
40 | | -from botocore import endpoint |
41 | 39 | from botocore.config import Config |
42 | 40 | from botocore.exceptions import ClientError, UnknownServiceError |
43 | 41 | from botocore.response import StreamingBody |
|
83 | 81 | s3re = re.compile(r"^(s3[an]?)://([^/]+)/(.+)$", re.IGNORECASE) |
84 | 82 |
|
85 | 83 |
|
86 | | -class CustomHTTPResponse(HTTPResponse): |
87 | | - def __init__(self, sock, **kwargs): |
88 | | - """ |
89 | | - Capture the local and remote ip address and ports if they may be of some use. |
90 | | - In exceptional scenarios, e.g. ReadTimeouts, the exception doesn't contain any response |
91 | | - information - it's an either/or - i.e. a response or an exception (see Endpoint._send_request()), |
92 | | - and there's nowhere to inject the socket information into the exception. |
93 | | - """ |
94 | | - super().__init__(sock, **kwargs) |
95 | | - |
96 | | - self.sock_name = sock.getsockname() |
97 | | - self.peer_name = sock.getpeername() |
98 | | - |
99 | | - |
100 | | -def convert_to_response_dict_with_socket_info(http_response, operation_model): |
101 | | - # Call the real implementation... |
102 | | - response_dict = original_convert_to_response_dict(http_response, operation_model) |
103 | | - |
104 | | - # Then add the custom socket info. Put in a try catch in case there's an edge case |
105 | | - # where there's no _original_response |
106 | | - with contextlib.suppress(Exception): |
107 | | - response_dict["socket"] = { |
108 | | - "sock_name": cast(CustomHTTPResponse, http_response.raw._original_response).sock_name, |
109 | | - "peer_name": cast(CustomHTTPResponse, http_response.raw._original_response).peer_name, |
110 | | - } |
111 | | - |
112 | | - return response_dict |
113 | | - |
114 | | - |
115 | | -HTTPConnection.response_class = CustomHTTPResponse # type: ignore |
116 | | - |
117 | | -# Proxy the botocore function to allow us to add socket information to the response_dict |
118 | | -original_convert_to_response_dict = endpoint.convert_to_response_dict |
119 | | -endpoint.convert_to_response_dict = convert_to_response_dict_with_socket_info |
120 | | - |
121 | | - |
122 | 84 | def s3_build_uri(bucket: str, key: str) -> str: |
123 | 85 | return f"s3://{bucket}/{key}" |
124 | 86 |
|
|
0 commit comments