Lazy-load Apache Thrift so the SEA/kernel paths never import it - #907
Lazy-load Apache Thrift so the SEA/kernel paths never import it#907vikrantpuppala wants to merge 2 commits into
Conversation
Sorry, something went wrong.
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).
Co-authored-by: Isaac
There was a problem hiding this comment.
Verdict: 1 Medium · 1 Low
Solid, well-scoped lazy-import refactor — the annotation-only moves are all correctly guarded by from __future__ import annotations, every runtime ttypes use on Thrift-only paths has a function-local import, and the patch(...) test seams are preserved via PEP 562 __getattr__. Two items: a Medium test-robustness bug where the new regression test's exit-code scheme can mask an import failure (notably a false-pass for the Thrift-backend counterpart test), and a Low back-compat note on dropped client.py re-exports.
Sorry, something went wrong.
| sys.exit(1 if "thrift" in sys.modules else 0) | ||
| """ | ||
| ) | ||
| result = subprocess.run([sys.executable, "-c", script]) |
There was a problem hiding this comment.
🟡 Medium — The exit-code scheme conflates "the module raised during import" with "thrift was imported". If import {module_name} throws (e.g. an ImportError/NameError from a future refactor), the child process exits with code 1 — the same code this helper uses to mean "thrift is in sys.modules". So result.returncode not in (0, 1) never trips on an import crash (Python uses exit code 1 for uncaught exceptions), and the function returns True.
Consequences:
- For the THRIFT_FREE_MODULES cases this only produces a misleading failure message ("pulled in the top-level 'thrift' package") when the real cause is an import error — the test still fails, so it's tolerable.
- For test_thrift_backend_still_imports_thrift it is a false pass: if the Thrift backend ever breaks at import time, the child exits 1, the helper returns True, and the assertion assert _thrift_loaded_after_importing(...) succeeds — exactly the regression this test claims to guard against ("a future 'fix' that hides thrift so aggressively the Thrift path breaks") slips through silently.
Consider making the child emit unambiguous sentinel exit codes only after a successful import (e.g. wrap the import in try/except and exit with a dedicated code on failure, or exit 2/3 for imported/not-imported and treat everything else as an error). Passing capture_output=True and surfacing the child's stderr on unexpected codes would also make failures debuggable.
Sorry, something went wrong.
| preserves the long-standing test seam | ||
| ``patch("databricks.sql.client.ThriftDatabricksClient")``. | ||
| """ | ||
| if name == "ThriftDatabricksClient": |
There was a problem hiding this comment.
🔵 Low — The module-level __getattr__ only re-exports ThriftDatabricksClient. Before this PR, databricks.sql.client also exposed TOpenSessionResp, TSparkParameter, TOperationState (real top-level imports) and ThriftResultSet (imported from result_set). After this change, TOpenSessionResp/TSparkParameter live only under TYPE_CHECKING, TOperationState/ThriftResultSet are dropped entirely, and none are handled by __getattr__ — so from databricks.sql.client import TOperationState (or the other three) now raises ImportError at runtime where it previously worked.
The canonical locations still resolve (parameters.native.TSparkParameter* via its own lazy __getattr__, result_set.ThriftResultSet), and these were never documented public API, so impact is likely small. But since the PR deliberately preserved ThriftDatabricksClient and the parameters.native re-exports for back-compat, it's worth either extending this __getattr__ to cover the other historically-importable names or confirming they aren't part of the surface any external caller relies on.
Sorry, something went wrong.
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>
Reviewers
At least 1 approving review is required to merge this pull request.
Problem
The connector imports the PyPI thrift package eagerly the moment connect() loads 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. The SEA and kernel backends never speak Thrift on the wire, yet they still drag it in at connect time.
This breaks build systems that vendor their own thrift under the thrift top-level namespace (e.g. Meta's Buck): two packages own the same namespace and the build is rejected, even though the connector's SEA/kernel path never uses Apache Thrift.
Reproduction (before this change):
What this changes
Makes the thrift import lazy — imported only when the Thrift backend is actually used. No public API changes and no install-semantics changes: thrift stays a base dependency; it is simply not imported on the SEA/kernel path.
Per module on the connect/execute chain:
New test
tests/unit/test_lazy_thrift_import.py 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). sys.modules is process-global, so the subprocess isolation is what makes the check reliable — a single stray module-level thrift import anywhere on the chain re-poisons the whole path, and this test catches that.
Verification
Notes for reviewers / follow-ups
This pull request and its description were written by Isaac.