poolmanager.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470
  1. from __future__ import absolute_import
  2. import collections
  3. import functools
  4. import logging
  5. from ._collections import RecentlyUsedContainer
  6. from .connectionpool import HTTPConnectionPool, HTTPSConnectionPool
  7. from .connectionpool import port_by_scheme
  8. from .exceptions import LocationValueError, MaxRetryError, ProxySchemeUnknown
  9. from .packages import six
  10. from .packages.six.moves.urllib.parse import urljoin
  11. from .request import RequestMethods
  12. from .util.url import parse_url
  13. from .util.retry import Retry
  14. __all__ = ["PoolManager", "ProxyManager", "proxy_from_url"]
  15. log = logging.getLogger(__name__)
  16. SSL_KEYWORDS = (
  17. "key_file",
  18. "cert_file",
  19. "cert_reqs",
  20. "ca_certs",
  21. "ssl_version",
  22. "ca_cert_dir",
  23. "ssl_context",
  24. "key_password",
  25. )
  26. # All known keyword arguments that could be provided to the pool manager, its
  27. # pools, or the underlying connections. This is used to construct a pool key.
  28. _key_fields = (
  29. "key_scheme", # str
  30. "key_host", # str
  31. "key_port", # int
  32. "key_timeout", # int or float or Timeout
  33. "key_retries", # int or Retry
  34. "key_strict", # bool
  35. "key_block", # bool
  36. "key_source_address", # str
  37. "key_key_file", # str
  38. "key_key_password", # str
  39. "key_cert_file", # str
  40. "key_cert_reqs", # str
  41. "key_ca_certs", # str
  42. "key_ssl_version", # str
  43. "key_ca_cert_dir", # str
  44. "key_ssl_context", # instance of ssl.SSLContext or urllib3.util.ssl_.SSLContext
  45. "key_maxsize", # int
  46. "key_headers", # dict
  47. "key__proxy", # parsed proxy url
  48. "key__proxy_headers", # dict
  49. "key_socket_options", # list of (level (int), optname (int), value (int or str)) tuples
  50. "key__socks_options", # dict
  51. "key_assert_hostname", # bool or string
  52. "key_assert_fingerprint", # str
  53. "key_server_hostname", # str
  54. )
  55. #: The namedtuple class used to construct keys for the connection pool.
  56. #: All custom key schemes should include the fields in this key at a minimum.
  57. PoolKey = collections.namedtuple("PoolKey", _key_fields)
  58. def _default_key_normalizer(key_class, request_context):
  59. """
  60. Create a pool key out of a request context dictionary.
  61. According to RFC 3986, both the scheme and host are case-insensitive.
  62. Therefore, this function normalizes both before constructing the pool
  63. key for an HTTPS request. If you wish to change this behaviour, provide
  64. alternate callables to ``key_fn_by_scheme``.
  65. :param key_class:
  66. The class to use when constructing the key. This should be a namedtuple
  67. with the ``scheme`` and ``host`` keys at a minimum.
  68. :type key_class: namedtuple
  69. :param request_context:
  70. A dictionary-like object that contain the context for a request.
  71. :type request_context: dict
  72. :return: A namedtuple that can be used as a connection pool key.
  73. :rtype: PoolKey
  74. """
  75. # Since we mutate the dictionary, make a copy first
  76. context = request_context.copy()
  77. context["scheme"] = context["scheme"].lower()
  78. context["host"] = context["host"].lower()
  79. # These are both dictionaries and need to be transformed into frozensets
  80. for key in ("headers", "_proxy_headers", "_socks_options"):
  81. if key in context and context[key] is not None:
  82. context[key] = frozenset(context[key].items())
  83. # The socket_options key may be a list and needs to be transformed into a
  84. # tuple.
  85. socket_opts = context.get("socket_options")
  86. if socket_opts is not None:
  87. context["socket_options"] = tuple(socket_opts)
  88. # Map the kwargs to the names in the namedtuple - this is necessary since
  89. # namedtuples can't have fields starting with '_'.
  90. for key in list(context.keys()):
  91. context["key_" + key] = context.pop(key)
  92. # Default to ``None`` for keys missing from the context
  93. for field in key_class._fields:
  94. if field not in context:
  95. context[field] = None
  96. return key_class(**context)
  97. #: A dictionary that maps a scheme to a callable that creates a pool key.
  98. #: This can be used to alter the way pool keys are constructed, if desired.
  99. #: Each PoolManager makes a copy of this dictionary so they can be configured
  100. #: globally here, or individually on the instance.
  101. key_fn_by_scheme = {
  102. "http": functools.partial(_default_key_normalizer, PoolKey),
  103. "https": functools.partial(_default_key_normalizer, PoolKey),
  104. }
  105. pool_classes_by_scheme = {"http": HTTPConnectionPool, "https": HTTPSConnectionPool}
  106. class PoolManager(RequestMethods):
  107. """
  108. Allows for arbitrary requests while transparently keeping track of
  109. necessary connection pools for you.
  110. :param num_pools:
  111. Number of connection pools to cache before discarding the least
  112. recently used pool.
  113. :param headers:
  114. Headers to include with all requests, unless other headers are given
  115. explicitly.
  116. :param \\**connection_pool_kw:
  117. Additional parameters are used to create fresh
  118. :class:`urllib3.connectionpool.ConnectionPool` instances.
  119. Example::
  120. >>> manager = PoolManager(num_pools=2)
  121. >>> r = manager.request('GET', 'http://google.com/')
  122. >>> r = manager.request('GET', 'http://google.com/mail')
  123. >>> r = manager.request('GET', 'http://yahoo.com/')
  124. >>> len(manager.pools)
  125. 2
  126. """
  127. proxy = None
  128. def __init__(self, num_pools=10, headers=None, **connection_pool_kw):
  129. RequestMethods.__init__(self, headers)
  130. self.connection_pool_kw = connection_pool_kw
  131. self.pools = RecentlyUsedContainer(num_pools, dispose_func=lambda p: p.close())
  132. # Locally set the pool classes and keys so other PoolManagers can
  133. # override them.
  134. self.pool_classes_by_scheme = pool_classes_by_scheme
  135. self.key_fn_by_scheme = key_fn_by_scheme.copy()
  136. def __enter__(self):
  137. return self
  138. def __exit__(self, exc_type, exc_val, exc_tb):
  139. self.clear()
  140. # Return False to re-raise any potential exceptions
  141. return False
  142. def _new_pool(self, scheme, host, port, request_context=None):
  143. """
  144. Create a new :class:`ConnectionPool` based on host, port, scheme, and
  145. any additional pool keyword arguments.
  146. If ``request_context`` is provided, it is provided as keyword arguments
  147. to the pool class used. This method is used to actually create the
  148. connection pools handed out by :meth:`connection_from_url` and
  149. companion methods. It is intended to be overridden for customization.
  150. """
  151. pool_cls = self.pool_classes_by_scheme[scheme]
  152. if request_context is None:
  153. request_context = self.connection_pool_kw.copy()
  154. # Although the context has everything necessary to create the pool,
  155. # this function has historically only used the scheme, host, and port
  156. # in the positional args. When an API change is acceptable these can
  157. # be removed.
  158. for key in ("scheme", "host", "port"):
  159. request_context.pop(key, None)
  160. if scheme == "http":
  161. for kw in SSL_KEYWORDS:
  162. request_context.pop(kw, None)
  163. return pool_cls(host, port, **request_context)
  164. def clear(self):
  165. """
  166. Empty our store of pools and direct them all to close.
  167. This will not affect in-flight connections, but they will not be
  168. re-used after completion.
  169. """
  170. self.pools.clear()
  171. def connection_from_host(self, host, port=None, scheme="http", pool_kwargs=None):
  172. """
  173. Get a :class:`ConnectionPool` based on the host, port, and scheme.
  174. If ``port`` isn't given, it will be derived from the ``scheme`` using
  175. ``urllib3.connectionpool.port_by_scheme``. If ``pool_kwargs`` is
  176. provided, it is merged with the instance's ``connection_pool_kw``
  177. variable and used to create the new connection pool, if one is
  178. needed.
  179. """
  180. if not host:
  181. raise LocationValueError("No host specified.")
  182. request_context = self._merge_pool_kwargs(pool_kwargs)
  183. request_context["scheme"] = scheme or "http"
  184. if not port:
  185. port = port_by_scheme.get(request_context["scheme"].lower(), 80)
  186. request_context["port"] = port
  187. request_context["host"] = host
  188. return self.connection_from_context(request_context)
  189. def connection_from_context(self, request_context):
  190. """
  191. Get a :class:`ConnectionPool` based on the request context.
  192. ``request_context`` must at least contain the ``scheme`` key and its
  193. value must be a key in ``key_fn_by_scheme`` instance variable.
  194. """
  195. scheme = request_context["scheme"].lower()
  196. pool_key_constructor = self.key_fn_by_scheme[scheme]
  197. pool_key = pool_key_constructor(request_context)
  198. return self.connection_from_pool_key(pool_key, request_context=request_context)
  199. def connection_from_pool_key(self, pool_key, request_context=None):
  200. """
  201. Get a :class:`ConnectionPool` based on the provided pool key.
  202. ``pool_key`` should be a namedtuple that only contains immutable
  203. objects. At a minimum it must have the ``scheme``, ``host``, and
  204. ``port`` fields.
  205. """
  206. with self.pools.lock:
  207. # If the scheme, host, or port doesn't match existing open
  208. # connections, open a new ConnectionPool.
  209. pool = self.pools.get(pool_key)
  210. if pool:
  211. return pool
  212. # Make a fresh ConnectionPool of the desired type
  213. scheme = request_context["scheme"]
  214. host = request_context["host"]
  215. port = request_context["port"]
  216. pool = self._new_pool(scheme, host, port, request_context=request_context)
  217. self.pools[pool_key] = pool
  218. return pool
  219. def connection_from_url(self, url, pool_kwargs=None):
  220. """
  221. Similar to :func:`urllib3.connectionpool.connection_from_url`.
  222. If ``pool_kwargs`` is not provided and a new pool needs to be
  223. constructed, ``self.connection_pool_kw`` is used to initialize
  224. the :class:`urllib3.connectionpool.ConnectionPool`. If ``pool_kwargs``
  225. is provided, it is used instead. Note that if a new pool does not
  226. need to be created for the request, the provided ``pool_kwargs`` are
  227. not used.
  228. """
  229. u = parse_url(url)
  230. return self.connection_from_host(
  231. u.host, port=u.port, scheme=u.scheme, pool_kwargs=pool_kwargs
  232. )
  233. def _merge_pool_kwargs(self, override):
  234. """
  235. Merge a dictionary of override values for self.connection_pool_kw.
  236. This does not modify self.connection_pool_kw and returns a new dict.
  237. Any keys in the override dictionary with a value of ``None`` are
  238. removed from the merged dictionary.
  239. """
  240. base_pool_kwargs = self.connection_pool_kw.copy()
  241. if override:
  242. for key, value in override.items():
  243. if value is None:
  244. try:
  245. del base_pool_kwargs[key]
  246. except KeyError:
  247. pass
  248. else:
  249. base_pool_kwargs[key] = value
  250. return base_pool_kwargs
  251. def urlopen(self, method, url, redirect=True, **kw):
  252. """
  253. Same as :meth:`urllib3.connectionpool.HTTPConnectionPool.urlopen`
  254. with custom cross-host redirect logic and only sends the request-uri
  255. portion of the ``url``.
  256. The given ``url`` parameter must be absolute, such that an appropriate
  257. :class:`urllib3.connectionpool.ConnectionPool` can be chosen for it.
  258. """
  259. u = parse_url(url)
  260. conn = self.connection_from_host(u.host, port=u.port, scheme=u.scheme)
  261. kw["assert_same_host"] = False
  262. kw["redirect"] = False
  263. if "headers" not in kw:
  264. kw["headers"] = self.headers.copy()
  265. if self.proxy is not None and u.scheme == "http":
  266. response = conn.urlopen(method, url, **kw)
  267. else:
  268. response = conn.urlopen(method, u.request_uri, **kw)
  269. redirect_location = redirect and response.get_redirect_location()
  270. if not redirect_location:
  271. return response
  272. # Support relative URLs for redirecting.
  273. redirect_location = urljoin(url, redirect_location)
  274. # RFC 7231, Section 6.4.4
  275. if response.status == 303:
  276. method = "GET"
  277. retries = kw.get("retries")
  278. if not isinstance(retries, Retry):
  279. retries = Retry.from_int(retries, redirect=redirect)
  280. # Strip headers marked as unsafe to forward to the redirected location.
  281. # Check remove_headers_on_redirect to avoid a potential network call within
  282. # conn.is_same_host() which may use socket.gethostbyname() in the future.
  283. if retries.remove_headers_on_redirect and not conn.is_same_host(
  284. redirect_location
  285. ):
  286. headers = list(six.iterkeys(kw["headers"]))
  287. for header in headers:
  288. if header.lower() in retries.remove_headers_on_redirect:
  289. kw["headers"].pop(header, None)
  290. try:
  291. retries = retries.increment(method, url, response=response, _pool=conn)
  292. except MaxRetryError:
  293. if retries.raise_on_redirect:
  294. raise
  295. return response
  296. kw["retries"] = retries
  297. kw["redirect"] = redirect
  298. log.info("Redirecting %s -> %s", url, redirect_location)
  299. return self.urlopen(method, redirect_location, **kw)
  300. class ProxyManager(PoolManager):
  301. """
  302. Behaves just like :class:`PoolManager`, but sends all requests through
  303. the defined proxy, using the CONNECT method for HTTPS URLs.
  304. :param proxy_url:
  305. The URL of the proxy to be used.
  306. :param proxy_headers:
  307. A dictionary containing headers that will be sent to the proxy. In case
  308. of HTTP they are being sent with each request, while in the
  309. HTTPS/CONNECT case they are sent only once. Could be used for proxy
  310. authentication.
  311. Example:
  312. >>> proxy = urllib3.ProxyManager('http://localhost:3128/')
  313. >>> r1 = proxy.request('GET', 'http://google.com/')
  314. >>> r2 = proxy.request('GET', 'http://httpbin.org/')
  315. >>> len(proxy.pools)
  316. 1
  317. >>> r3 = proxy.request('GET', 'https://httpbin.org/')
  318. >>> r4 = proxy.request('GET', 'https://twitter.com/')
  319. >>> len(proxy.pools)
  320. 3
  321. """
  322. def __init__(
  323. self,
  324. proxy_url,
  325. num_pools=10,
  326. headers=None,
  327. proxy_headers=None,
  328. **connection_pool_kw
  329. ):
  330. if isinstance(proxy_url, HTTPConnectionPool):
  331. proxy_url = "%s://%s:%i" % (
  332. proxy_url.scheme,
  333. proxy_url.host,
  334. proxy_url.port,
  335. )
  336. proxy = parse_url(proxy_url)
  337. if not proxy.port:
  338. port = port_by_scheme.get(proxy.scheme, 80)
  339. proxy = proxy._replace(port=port)
  340. if proxy.scheme not in ("http", "https"):
  341. raise ProxySchemeUnknown(proxy.scheme)
  342. self.proxy = proxy
  343. self.proxy_headers = proxy_headers or {}
  344. connection_pool_kw["_proxy"] = self.proxy
  345. connection_pool_kw["_proxy_headers"] = self.proxy_headers
  346. super(ProxyManager, self).__init__(num_pools, headers, **connection_pool_kw)
  347. def connection_from_host(self, host, port=None, scheme="http", pool_kwargs=None):
  348. if scheme == "https":
  349. return super(ProxyManager, self).connection_from_host(
  350. host, port, scheme, pool_kwargs=pool_kwargs
  351. )
  352. return super(ProxyManager, self).connection_from_host(
  353. self.proxy.host, self.proxy.port, self.proxy.scheme, pool_kwargs=pool_kwargs
  354. )
  355. def _set_proxy_headers(self, url, headers=None):
  356. """
  357. Sets headers needed by proxies: specifically, the Accept and Host
  358. headers. Only sets headers not provided by the user.
  359. """
  360. headers_ = {"Accept": "*/*"}
  361. netloc = parse_url(url).netloc
  362. if netloc:
  363. headers_["Host"] = netloc
  364. if headers:
  365. headers_.update(headers)
  366. return headers_
  367. def urlopen(self, method, url, redirect=True, **kw):
  368. "Same as HTTP(S)ConnectionPool.urlopen, ``url`` must be absolute."
  369. u = parse_url(url)
  370. if u.scheme == "http":
  371. # For proxied HTTPS requests, httplib sets the necessary headers
  372. # on the CONNECT to the proxy. For HTTP, we'll definitely
  373. # need to set 'Host' at the very least.
  374. headers = kw.get("headers", self.headers)
  375. kw["headers"] = self._set_proxy_headers(url, headers)
  376. return super(ProxyManager, self).urlopen(method, url, redirect=redirect, **kw)
  377. def proxy_from_url(url, **kw):
  378. return ProxyManager(proxy_url=url, **kw)