(iceberg-python) 02/11: Fix thrift client connection for Kerberos Hive Client (#1747)

[fo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

fokko pushed a commit to branch pyiceberg-0.9.x
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git

commit b989197f94093fcbd0bbe8d58928dda9d1888b27
Author: Kevin Liu <ke...@users.noreply.github.com>
AuthorDate: Wed Apr 16 14:37:23 2025 -0400

    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.
---
 pyiceberg/catalog/hive.py | 36 ++++++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/pyiceberg/catalog/hive.py b/pyiceberg/catalog/hive.py
index 7189b12c..cad89e10 100644
--- a/pyiceberg/catalog/hive.py
+++ b/pyiceberg/catalog/hive.py
@@ -18,6 +18,7 @@ import getpass
 import logging
 import socket
 import time
+from functools import cached_property
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
@@ -142,40 +143,47 @@ class _HiveClient:
     """Helper class to nicely open and close the transport."""
 
     _transport: TTransport
-    _client: Client
     _ugi: Optional[List[str]]
 
     def __init__(self, uri: str, ugi: Optional[str] = None, kerberos_auth: Optional[bool] = HIVE_KERBEROS_AUTH_DEFAULT):
         self._uri = uri
         self._kerberos_auth = kerberos_auth
         self._ugi = ugi.split(":") if ugi else None
+        self._transport = self._init_thrift_transport()
 
-        self._init_thrift_client()
-
-    def _init_thrift_client(self) -> None:
+    def _init_thrift_transport(self) -> TTransport:
         url_parts = urlparse(self._uri)
-
         socket = TSocket.TSocket(url_parts.hostname, url_parts.port)
-
         if not self._kerberos_auth:
-            self._transport = TTransport.TBufferedTransport(socket)
+            return TTransport.TBufferedTransport(socket)
         else:
-            self._transport = TTransport.TSaslClientTransport(socket, host=url_parts.hostname, service="hive")
+            return TTransport.TSaslClientTransport(socket, host=url_parts.hostname, service="hive")
 
+    @cached_property
+    def _client(self) -> Client:
         protocol = TBinaryProtocol.TBinaryProtocol(self._transport)
-
-        self._client = Client(protocol)
+        client = Client(protocol)
+        if self._ugi:
+            client.set_ugi(*self._ugi)
+        return client
 
     def __enter__(self) -> Client:
-        self._transport.open()
-        if self._ugi:
-            self._client.set_ugi(*self._ugi)
+        """Make sure the transport is initialized and open."""
+        if not self._transport.isOpen():
+            try:
+                self._transport.open()
+            except TTransport.TTransportException:
+                # reinitialize _transport
+                self._transport = self._init_thrift_transport()
+                self._transport.open()
         return self._client
 
     def __exit__(
         self, exctype: Optional[Type[BaseException]], excinst: Optional[BaseException], exctb: Optional[TracebackType]
     ) -> None:
-        self._transport.close()
+        """Close transport if it was opened."""
+        if self._transport.isOpen():
+            self._transport.close()
 
 
 def _construct_hive_storage_descriptor(