1+ import base64
12import logging
2- from typing import Dict
3-
3+ import urllib . parse
4+ from typing import Dict , Union
45
6+ import six
57import thrift
68
7- import urllib .parse ,six ,base64
8-
99logger = logging .getLogger (__name__ )
1010
11+ import ssl
12+ import warnings
13+ from http .client import HTTPResponse
14+ from io import BytesIO
15+
16+ from urllib3 import HTTPConnectionPool ,HTTPSConnectionPool ,ProxyManager
17+
1118
1219class THttpClient (thrift .transport .THttpClient .THttpClient ):
1320def __init__ (
@@ -20,22 +27,152 @@ def __init__(
2027cert_file = None ,
2128key_file = None ,
2229ssl_context = None ,
30+ max_connections :int = 1 ,
2331 ):
24- super ().__init__ (
25- uri_or_host ,port ,path ,cafile ,cert_file ,key_file ,ssl_context
26- )
32+ if port is not None :
33+ warnings .warn (
34+ "Please use the THttpClient('http{s}://host:port/path') constructor" ,
35+ DeprecationWarning ,
36+ stacklevel = 2 ,
37+ )
38+ self .host = uri_or_host
39+ self .port = port
40+ assert path
41+ self .path = path
42+ self .scheme = "http"
43+ else :
44+ parsed = urllib .parse .urlsplit (uri_or_host )
45+ self .scheme = parsed .scheme
46+ assert self .scheme in ("http" ,"https" )
47+ if self .scheme == "https" :
48+ self .certfile = cert_file
49+ self .keyfile = key_file
50+ self .context = (
51+ ssl .create_default_context (cafile = cafile )
52+ if (cafile and not ssl_context )
53+ else ssl_context
54+ )
55+ self .port = parsed .port
56+ self .host = parsed .hostname
57+ self .path = parsed .path
58+ if parsed .query :
59+ self .path += "?%s" % parsed .query
60+ try :
61+ proxy = urllib .request .getproxies ()[self .scheme ]
62+ except KeyError :
63+ proxy = None
64+ else :
65+ if urllib .request .proxy_bypass (self .host ):
66+ proxy = None
67+ if proxy :
68+ parsed = urllib .parse .urlparse (proxy )
69+
70+ # realhost and realport are the host and port of the actual request
71+ self .realhost = self .host
72+ self .realport = self .port
73+
74+ # this is passed to ProxyManager
75+ self .proxy_uri :str = proxy
76+ self .host = parsed .hostname
77+ self .port = parsed .port
78+ self .proxy_auth = self .basic_proxy_auth_header (parsed )
79+ else :
80+ self .realhost = self .realport = self .proxy_auth = None
81+
82+ self .max_connections = max_connections
83+
84+ self .__wbuf = BytesIO ()
85+ self .__resp :Union [None ,HTTPResponse ]= None
86+ self .__timeout = None
87+ self .__custom_headers = None
88+
2789self .__auth_provider = auth_provider
2890
2991def setCustomHeaders (self ,headers :Dict [str ,str ]):
3092self ._headers = headers
3193super ().setCustomHeaders (headers )
3294
95+ def open (self ):
96+
97+ # self.__pool replaces the self.__http used by the original THttpClient
98+ if self .scheme == "http" :
99+ pool_class = HTTPConnectionPool
100+ elif self .scheme == "https" :
101+ pool_class = HTTPSConnectionPool
102+
103+ _pool_kwargs = {"maxsize" :self .max_connections }
104+
105+ if self .using_proxy ():
106+ proxy_manager = ProxyManager (
107+ self .proxy_uri ,
108+ num_pools = 1 ,
109+ headers = {"Proxy-Authorization" :self .proxy_auth },
110+ )
111+ self .__pool = proxy_manager .connection_from_host (
112+ self .host ,self .port ,pool_kwargs = _pool_kwargs
113+ )
114+ else :
115+ self .__pool = pool_class (self .host ,self .port ,** _pool_kwargs )
116+
117+ def close (self ):
118+ self .__resp and self .__resp .release_conn ()
119+ self .__resp = None
120+
121+ def read (self ,sz ):
122+ return self .__resp .read (sz )
123+
124+ def isOpen (self ):
125+ return self .__resp is not None
126+
33127def flush (self ):
128+
129+ # Pull data out of buffer that will be sent in this request
130+ data = self .__wbuf .getvalue ()
131+ self .__wbuf = BytesIO ()
132+
133+ # Header handling
134+
34135headers = dict (self ._headers )
35136self .__auth_provider .add_headers (headers )
36137self ._headers = headers
37138self .setCustomHeaders (self ._headers )
38- super ().flush ()
139+
140+ # Note: we don't set User-Agent explicitly in this class because PySQL
141+ # should always provide one. Unlike the original THttpClient class, our version
142+ # doesn't define a default User-Agent and so should raise an exception if one
143+ # isn't provided.
144+ assert self .__custom_headers and "User-Agent" in self .__custom_headers
145+
146+ headers = {
147+ "Content-Type" :"application/x-thrift" ,
148+ "Content-Length" :str (len (data )),
149+ }
150+
151+ if self .using_proxy ()and self .scheme == "http" and self .proxy_auth is not None :
152+ headers ["Proxy-Authorization" :self .proxy_auth ]
153+
154+ if self .__custom_headers :
155+ custom_headers = {key :val for key ,val in self .__custom_headers .items ()}
156+ headers .update (** custom_headers )
157+
158+ # HTTP request
159+ self .__resp = self .__pool .request (
160+ "POST" ,
161+ url = self .path ,
162+ body = data ,
163+ headers = headers ,
164+ preload_content = False ,
165+ timeout = self .__timeout ,
166+ )
167+
168+ # Get reply to flush the request
169+ self .code = self .__resp .status
170+ self .message = self .__resp .reason
171+ self .headers = self .__resp .headers
172+
173+ # Saves the cookie sent by the server response
174+ if "Set-Cookie" in self .headers :
175+ self .setCustomHeaders (dict ("Cookie" ,self .headers ["Set-Cookie" ]))
39176
40177@staticmethod
41178def basic_proxy_auth_header (proxy ):
