|
18 | 18 | import logging |
19 | 19 | import socket |
20 | 20 | import time |
| 21 | +from functools import cached_property |
21 | 22 | from types import TracebackType |
22 | 23 | from typing import ( |
23 | 24 | TYPE_CHECKING, |
@@ -142,40 +143,47 @@ class _HiveClient: |
142 | 143 | """Helper class to nicely open and close the transport.""" |
143 | 144 |
|
144 | 145 | _transport: TTransport |
145 | | - _client: Client |
146 | 146 | _ugi: Optional[List[str]] |
147 | 147 |
|
148 | 148 | def __init__(self, uri: str, ugi: Optional[str] = None, kerberos_auth: Optional[bool] = HIVE_KERBEROS_AUTH_DEFAULT): |
149 | 149 | self._uri = uri |
150 | 150 | self._kerberos_auth = kerberos_auth |
151 | 151 | self._ugi = ugi.split(":") if ugi else None |
| 152 | + self._transport = self._init_thrift_transport() |
152 | 153 |
|
153 | | - self._init_thrift_client() |
154 | | - |
155 | | - def _init_thrift_client(self) -> None: |
| 154 | + def _init_thrift_transport(self) -> TTransport: |
156 | 155 | url_parts = urlparse(self._uri) |
157 | | - |
158 | 156 | socket = TSocket.TSocket(url_parts.hostname, url_parts.port) |
159 | | - |
160 | 157 | if not self._kerberos_auth: |
161 | | - self._transport = TTransport.TBufferedTransport(socket) |
| 158 | + return TTransport.TBufferedTransport(socket) |
162 | 159 | else: |
163 | | - self._transport = TTransport.TSaslClientTransport(socket, host=url_parts.hostname, service="hive") |
| 160 | + return TTransport.TSaslClientTransport(socket, host=url_parts.hostname, service="hive") |
164 | 161 |
|
| 162 | + @cached_property |
| 163 | + def _client(self) -> Client: |
165 | 164 | protocol = TBinaryProtocol.TBinaryProtocol(self._transport) |
166 | | - |
167 | | - self._client = Client(protocol) |
| 165 | + client = Client(protocol) |
| 166 | + if self._ugi: |
| 167 | + client.set_ugi(*self._ugi) |
| 168 | + return client |
168 | 169 |
|
169 | 170 | def __enter__(self) -> Client: |
170 | | - self._transport.open() |
171 | | - if self._ugi: |
172 | | - self._client.set_ugi(*self._ugi) |
| 171 | + """Make sure the transport is initialized and open.""" |
| 172 | + if not self._transport.isOpen(): |
| 173 | + try: |
| 174 | + self._transport.open() |
| 175 | + except TTransport.TTransportException: |
| 176 | + # reinitialize _transport |
| 177 | + self._transport = self._init_thrift_transport() |
| 178 | + self._transport.open() |
173 | 179 | return self._client |
174 | 180 |
|
175 | 181 | def __exit__( |
176 | 182 | self, exctype: Optional[Type[BaseException]], excinst: Optional[BaseException], exctb: Optional[TracebackType] |
177 | 183 | ) -> None: |
178 | | - self._transport.close() |
| 184 | + """Close transport if it was opened.""" |
| 185 | + if self._transport.isOpen(): |
| 186 | + self._transport.close() |
179 | 187 |
|
180 | 188 |
|
181 | 189 | def _construct_hive_storage_descriptor( |
|
0 commit comments