Coverage for netbox_data_import/inference_transport.py: 95.26%
211 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 09:47 +0000
1# SPDX-License-Identifier: Apache-2.0
2# SPDX-FileCopyrightText: 2026 Marcin Zieba <marcinpsk@gmail.com>
3"""Address-pinned HTTP transport for inference and credential requests."""
5import ipaddress
6import socket
7import time
9from collections import OrderedDict
10from contextlib import suppress
11from dataclasses import dataclass
12from threading import Event, Lock, RLock, Timer
13from typing import Any, Callable, TypeVar, cast
14from urllib.parse import urlsplit, urlunsplit
15from weakref import WeakKeyDictionary
17import requests
18from urllib3.connection import HTTPConnection, HTTPSConnection
19from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool
20from urllib3.exceptions import MaxRetryError, NewConnectionError
21from urllib3.response import HTTPResponse
22from urllib3.util import Timeout as Urllib3Timeout
25_SESSION_LOCKS: WeakKeyDictionary[requests.Session, RLock] = WeakKeyDictionary()
26_SESSION_LOCKS_GUARD = Lock()
27_Result = TypeVar("_Result")
30class WallClockDeadlineExceeded(requests.Timeout):
31 """One foreground operation exhausted its shared wall-clock budget."""
34@dataclass(frozen=True)
35class WallClockDeadline:
36 """Share one time budget across DNS and every HTTP request in an operation."""
38 expires_at: float
40 @classmethod
41 def after(cls, seconds: float) -> "WallClockDeadline":
42 """Return a deadline that expires after the supplied positive interval."""
43 if seconds <= 0:
44 raise ValueError("A wall-clock deadline must be positive.")
45 return cls(time.monotonic() + seconds)
47 def remaining(self) -> float:
48 """Return the remaining seconds, or fail when the shared budget is exhausted."""
49 remaining = self.expires_at - time.monotonic()
50 if remaining <= 0:
51 raise WallClockDeadlineExceeded("The operation exceeded its overall time limit.")
52 return remaining
54 def request_timeout(self, timeout: float | tuple[float | None, float | None] | None) -> Urllib3Timeout:
55 """Cap one Requests connect and initial read to the remaining shared budget."""
56 remaining = self.remaining()
57 if isinstance(timeout, tuple):
58 try:
59 connect, read = timeout
60 except ValueError as exc:
61 raise ValueError("A request timeout tuple must contain connect and read values.") from exc
62 else:
63 connect = read = timeout
64 try:
65 return Urllib3Timeout(
66 total=remaining,
67 connect=remaining if connect is None else min(connect, remaining),
68 read=remaining if read is None else min(read, remaining),
69 )
70 except TypeError as exc:
71 raise ValueError("A request timeout must be a number, None, or a two-value tuple.") from exc
74def _run_connection_until_deadline(
75 connection: HTTPConnection,
76 deadline: WallClockDeadline,
77 operation: Callable[[], _Result],
78) -> _Result:
79 """Interrupt connection or response-header I/O when its wall-clock budget expires."""
80 expired = Event()
82 def abort_socket() -> None:
83 expired.set()
84 connected_socket = connection.sock or getattr(connection, "_deadline_response_socket", None)
85 if connected_socket is None:
86 return
87 with suppress(OSError):
88 connected_socket.shutdown(socket.SHUT_RDWR)
89 with suppress(OSError):
90 connected_socket.close()
92 timer = Timer(deadline.remaining(), abort_socket)
93 timer.daemon = True
94 timer.start()
95 try:
96 result = operation()
97 except Exception:
98 if expired.is_set():
99 raise WallClockDeadlineExceeded("The operation exceeded its overall time limit.") from None
100 raise
101 else:
102 deadline.remaining()
103 return result
104 finally:
105 timer.cancel()
108def _deadline_pool_classes(deadline: WallClockDeadline):
109 """Return urllib3 pools whose connections enforce one operation deadline."""
111 class DeadlineHTTPConnection(HTTPConnection):
112 _deadline_response_socket: socket.socket | None = None
114 def connect(self) -> None:
115 _run_connection_until_deadline(self, deadline, super().connect)
117 def getresponse(self) -> HTTPResponse: # type: ignore[override]
118 self._deadline_response_socket = self.sock
119 return _run_connection_until_deadline(self, deadline, super().getresponse)
121 class DeadlineHTTPSConnection(HTTPSConnection):
122 _deadline_response_socket: socket.socket | None = None
124 def connect(self) -> None:
125 _run_connection_until_deadline(self, deadline, super().connect)
127 def getresponse(self) -> HTTPResponse: # type: ignore[override]
128 self._deadline_response_socket = self.sock
129 return _run_connection_until_deadline(self, deadline, super().getresponse)
131 class DeadlineHTTPConnectionPool(HTTPConnectionPool):
132 ConnectionCls = DeadlineHTTPConnection
134 class DeadlineHTTPSConnectionPool(HTTPSConnectionPool):
135 ConnectionCls = DeadlineHTTPSConnection
137 return {"http": DeadlineHTTPConnectionPool, "https": DeadlineHTTPSConnectionPool}
140class ResponseProcessingFailure(requests.RequestException):
141 """A response arrived, but Requests failed while processing it."""
143 response: requests.Response
145 def __init__(self, cause: Exception, response: requests.Response):
146 super().__init__(str(cause), response=response)
147 self.cause = cause
150class ResponseBodyTooLarge(requests.RequestException):
151 """The peer sent more response bytes than this request permits."""
154def is_preconnect_failure(exc: requests.RequestException) -> bool:
155 """Return whether another address can be tried without replaying a sent request."""
156 if isinstance(exc, requests.ConnectTimeout):
157 return True
158 reason = exc.args[0] if exc.args else None
159 return isinstance(reason, MaxRetryError) and isinstance(reason.reason, NewConnectionError)
162def _session_lock(session: requests.Session) -> RLock:
163 """Return the exclusive transport lock owned by one session."""
164 with _SESSION_LOCKS_GUARD:
165 lock = _SESSION_LOCKS.get(session)
166 if lock is None:
167 lock = RLock()
168 _SESSION_LOCKS[session] = lock
169 return lock
172def _consume_response(response, response_body_limit, deadline: WallClockDeadline | None) -> None:
173 """Read one response under its size and wall-clock limits."""
174 chunks = bytearray()
175 chunk_size = min(response_body_limit + 1, 65_536) if response_body_limit is not None else 65_536
176 iterator = response.iter_content(chunk_size=chunk_size)
177 connection = cast(HTTPConnection | None, getattr(response.raw, "connection", None))
179 def next_chunk():
180 if deadline is None:
181 return next(iterator)
182 if connection is None:
183 raise requests.ConnectionError("The streamed response has no active connection.")
184 return _run_connection_until_deadline(connection, deadline, lambda: next(iterator))
186 while True:
187 try:
188 chunk = next_chunk()
189 except StopIteration:
190 break
191 if not chunk:
192 continue
193 chunks.extend(chunk)
194 if response_body_limit is not None and len(chunks) > response_body_limit:
195 response.close()
196 raise ResponseBodyTooLarge(f"The response body exceeded {response_body_limit} bytes.", response=response)
197 if deadline is not None:
198 deadline.remaining()
199 response._content = bytes(chunks)
200 response._content_consumed = True
203class _AddressPinnedAdapter(requests.adapters.HTTPAdapter):
204 """Connect to one numeric address while preserving the origin hostname."""
206 def __init__(
207 self,
208 origin_url: str,
209 resolved_address: str,
210 deadline: WallClockDeadline | None = None,
211 response_body_limit: int | None = None,
212 ):
213 super().__init__()
214 parts = urlsplit(origin_url)
215 self._origin = (parts.scheme.lower(), parts.hostname, parts.port)
216 address = ipaddress.ip_address(resolved_address)
217 self._address = f"[{address}]" if address.version == 6 else str(address)
218 self._deadline = deadline
219 self._response_body_limit = response_body_limit
220 if deadline is not None:
221 self.poolmanager.pool_classes_by_scheme = _deadline_pool_classes(deadline)
223 def build_connection_pool_key_attributes(self, request, verify, cert=None):
224 """Keep TLS SNI and certificate verification bound to the origin hostname."""
225 host_params, pool_kwargs = super().build_connection_pool_key_attributes(request, verify, cert)
226 if self._origin[0] == "https":
227 pool_kwargs["assert_hostname"] = self._origin[1]
228 pool_kwargs["server_hostname"] = self._origin[1]
229 return host_params, pool_kwargs
231 def send(self, request, *args, **kwargs):
232 """Rewrite only the connection URL, then restore the prepared request."""
233 parts = urlsplit(request.url)
234 request_origin = (parts.scheme.lower(), parts.hostname, parts.port)
235 if request_origin != self._origin:
236 raise requests.exceptions.InvalidURL("The pinned transport received a request for another origin.")
237 original_url = request.url
238 original_host = request.headers.get("Host")
239 port = f":{parts.port}" if parts.port is not None else ""
240 request.url = urlunsplit((parts.scheme, f"{self._address}{port}", parts.path, parts.query, parts.fragment))
241 request.headers["Host"] = parts.netloc
242 try:
243 try:
244 # A proxy would resolve the hostname again and bypass the pinned address.
245 response = super().send(request, *args, **{**kwargs, "proxies": {}})
246 except requests.RequestException:
247 if self._deadline is not None:
248 self._deadline.remaining()
249 raise
250 response.url = original_url
251 if self._response_body_limit is not None or self._deadline is not None:
252 try:
253 _consume_response(response, self._response_body_limit, self._deadline)
254 except (ValueError, requests.RequestException) as exc:
255 response.close()
256 raise ResponseProcessingFailure(exc, response) from exc
257 return response
258 finally:
259 request.url = original_url
260 if original_host is None:
261 request.headers.pop("Host", None)
262 else:
263 request.headers["Host"] = original_host
266def request_to_resolved_address(
267 session: requests.Session,
268 method: str,
269 url: str,
270 resolved_address: str,
271 response_body_limit: int | None = None,
272 deadline: WallClockDeadline | None = None,
273 **kwargs: Any,
274) -> requests.Response:
275 """Send one request to a resolved address without changing its HTTP or TLS hostname."""
276 if response_body_limit is not None and response_body_limit < 1:
277 raise ValueError("A response body limit must be positive.")
278 if response_body_limit is not None or deadline is not None:
279 kwargs["stream"] = True
280 with _session_lock(session):
281 previous_adapters = OrderedDict(session.adapters)
282 adapter = _AddressPinnedAdapter(url, resolved_address, deadline, response_body_limit)
283 captured_response = None
285 def capture_response(response, *_args, **_kwargs):
286 nonlocal captured_response
287 captured_response = response
288 return response
290 hooks = dict(kwargs.pop("hooks", {}) or {})
291 response_hooks = hooks.get("response", ())
292 if response_hooks is None:
293 response_hooks = ()
294 elif callable(response_hooks):
295 response_hooks = (response_hooks,)
296 hooks["response"] = (capture_response, *response_hooks)
297 session.mount(url, adapter)
298 try:
299 try:
300 if deadline is not None:
301 kwargs["timeout"] = deadline.request_timeout(kwargs.get("timeout"))
302 response = session.request(method, url, hooks=hooks, **kwargs)
303 except (ValueError, requests.RequestException) as exc:
304 if captured_response is not None:
305 captured_response.close()
306 raise ResponseProcessingFailure(exc, captured_response) from exc
307 raise
308 else:
309 return response
310 finally:
311 adapter.close()
312 session.adapters.clear()
313 session.adapters.update(previous_adapters)
316__all__ = (
317 "ResponseBodyTooLarge",
318 "ResponseProcessingFailure",
319 "WallClockDeadline",
320 "WallClockDeadlineExceeded",
321 "is_preconnect_failure",
322 "request_to_resolved_address",
323)