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

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__)