Lazy-load Apache Thrift so the SEA/kernel paths never import it (#907)

* Lazy-load Apache Thrift so SEA/kernel paths never import it

The connector declared `thrift` as a hard dependency and imported it
eagerly the moment `connect()` loaded `databricks.sql.client` -- 16
thrift modules, including the top-level `thrift` package -- regardless
of whether the caller selected `use_sea=True` or `use_kernel=True`.
Build systems that vendor their own `thrift` (e.g. Meta's Buck) then hit
a namespace collision even on the SEA / kernel code paths, which never
speak Thrift on the wire.

This makes the thrift import lazy without changing any public API or
install semantics: `thrift` stays a base dependency, but it is only
imported when the Thrift backend is actually used.

Mechanism, per module on the connect/execute chain:
- Annotation-only uses (the cloud-fetch download manager/downloader,
  the DatabricksClient ABC's `execute_command`, and the SEA/kernel
  `TSparkParameter` annotations): add `from __future__ import
  annotations` and move the thrift import under `TYPE_CHECKING`, so the
  annotations are never evaluated at runtime.
- Runtime uses on Thrift-only paths (`from_thrift_state`, the queue
  factory's `TSparkRowSetType`, `TProtocolVersion`, the SEA
  `_convert_to_thrift_link` construction, and every `TSparkParameter*`
  construction in parameters/native): move the import into the function
  body.
- Backend selection in `Session.open`: resolve `ThriftDatabricksClient`
  / `SeaDatabricksClient` via a module-level `__getattr__` (PEP 562) and
  reference them through the module namespace, so thrift is imported
  only on the branch that needs it. This also preserves the
  `patch("...session.ThriftDatabricksClient")` test seam.
- Preserve the historical re-exports `parameters.native.TSparkParameter*`
  and `client.ThriftDatabricksClient` via lazy `__getattr__` so existing
  importers (and tests) keep working without importing thrift at load.

Add tests/unit/test_lazy_thrift_import.py, which imports the connector
and each non-Thrift backend in a fresh subprocess and asserts the
top-level `thrift` package is absent from sys.modules (and, conversely,
that the Thrift backend still imports it). This locks the invariant --
a single stray module-level thrift import re-poisons the whole path.

Verified empirically: importing `databricks.sql.client` and the SEA /
kernel backend modules loads zero thrift modules, while the Thrift
backend still loads thrift. Full unit suite passes (the two failures
present also reproduce unchanged on main: a kernel test-ordering issue
and a realkernel-marked test).

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Co-authored-by: Isaac

* Address PR review: robust guard test + preserve client.py re-exports

Follow-up to the lazy-thrift change addressing two review-bot findings.

Medium (test robustness / CI false-positive): the guard test's child
subprocess exited with code 1 on any failure, which is also Python's
generic uncaught-exception code -- so an import failure was
indistinguishable from "thrift was imported". This caused a false CI
failure: in the "default deps" job (no pyarrow), importing
`kernel.type_mapping` raises `ModuleNotFoundError: pyarrow` (exit 1),
which the test misread as a thrift leak. It was also a false *pass* risk
for the Thrift-backend counterpart test. The child now emits dedicated
sentinel exit codes only after the import completes, captures stderr,
and reports import failure separately; the parametrized test skips
modules that can't import due to a missing optional dependency (rather
than failing), while still detecting a genuine thrift leak.

Low (back-compat): `client.py`'s `__getattr__` only re-exported
`ThriftDatabricksClient`. Extend it to also lazily resolve the other
names `client.py` historically exposed as importable
(`ThriftResultSet`, `TOpenSessionResp`, `TSparkParameter`,
`TOperationState`), so `from databricks.sql.client import <name>` keeps
working without importing the `thrift` package at module load.

Verified: importing `databricks.sql.client` still loads zero thrift
modules; touching any re-export resolves correctly (and only then pulls
thrift). Guard test passes with full deps (14/14) and correctly skips
the kernel modules when pyarrow/kernel are absent.

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Co-authored-by: Isaac

---------

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

src/databricks/sql/backend/databricks_client.py

@@ -7,7 +7,13 @@
     from databricks.sql.client import Cursor
     from databricks.sql.result_set import ResultSet
 
-from databricks.sql.thrift_api.TCLIService import ttypes
+    # Type-annotation-only import (deferred by ``from __future__ import
+    # annotations``). ``execute_command`` is typed with ``TSparkParameter`` for
+    # backwards compatibility, but this abstract base -- and the SEA/kernel
+    # implementations of it -- never import the Apache Thrift ``thrift`` package
+    # at load time. See ``test_lazy_thrift_import``.
+    from databricks.sql.thrift_api.TCLIService import ttypes
+
 from databricks.sql.backend.types import SessionId, CommandId, CommandState
 
 

src/databricks/sql/backend/kernel/client.py

@@ -47,12 +47,19 @@
     NotSupportedError,
     ProgrammingError,
 )
-from databricks.sql.thrift_api.TCLIService import ttypes
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
     from databricks.sql.result_set import ResultSet
 
+    # Type-annotation-only import (deferred by ``from __future__ import
+    # annotations``). ``execute_command`` accepts the Thrift-shaped
+    # ``TSparkParameter`` for interface compatibility and forwards it to
+    # ``bind_tspark_params``, which only reads its attributes; the kernel
+    # backend never imports the Apache Thrift ``thrift`` package. See
+    # ``test_lazy_thrift_import``.
+    from databricks.sql.thrift_api.TCLIService import ttypes
+
 logger = logging.getLogger(__name__)
 
 # Headers the kernel manages itself and that the connector must NOT

src/databricks/sql/backend/kernel/type_mapping.py

@@ -21,13 +21,19 @@
 
 from __future__ import annotations
 
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, TYPE_CHECKING
 
 import pyarrow
 
 from databricks.sql.backend.sea.utils.conversion import SqlType
 from databricks.sql.exc import NotSupportedError
-from databricks.sql.thrift_api.TCLIService import ttypes
+
+if TYPE_CHECKING:
+    # Type-annotation-only import (deferred by ``from __future__ import
+    # annotations``). ``bind_tspark_params`` only reads ``TSparkParameter``
+    # attributes (duck-typed) at runtime, so the kernel backend never imports
+    # the Apache Thrift ``thrift`` package. See ``test_lazy_thrift_import``.
+    from databricks.sql.thrift_api.TCLIService import ttypes
 
 # Type names that the connector emits as compound TSparkParameter
 # shapes (payload on ``arguments``, not ``value``). The kernel's

src/databricks/sql/backend/sea/backend.py

@@ -20,11 +20,17 @@
     MetadataCommands,
 )
 from databricks.sql.backend.sea.utils.normalize import normalize_sea_type_to_thrift
-from databricks.sql.thrift_api.TCLIService import ttypes
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
 
+    # Type-annotation-only import (deferred by ``from __future__ import
+    # annotations``). ``execute_command`` accepts the Thrift-shaped
+    # ``TSparkParameter`` for interface compatibility, but only reads its
+    # attributes (duck-typed) at runtime, so the SEA backend never imports the
+    # Apache Thrift ``thrift`` package. See ``test_lazy_thrift_import``.
+    from databricks.sql.thrift_api.TCLIService import ttypes
+
 from databricks.sql.backend.sea.result_set import SeaResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient

src/databricks/sql/backend/sea/queue.py

@@ -23,9 +23,16 @@
         ResultData,
         ResultManifest,
     )
+
+    # Type-annotation-only import (deferred by ``from __future__ import
+    # annotations``). The SEA backend reuses the Thrift ``TSparkArrowResultLink``
+    # only as the payload the shared cloud-fetch download manager expects; it is
+    # constructed via a function-local import in ``_convert_to_thrift_link`` so
+    # importing the SEA backend never imports the Apache Thrift ``thrift``
+    # package. See ``test_lazy_thrift_import``.
+    from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.backend.sea.utils.constants import ResultFormat
 from databricks.sql.exc import ProgrammingError, ServerOperationError
-from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.types import SSLOptions
 from databricks.sql.utils import (
     ArrowQueue,
@@ -262,6 +269,10 @@ def get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]:
     @staticmethod
     def _convert_to_thrift_link(link: ExternalLink) -> TSparkArrowResultLink:
         """Convert SEA external links to Thrift format for compatibility with existing download manager."""
+        from databricks.sql.thrift_api.TCLIService.ttypes import (
+            TSparkArrowResultLink,
+        )
+
         # Parse the ISO format expiration time
         expiry_time = int(dateutil.parser.parse(link.expiration).timestamp())
         return TSparkArrowResultLink(

src/databricks/sql/backend/types.py

@@ -1,11 +1,21 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, List, Optional, Any, Tuple
+from typing import Dict, List, Optional, Any, Tuple, TYPE_CHECKING
 import logging
 
 from databricks.sql.backend.utils.guid_utils import guid_to_hex_id
 from databricks.sql.telemetry.models.enums import StatementType
-from databricks.sql.thrift_api.TCLIService import ttypes
+
+if TYPE_CHECKING:
+    # Type-annotation-only import (evaluated lazily thanks to
+    # ``from __future__ import annotations``). The runtime uses of ``ttypes``
+    # in this module are function-local imports inside the Thrift-only code
+    # paths (``from_thrift_state``, ``to_thrift_handle``,
+    # ``to_operation_handle``), so importing this module never pulls in the
+    # Apache Thrift ``thrift`` package. See ``test_lazy_thrift_import``.
+    from databricks.sql.thrift_api.TCLIService import ttypes
 
 logger = logging.getLogger(__name__)
 
@@ -60,6 +70,11 @@ def from_thrift_state(
             - CANCELED_STATE -> CANCELLED
         """
 
+        # Function-local import: this classmethod is only ever called from the
+        # Thrift backend, so deferring the import keeps ``thrift`` out of the
+        # SEA/kernel load path.
+        from databricks.sql.thrift_api.TCLIService import ttypes
+
         if state in (
             ttypes.TOperationState.INITIALIZED_STATE,
             ttypes.TOperationState.PENDING_STATE,

src/databricks/sql/client.py

@@ -1,5 +1,17 @@
+from __future__ import annotations
+
 import time
-from typing import Dict, Tuple, List, Optional, Any, Union, Sequence, BinaryIO
+from typing import (
+    Dict,
+    Tuple,
+    List,
+    Optional,
+    Any,
+    Union,
+    Sequence,
+    BinaryIO,
+    TYPE_CHECKING,
+)
 import pandas
 
 try:
@@ -25,8 +37,6 @@
     DatabaseError,
 )
 
-from databricks.sql.thrift_api.TCLIService import ttypes
-from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.utils import (
     ParamEscaper,
@@ -49,7 +59,7 @@
     ParameterApproach,
 )
 
-from databricks.sql.result_set import ResultSet, ThriftResultSet
+from databricks.sql.result_set import ResultSet
 from databricks.sql.types import Row, SSLOptions
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
 from databricks.sql.experimental.oauth_persistence import OAuthPersistence
@@ -60,11 +70,18 @@
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
 from databricks.sql.common.http import HttpMethod
 
-from databricks.sql.thrift_api.TCLIService.ttypes import (
-    TOpenSessionResp,
-    TSparkParameter,
-    TOperationState,
-)
+if TYPE_CHECKING:
+    # Type-annotation-only imports (deferred by ``from __future__ import
+    # annotations``). ``get_protocol_version`` and ``_prepare_native_parameters``
+    # are typed with these Thrift-generated types, but the Thrift backend and
+    # its result set are imported lazily (only on the Thrift connect path), so
+    # importing this module -- and connecting with the SEA or kernel backend --
+    # never imports the Apache Thrift ``thrift`` package. See
+    # ``test_lazy_thrift_import``.
+    from databricks.sql.thrift_api.TCLIService.ttypes import (
+        TOpenSessionResp,
+        TSparkParameter,
+    )
 from databricks.sql.telemetry.telemetry_client import (
     TelemetryHelper,
     TelemetryClientFactory,
@@ -95,6 +112,33 @@
 TRANSACTION_ISOLATION_LEVEL_REPEATABLE_READ = "REPEATABLE_READ"
 
 
+def __getattr__(name: str) -> Any:
+    """Lazily resolve Thrift-related names that ``client.py`` used to expose as
+    real top-level imports.
+
+    ``client.py`` itself never instantiates these (the backend is chosen in
+    ``Session.open``), but they are resolved here as module attributes so
+    ``from databricks.sql.client import <name>`` keeps working for any existing
+    caller -- and the ``patch("databricks.sql.client.ThriftDatabricksClient")``
+    test seam is preserved -- without importing the Apache Thrift ``thrift``
+    package at module load, which is what keeps the SEA/kernel connect path
+    Thrift-free (see ``test_lazy_thrift_import``).
+    """
+    if name == "ThriftDatabricksClient":
+        from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
+
+        return ThriftDatabricksClient
+    if name == "ThriftResultSet":
+        from databricks.sql.result_set import ThriftResultSet
+
+        return ThriftResultSet
+    if name in ("TOpenSessionResp", "TSparkParameter", "TOperationState"):
+        from databricks.sql.thrift_api.TCLIService import ttypes
+
+        return getattr(ttypes, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 class Connection:
     def __init__(
         self,

src/databricks/sql/cloudfetch/download_manager.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import logging
 
 from concurrent.futures import ThreadPoolExecutor, Future
-from typing import List, Union, Tuple, Optional
+from typing import List, Union, Tuple, Optional, TYPE_CHECKING
 
 from databricks.sql.cloudfetch.downloader import (
     ResultSetDownloadHandler,
@@ -10,7 +12,12 @@
 )
 from databricks.sql.types import SSLOptions
 from databricks.sql.telemetry.models.event import StatementType
-from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
+
+if TYPE_CHECKING:
+    # Type-annotation-only import; see the note in downloader.py. Keeping the
+    # ``thrift`` package out of this module lets the SEA/kernel backends use the
+    # cloud-fetch download manager without importing Apache Thrift.
+    from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 
 logger = logging.getLogger(__name__)
 

src/databricks/sql/cloudfetch/downloader.py

@@ -1,17 +1,26 @@
+from __future__ import annotations
+
 import logging
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 import lz4.frame
 import time
 from databricks.sql.common.http import HttpMethod
-from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.exc import Error
 from databricks.sql.types import SSLOptions
 from databricks.sql.telemetry.latency_logger import log_latency
 from databricks.sql.telemetry.models.event import StatementType
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
 
+if TYPE_CHECKING:
+    # Imported for type annotations only. ``from __future__ import annotations``
+    # makes every annotation a string, so this import is never evaluated at
+    # runtime -- which keeps the (Apache Thrift) ``thrift`` package out of the
+    # cloud-fetch code path used by the SEA and kernel backends. See the
+    # ``test_lazy_thrift_import`` regression test.
+    from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
+
 logger = logging.getLogger(__name__)