|
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