Skip to content

Commit 881b2d5

Browse files
authored
Fix thrift client connection for Kerberos Hive Client (#1747)
Closes #1744 `TSaslClientTransport` cannot be reopen. This PR changes the behavior to recreate a `TSaslClientTransport` when its already closed. Note, `_HiveClient` should be used with context manager, but can be used without.
1 parent 5f10bbc commit 881b2d5

File tree

1 file changed

+22
-14
lines changed

1 file changed

+22
-14
lines changed

pyiceberg/catalog/hive.py

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import logging
1919
import socket
2020
import time
21+
from functools import cached_property
2122
from types import TracebackType
2223
from typing import (
2324
TYPE_CHECKING,
@@ -143,40 +144,47 @@ class _HiveClient:
143144
"""Helper class to nicely open and close the transport."""
144145

145146
_transport: TTransport
146-
_client: Client
147147
_ugi: Optional[List[str]]
148148

149149
def __init__(self, uri: str, ugi: Optional[str] = None, kerberos_auth: Optional[bool] = HIVE_KERBEROS_AUTH_DEFAULT):
150150
self._uri = uri
151151
self._kerberos_auth = kerberos_auth
152152
self._ugi = ugi.split(":") if ugi else None
153+
self._transport = self._init_thrift_transport()
153154

154-
self._init_thrift_client()
155-
156-
def _init_thrift_client(self) -> None:
155+
def _init_thrift_transport(self) -> TTransport:
157156
url_parts = urlparse(self._uri)
158-
159157
socket = TSocket.TSocket(url_parts.hostname, url_parts.port)
160-
161158
if not self._kerberos_auth:
162-
self._transport = TTransport.TBufferedTransport(socket)
159+
return TTransport.TBufferedTransport(socket)
163160
else:
164-
self._transport = TTransport.TSaslClientTransport(socket, host=url_parts.hostname, service="hive")
161+
return TTransport.TSaslClientTransport(socket, host=url_parts.hostname, service="hive")
165162

163+
@cached_property
164+
def _client(self) -> Client:
166165
protocol = TBinaryProtocol.TBinaryProtocol(self._transport)
167-
168-
self._client = Client(protocol)
166+
client = Client(protocol)
167+
if self._ugi:
168+
client.set_ugi(*self._ugi)
169+
return client
169170

170171
def __enter__(self) -> Client:
171-
self._transport.open()
172-
if self._ugi:
173-
self._client.set_ugi(*self._ugi)
172+
"""Make sure the transport is initialized and open."""
173+
if not self._transport.isOpen():
174+
try:
175+
self._transport.open()
176+
except TTransport.TTransportException:
177+
# reinitialize _transport
178+
self._transport = self._init_thrift_transport()
179+
self._transport.open()
174180
return self._client
175181

176182
def __exit__(
177183
self, exctype: Optional[Type[BaseException]], excinst: Optional[BaseException], exctb: Optional[TracebackType]
178184
) -> None:
179-
self._transport.close()
185+
"""Close transport if it was opened."""
186+
if self._transport.isOpen():
187+
self._transport.close()
180188

181189

182190
def _construct_hive_storage_descriptor(

0 commit comments

Comments
 (0)