Preserve _HiveClient.__init__ transport setup; reinit on reuse via flag

Author: barking-code

pyiceberg/catalog/hive.py

@@ -142,18 +142,11 @@
 
 
 class _HiveClient:
-    """Helper class to nicely open and close the transport.
+    """Helper class to nicely open and close the transport."""
 
-    ``TSaslClientTransport`` instances are single-use: ``close()`` disposes
-    the SASL session and the same instance cannot be reopened. The
-    transport is therefore created lazily in ``__enter__`` rather than held
-    on the class, which lets the ``_HiveClient`` itself live for the
-    duration of a ``HiveCatalog`` while still honouring the SASL transport
-    lifecycle.
-    """
-
-    _transport: "TTransport | None"
+    _transport: TTransport
     _ugi: list[str] | None
+    _was_opened: bool
 
     def __init__(
         self,
@@ -166,7 +159,8 @@ def __init__(
         self._kerberos_auth = kerberos_auth
         self._kerberos_service_name = kerberos_service_name
         self._ugi = ugi.split(":") if ugi else None
-        self._transport = None
+        self._transport = self._init_thrift_transport()
+        self._was_opened = False
 
     def _init_thrift_transport(self) -> TTransport:
         url_parts = urlparse(self._uri)
@@ -177,28 +171,23 @@ def _init_thrift_transport(self) -> TTransport:
             return TTransport.TSaslClientTransport(socket, host=url_parts.hostname, service=self._kerberos_service_name)
 
     def _client(self) -> Client:
-        assert self._transport is not None  # set by __enter__
         protocol = TBinaryProtocol.TBinaryProtocol(self._transport)
         client = Client(protocol)
         if self._ugi:
             client.set_ugi(*self._ugi)
         return client
 
     def __enter__(self) -> Client:
-        """Initialize and open a fresh transport for this entry.
-
-        A new ``TSaslClientTransport`` is created on every entry because
-        the underlying SASL session cannot be reused after ``close()``.
-        Reusing a closed transport would leave a half-opened TCP socket
-        and SASL handshake EOF noise on the server.
-        """
-        self._transport = self._init_thrift_transport()
+        """Make sure the transport is initialized and open."""
+        if self._was_opened:
+            self._transport = self._init_thrift_transport()
         self._transport.open()
+        self._was_opened = True
         return self._client()
 
     def __exit__(self, exctype: type[BaseException] | None, excinst: BaseException | None, exctb: TracebackType | None) -> None:
         """Close transport if it was opened."""
-        if self._transport is not None and self._transport.isOpen():
+        if self._transport.isOpen():
             self._transport.close()
 
 

tests/catalog/test_hive.py

@@ -1412,22 +1412,14 @@ def test_create_hive_client_with_kerberos_using_context_manager(
             assert open_client._iprot.trans.isOpen()
 
 
-def test_kerberized_client_uses_fresh_transport_per_entry(
+def test_kerberized_client_uses_fresh_transport_on_reuse(
     kerberized_hive_metastore_fake_url: str,
 ) -> None:
-    """TSaslClientTransport is single-use — each entry must create a new one.
-
-    The underlying SASL session cannot be reused after close(). Each
-    context-manager entry must therefore produce a fresh transport
-    instance rather than reopening the previous one (which would leave
-    a half-opened TCP socket and SASL handshake EOF noise on the server).
-    """
+    """Reusing the context manager must reinitialize the transport."""
     client = _HiveClient(
         uri=kerberized_hive_metastore_fake_url,
         kerberos_auth=True,
     )
-    # No transport should exist until the context manager is entered.
-    assert client._transport is None
     with (
         patch(
             "puresasl.mechanisms.kerberos.authGSSClientStep",