diff --git a/src/databricks/sql/backend/databricks_client.py b/src/databricks/sql/backend/databricks_client.py index b772e7ddd..0c22b9e07 100644 --- a/src/databricks/sql/backend/databricks_client.py +++ b/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 diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index c6e66d661..dfa47d5f6 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/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 diff --git a/src/databricks/sql/backend/kernel/type_mapping.py b/src/databricks/sql/backend/kernel/type_mapping.py index fdbe1dde5..962e3d1a8 100644 --- a/src/databricks/sql/backend/kernel/type_mapping.py +++ b/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 diff --git a/src/databricks/sql/backend/sea/backend.py b/src/databricks/sql/backend/sea/backend.py index 04c79a18b..d9a991a49 100644 --- a/src/databricks/sql/backend/sea/backend.py +++ b/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 diff --git a/src/databricks/sql/backend/sea/queue.py b/src/databricks/sql/backend/sea/queue.py index 4a319c442..5ca070bc3 100644 --- a/src/databricks/sql/backend/sea/queue.py +++ b/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( diff --git a/src/databricks/sql/backend/types.py b/src/databricks/sql/backend/types.py index d080d0187..402cf6beb 100644 --- a/src/databricks/sql/backend/types.py +++ b/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, diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index d45d51181..2a86fb4ca 100755 --- a/src/databricks/sql/client.py +++ b/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 `` 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, diff --git a/src/databricks/sql/cloudfetch/download_manager.py b/src/databricks/sql/cloudfetch/download_manager.py index 27265720f..ece7406b7 100644 --- a/src/databricks/sql/cloudfetch/download_manager.py +++ b/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__) diff --git a/src/databricks/sql/cloudfetch/downloader.py b/src/databricks/sql/cloudfetch/downloader.py index e6d1c6d10..295f147cd 100644 --- a/src/databricks/sql/cloudfetch/downloader.py +++ b/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__) diff --git a/src/databricks/sql/parameters/native.py b/src/databricks/sql/parameters/native.py index d0fb8d82c..d0d811940 100644 --- a/src/databricks/sql/parameters/native.py +++ b/src/databricks/sql/parameters/native.py @@ -1,20 +1,45 @@ +from __future__ import annotations + import datetime import decimal from enum import Enum, auto -from typing import Optional, Sequence, Any +from typing import Optional, Sequence, Any, TYPE_CHECKING from databricks.sql.exc import NotSupportedError -from databricks.sql.thrift_api.TCLIService.ttypes import ( - TSparkParameter, - TSparkParameterValue, - TSparkParameterValueArg, -) -import datetime -import decimal -from enum import Enum, auto +if TYPE_CHECKING: + # Type-annotation-only imports. ``from __future__ import annotations`` keeps + # these out of the runtime import graph; the ``TSparkParameter*`` objects + # are constructed via function-local imports inside the ``as_tspark_param`` + # / ``_tspark_*`` helpers below, which only run on the Thrift execute path. + # This keeps the Apache Thrift ``thrift`` package out of the SEA/kernel + # load path (see ``test_lazy_thrift_import``). + from databricks.sql.thrift_api.TCLIService.ttypes import ( + TSparkParameter, + TSparkParameterValue, + TSparkParameterValueArg, + ) + from typing import Dict, List, Union +# Names historically re-exported from this module that actually live in the +# Thrift-generated ``ttypes``. They are resolved lazily via ``__getattr__`` +# (PEP 562) so ``from databricks.sql.parameters.native import TSparkParameter`` +# keeps working for existing callers without importing the Apache Thrift +# ``thrift`` package at module load -- which is what keeps the SEA/kernel path +# Thrift-free. See ``test_lazy_thrift_import``. +_LAZY_THRIFT_REEXPORTS = frozenset( + {"TSparkParameter", "TSparkParameterValue", "TSparkParameterValueArg"} +) + + +def __getattr__(name: str) -> Any: + if name in _LAZY_THRIFT_REEXPORTS: + from databricks.sql.thrift_api.TCLIService import ttypes + + return getattr(ttypes, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + class ParameterApproach(Enum): INLINE = 1 @@ -98,6 +123,8 @@ class DbsqlParameterBase: def as_tspark_param(self, named: bool) -> TSparkParameter: """Returns a TSparkParameter object that can be passed to the DBR thrift server.""" + from databricks.sql.thrift_api.TCLIService.ttypes import TSparkParameter + tsp = TSparkParameter(value=self._tspark_param_value(), type=self._cast_expr()) if named: @@ -108,10 +135,18 @@ def as_tspark_param(self, named: bool) -> TSparkParameter: return tsp def _tspark_param_value(self): + from databricks.sql.thrift_api.TCLIService.ttypes import ( + TSparkParameterValue, + ) + return TSparkParameterValue(stringValue=str(self.value)) def _tspark_value_arg(self): """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server.""" + from databricks.sql.thrift_api.TCLIService.ttypes import ( + TSparkParameterValueArg, + ) + return TSparkParameterValueArg(value=str(self.value), type=self._cast_expr()) def _cast_expr(self): @@ -470,6 +505,8 @@ def __init__(self, value: Sequence[Any], name: Optional[str] = None): def as_tspark_param(self, named: bool = False) -> TSparkParameter: """Returns a TSparkParameter object that can be passed to the DBR thrift server.""" + from databricks.sql.thrift_api.TCLIService.ttypes import TSparkParameter + tsp = TSparkParameter(type=self._cast_expr()) tsp.arguments = [val._tspark_value_arg() for val in self.value] @@ -482,6 +519,10 @@ def as_tspark_param(self, named: bool = False) -> TSparkParameter: def _tspark_value_arg(self): """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server.""" + from databricks.sql.thrift_api.TCLIService.ttypes import ( + TSparkParameterValueArg, + ) + tva = TSparkParameterValueArg(type=self._cast_expr()) tva.arguments = [val._tspark_value_arg() for val in self.value] return tva @@ -519,6 +560,8 @@ def __init__(self, value: dict, name: Optional[str] = None): def as_tspark_param(self, named: bool = False) -> TSparkParameter: """Returns a TSparkParameter object that can be passed to the DBR thrift server.""" + from databricks.sql.thrift_api.TCLIService.ttypes import TSparkParameter + tsp = TSparkParameter(type=self._cast_expr()) tsp.arguments = [val._tspark_value_arg() for val in self.value] if named: @@ -530,6 +573,10 @@ def as_tspark_param(self, named: bool = False) -> TSparkParameter: def _tspark_value_arg(self): """Returns a TSparkParameterValueArg object that can be passed to the DBR thrift server.""" + from databricks.sql.thrift_api.TCLIService.ttypes import ( + TSparkParameterValueArg, + ) + tva = TSparkParameterValueArg(type=self._cast_expr()) tva.arguments = [val._tspark_value_arg() for val in self.value] return tva diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index c382103ea..69c2dcb1c 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -1,8 +1,7 @@ import logging import re -from typing import Dict, Tuple, List, Optional, Any, Type +from typing import Dict, Tuple, List, Optional, Any, Type, TYPE_CHECKING -from databricks.sql.thrift_api.TCLIService import ttypes from databricks.sql.types import SSLOptions from databricks.sql.auth.auth import get_python_sql_connector_auth_provider from databricks.sql.auth.authenticators import AccessTokenAuthProvider @@ -10,16 +9,38 @@ from databricks.sql.exc import SessionAlreadyClosedError, DatabaseError, RequestError from databricks.sql import __version__ from databricks.sql import USER_AGENT_NAME -from databricks.sql.backend.thrift_backend import ThriftDatabricksClient -from databricks.sql.backend.sea.backend import SeaDatabricksClient from databricks.sql.backend.databricks_client import DatabricksClient from databricks.sql.backend.types import SessionId, BackendType from databricks.sql.common.unified_http_client import UnifiedHttpClient from databricks.sql.common.agent import detect as detect_agent +if TYPE_CHECKING: + from databricks.sql.backend.thrift_backend import ThriftDatabricksClient + from databricks.sql.backend.sea.backend import SeaDatabricksClient + logger = logging.getLogger(__name__) +# The backend client classes are resolved lazily (PEP 562) rather than imported +# at module load. ``ThriftDatabricksClient`` pulls in the Apache Thrift +# ``thrift`` package, so importing it eagerly would drag ``thrift`` into the +# SEA and kernel connect paths (breaking build systems that ship their own +# ``thrift``, e.g. Buck). Exposing them as module attributes -- rather than as +# function-local imports inside ``open`` -- also keeps the long-standing test +# seam ``patch("databricks.sql.session.ThriftDatabricksClient")`` working. +# See ``test_lazy_thrift_import``. +def __getattr__(name: str) -> Any: + if name == "ThriftDatabricksClient": + from databricks.sql.backend.thrift_backend import ThriftDatabricksClient + + return ThriftDatabricksClient + if name == "SeaDatabricksClient": + from databricks.sql.backend.sea.backend import SeaDatabricksClient + + return SeaDatabricksClient + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + class Session: def __init__( self, @@ -205,13 +226,19 @@ def _create_backend( retry_options=kernel_retry_options, ) + # These reference the lazily-resolved module attributes defined via + # ``__getattr__`` above (or a test's ``patch(...)`` of them), so neither + # backend class -- and in particular ``thrift`` -- is imported until the + # branch that actually needs it runs. + import databricks.sql.session as _session_module + databricks_client_class: Type[DatabricksClient] if self.use_sea: logger.debug("Creating SEA backend client") - databricks_client_class = SeaDatabricksClient + databricks_client_class = _session_module.SeaDatabricksClient else: logger.debug("Creating Thrift backend client") - databricks_client_class = ThriftDatabricksClient + databricks_client_class = _session_module.ThriftDatabricksClient common_args = { "server_hostname": server_hostname, @@ -315,6 +342,12 @@ def get_protocol_version(session_id: SessionId): @staticmethod def server_parameterized_queries_enabled(protocolVersion): + # Function-local import: the protocol-version constant lives in the + # Thrift-generated ttypes, but this check only ever runs with a + # Thrift-negotiated protocol version, so deferring keeps ``thrift`` out + # of the SEA/kernel load path. + from databricks.sql.thrift_api.TCLIService import ttypes + if ( protocolVersion and protocolVersion >= ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8 diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 7e2abc07d..0914bb168 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -1,5 +1,5 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional, Tuple, Union, Sequence +from typing import Any, Dict, List, Optional, Tuple, Union, Sequence, TYPE_CHECKING from dateutil import parser import datetime @@ -21,11 +21,18 @@ from databricks.sql import OperationalError from databricks.sql.exc import ProgrammingError from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager -from databricks.sql.thrift_api.TCLIService.ttypes import ( - TRowSet, - TSparkArrowResultLink, - TSparkRowSetType, -) + +if TYPE_CHECKING: + # Type-annotation-only imports (deferred by ``from __future__ import + # annotations``). ``TSparkRowSetType`` is used at runtime only inside + # ``ThriftResultSetQueueFactory.build_queue``, which is a Thrift-only code + # path and imports it locally. This keeps the Apache Thrift ``thrift`` + # package out of the SEA/kernel load path (see ``test_lazy_thrift_import``). + from databricks.sql.thrift_api.TCLIService.ttypes import ( + TRowSet, + TSparkArrowResultLink, + TSparkRowSetType, + ) from databricks.sql.types import SSLOptions from databricks.sql.backend.types import CommandId from databricks.sql.telemetry.models.event import StatementType @@ -99,6 +106,10 @@ def build_queue( ResultSetQueue """ + # Function-local import: this factory is only used by the Thrift + # backend, so deferring keeps ``thrift`` out of the SEA/kernel path. + from databricks.sql.thrift_api.TCLIService.ttypes import TSparkRowSetType + if row_set_type == TSparkRowSetType.ARROW_BASED_SET: arrow_table, n_valid_rows = convert_arrow_based_set_to_arrow_table( t_row_set.arrowBatches, lz4_compressed, arrow_schema_bytes diff --git a/tests/unit/test_lazy_thrift_import.py b/tests/unit/test_lazy_thrift_import.py new file mode 100644 index 000000000..5763a91c6 --- /dev/null +++ b/tests/unit/test_lazy_thrift_import.py @@ -0,0 +1,164 @@ +"""Regression tests locking in lazy loading of the Apache Thrift runtime. + +The Python SQL connector historically imported the PyPI ``thrift`` package +eagerly whenever ``connect()`` (i.e. ``databricks.sql.client``) was imported, +regardless of which backend the caller selected. 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. + +These tests assert that importing the connector and the non-Thrift backends +does NOT import ``thrift``, while the Thrift backend still does. The invariant +is easy to regress silently -- a single top-level ``from ...ttypes import X`` +in any module on the ``client.py`` import chain re-poisons the whole path -- so +each check runs in a *fresh* subprocess interpreter (``sys.modules`` is +process-global; other tests in the same process may already have imported +thrift, which would mask a regression). +""" + +import subprocess +import sys +import textwrap + +import pytest + +# Sentinel exit codes emitted by the child *only after* the import completes, +# so an interpreter crash / uncaught exception (Python's generic exit code 1) +# can never be misread as a definitive "thrift (not) loaded" answer. +_EXIT_IMPORTED_NO_THRIFT = 10 +_EXIT_IMPORTED_WITH_THRIFT = 11 +_EXIT_IMPORT_FAILED = 12 + + +class _ImportProbeResult: + """Outcome of importing a module in a clean subprocess.""" + + def __init__(self, returncode: int, stderr: str): + self.returncode = returncode + self.stderr = stderr + + @property + def imported(self) -> bool: + return self.returncode in ( + _EXIT_IMPORTED_NO_THRIFT, + _EXIT_IMPORTED_WITH_THRIFT, + ) + + @property + def thrift_loaded(self) -> bool: + return self.returncode == _EXIT_IMPORTED_WITH_THRIFT + + @property + def import_failed(self) -> bool: + return self.returncode == _EXIT_IMPORT_FAILED + + +def _probe_import(module_name: str) -> _ImportProbeResult: + """Import ``module_name`` in a clean subprocess and report, via a dedicated + sentinel exit code, whether the top-level ``thrift`` package ended up in + ``sys.modules`` -- distinguishing that from an import failure (e.g. a + missing optional dependency such as pyarrow), which is reported separately + rather than being conflated with "thrift was imported".""" + script = textwrap.dedent( + f""" + import sys + + try: + import {module_name} # noqa: F401 + except BaseException: + import traceback + traceback.print_exc() + sys.exit({_EXIT_IMPORT_FAILED}) + + sys.exit( + {_EXIT_IMPORTED_WITH_THRIFT} + if "thrift" in sys.modules + else {_EXIT_IMPORTED_NO_THRIFT} + ) + """ + ) + proc = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + ) + if proc.returncode not in ( + _EXIT_IMPORTED_NO_THRIFT, + _EXIT_IMPORTED_WITH_THRIFT, + _EXIT_IMPORT_FAILED, + ): + raise AssertionError( + f"subprocess importing {module_name!r} exited with unexpected code " + f"{proc.returncode}. stderr:\n{proc.stderr}" + ) + return _ImportProbeResult(proc.returncode, proc.stderr) + + +# Modules on the connect()/execute() path that must stay Thrift-free so the +# SEA and kernel backends can be used without the ``thrift`` package present. +THRIFT_FREE_MODULES = [ + "databricks.sql", + "databricks.sql.client", + "databricks.sql.session", + "databricks.sql.utils", + "databricks.sql.parameters.native", + "databricks.sql.backend.types", + "databricks.sql.backend.databricks_client", + "databricks.sql.backend.sea.backend", + "databricks.sql.backend.sea.queue", + "databricks.sql.backend.kernel.type_mapping", + "databricks.sql.backend.kernel.client", + "databricks.sql.cloudfetch.downloader", + "databricks.sql.cloudfetch.download_manager", +] + + +@pytest.mark.parametrize("module_name", THRIFT_FREE_MODULES) +def test_module_does_not_import_thrift(module_name): + """Importing the connector and its non-Thrift backends must not import the + Apache Thrift runtime. + + This is what unblocks callers (e.g. Buck-based builds) that ship their own + ``thrift`` package and use only the SEA or kernel backend. + """ + result = _probe_import(module_name) + + if result.import_failed: + # A module that can't even be imported in this environment (typically a + # missing *optional* dependency, e.g. pyarrow in the "default deps" CI + # job) can't leak thrift. Skip rather than fail so this test stays + # focused on the thrift invariant and doesn't double as an + # optional-dependency presence check. + pytest.skip( + f"{module_name!r} could not be imported in this environment " + f"(likely a missing optional dependency); import error:\n" + f"{result.stderr}" + ) + + assert not result.thrift_loaded, ( + f"Importing {module_name!r} pulled in the top-level 'thrift' package. " + f"Something on this import chain grew a module-level " + f"'from databricks.sql.thrift_api...' / 'import thrift' statement (or a " + f"non-deferred type annotation). Move it under TYPE_CHECKING (with " + f"'from __future__ import annotations') or into a function body on the " + f"Thrift-only code path." + ) + + +def test_thrift_backend_still_imports_thrift(): + """Sanity check the counterpart invariant: the Thrift backend legitimately + depends on the Thrift runtime, so it must still import it. This guards + against a future 'fix' that hides thrift so aggressively the Thrift path + breaks. + + A failure to import the module (as opposed to importing it without thrift) + is surfaced explicitly rather than being treated as a pass.""" + result = _probe_import("databricks.sql.backend.thrift_backend") + + assert result.imported, ( + "The Thrift backend could not be imported at all -- the Thrift code " + f"path is broken. Import error:\n{result.stderr}" + ) + assert result.thrift_loaded, ( + "The Thrift backend no longer imports the 'thrift' package; the Thrift " + "code path is likely broken." + )