feat(backend/kernel): add use_kernel=True flag — route through the Rust kernel via PyO3 - #787
feat(backend/kernel): add use_kernel=True flag — route through the Rust kernel via PyO3#787vikrantpuppala merged 16 commits into
Conversation
Sorry, something went wrong.
|
Two updates since the initial PR: 1. Dropped External auth → PAT-only on the kernel backend (25723627). auth_bridge.py now routes PAT (including TokenFederationProvider-wrapped PAT) through the kernel's PAT path; everything else raises NotSupportedError pointing the user at use_sea=False. The kernel-side PR is no longer on this PR's critical path. Why: routing OAuth into the kernel requires per-request token resolution to keep refresh working — two viable mechanisms (kernel-native OAuth, or the External callback), both with costs. Punting the decision until there's pressure. 2. Live e2e tests moved into the repo (6b308156). The previous ad-hoc /tmp/connector_smoke.py is now tests/e2e/test_kernel_backend.py — a proper pytest module using the existing connection_details fixture. 11 tests cover: connect, SELECT 1, range(10000), fetchmany pacing, fetchall_arrow, all four metadata methods, session_configuration round-trip, structured DatabaseError on bad SQL. All 11 pass against dogfood in ~20s. Module skips cleanly when databricks_sql_kernel isn't installed or creds aren't set. The auth_bridge unit tests are updated: OAuth providers / ExternalAuthProvider now assert NotSupportedError. 39 unit tests pass. |
Sorry, something went wrong.
|
CI status, final: 33 / 34 checks pass. The one failing check (test-with-coverage, the e2e + coverage job) reports 7 test failures — but all 7 are pre-existing on main, unrelated to this PR. Verified by reproducing them on an independent recent PR (PR for thrift-result-set-heartbeat, run 25725422361, 2 days ago) and confirming the same 6 TestMstMetadata cases fail there too with identical error messages. Failure breakdown:
My PR contributions to this job (all working as intended):
I don't believe my PR should be responsible for fixing the 7 pre-existing failures — they need their own fixes (server-side investigation for MST metadata; switching fetch_rows to pytest assertions for the large-result test). Happy to file follow-up issues if you want. |
Sorry, something went wrong.
Code Review Squad — Failed Inline CommentsCould not post inline comments for: F2, F4, F7, F10, F11, F13, F14, F19, F22, F25 — see body below. F2 — Federated-PAT token refresh is dead — single snapshot at construct timeHigh severity — flagged by security + devils-advocate kernel_auth_kwargs extracts the bearer token once via auth_provider.add_headers({}) at construct time. The result is stored on KernelDatabricksClient._auth_kwargs and spread into _kernel.Session(...) at open_session. After that, the kernel uses that frozen token for the session lifetime. get_python_sql_connector_auth_provider wraps every provider in TokenFederationProvider. For a PAT whose issuer host differs from the workspace host, TokenFederationProvider._get_token (src/databricks/sql/auth/token_federation.py:131-153) performs an OAuth token-exchange and caches the exchanged token, refreshing it on subsequent add_headers calls when it expires. The bridge captures only the first exchanged token and never re-extracts. Long-running kernel sessions outlive the exchanged token's TTL and start failing Unauthenticated mid-session — even though the connector-side TokenFederationProvider would happily mint a fresh one. The bridge's own docstring acknowledges this failure mode while justifying OAuth rejection: "routing OAuth through PAT would silently break token refresh during long-running sessions." That exact failure mode applies to the federated-PAT case the bridge accepts. Recommend: either (a) propagate a refresh callback into the kernel Session, (b) block federation-wrapped PAT until the kernel exposes a refresh hook, or (c) refresh _auth_kwargs immediately before each kernel call (defeats the purpose of caching but is correct). F4 — get_tables with table_types silently returns wrong-answer resultHigh severity — flagged by devils-advocate, agent-compat, ops, security The warning says "client-side table_types filter not yet implemented … returning unfiltered rows for %r" — and then the code passes table_types=table_types to _kernel_session.metadata().list_tables(...) anyway on the next line. Either the kernel filters server-side (in which case the warning is wrong and misleading — callers think they're getting unfiltered rows when they aren't) or it doesn't (the kwarg passthrough is dead). Both states are bad:
logger.warning is also invisible to programmatic consumers — it's not actionable. Recommend: until the filter is correctly implemented, raise NotSupportedError when table_types is non-empty. Same pattern as parameters and query_tags rejection elsewhere in this client. A logger.warning plus silent wrong-answer is the worst of both worlds. F7 — get_tables warning is log spam at metadata-call frequencyHigh severity — flagged by ops, agent-compat, architecture logger.warning fires on every get_tables call with a non-empty table_types. Metadata calls are made by every BI tool that browses a workspace: Power BI, Tableau, dbt, DataGrip, etc. — all call getTables(table_types=["TABLE", "VIEW"]) on connection open and every catalog-tree refresh. At Databricks customer scale this is easily 100s of QPS per workspace. WARN-level log spam at that rate triggers alert fatigue for any operator with a "WARN-or-above" alarm and dominates the kernel-backend portion of stdout/stderr at scale. Recommend: demote to logger.debug, or move to a once-per-session lazy-init flag (operator sees it once during the session, not on every call). Even logger.info is borderline at metadata-call frequency. F10 — Telemetry is structurally degraded for kernel-backed cursorsMedium severity — flagged by ops, architecture Three places the per-statement / per-chunk telemetry that downstream dashboards depend on is silently broken for use_sea=True:
Plus: connection-level DriverConnectionParameters still reports DatabricksClientType.SEA for kernel sessions (src/databricks/sql/client.py:379-381). Recommend: either (a) plumb a telemetry hook from the kernel through KernelResultSet (kernel exposes chunk-fetch / retry-count summaries that the connector forwards into SqlExecutionEvent), or (b) document this as a known-limitation banner in the PR description and acknowledge dashboards keyed on these fields under-report for use_sea=True. F11 — cancel_command is a no-op on the sync-execute path; get_query_state reports SUCCEEDEDMedium severity — flagged by ops, devils-advocate, agent-compat Sync executes never populate _async_handles. So calling cancel_command for a still-draining sync cursor logs at DEBUG and returns. Combined with get_query_state (client.py:543-549) returning CommandState.SUCCEEDED for unknown command_ids, a cursor.cancel() followed by get_query_state() falsely reports success — the query continues running server-side. The inline comment claims this matches Thrift's tolerant behavior — but Thrift's tolerance is for race-y double-cancels, not "cancel was called for a command we never tracked." Recommend: either (a) for sync-execute results, store the executed handle in a separate _sync_handles map keyed on command_id so cancel_command can call handle.cancel() on it, or (b) explicitly document that cursor.cancel() after execute() returned (i.e., during result drain) is a no-op on use_sea=True. Today it's both silently broken and not documented. F13 — get_columns(catalog_name=None) raises ProgrammingError; Thrift accepts None — silent semantic driftMedium severity — flagged by devils-advocate, agent-compat, security, architecture KernelDatabricksClient.get_columns raises ProgrammingError("get_columns requires catalog_name on the kernel backend.") when catalog_name is None. The Thrift backend (src/databricks/sql/backend/thrift_backend.py:1217-1244) accepts catalog_name=None and forwards it to the server. Cursor.columns() (src/databricks/sql/client.py:1567-1593) keeps catalog_name: Optional[str] = None in its public signature. Cross-backend code (third-party tooling doing schema discovery) that worked on use_sea=False will start raising on the same call under use_sea=True. The error message is actionable, but the divergence isn't documented in the Cursor.columns() docstring. Recommend: either (a) narrow the Cursor.columns() signature for kernel mode, (b) fall back to a per-catalog scan inside the kernel client, or (c) document the divergence on Cursor.columns() so users know to branch on backend. F14 — Synthetic metadata-{uuid} CommandIds pollute telemetry and cursor.query_idMedium severity — flagged by architecture, maintainability, devils-advocate _synthetic_command_id returns CommandId.from_sea_statement_id(f"metadata-{uuid.uuid4()}") — i.e., a CommandId(backend_type=SEA, guid="metadata-..."). Followed the consumers:
The synthetic ID also misrepresents the backend (BackendType.SEA when this is the kernel). Recommend: either (a) make the synthetic ID look like a UUID (uuid.uuid4().hex with no metadata- prefix — lose human-readability for telemetry safety), (b) add a CommandId.for_kernel_metadata() factory in backend/types.py so the synthetic-ness is explicit, or (c) don't set active_command_id for metadata calls (let query_id stay None as it does for any cursor before execute). F19 — Three near-identical return KernelResultSet(...) blocks should call _metadata_resultMedium severity — flagged by maintainability client.py:510-517 (sync execute), client.py:577-584 (get_execution_result), and client.py:588-596 (metadata) all build KernelResultSet with the same 6 args. The metadata path is already factored out as _metadata_result(stream, cursor, command_id). The other two should call it too — saves 12 lines and removes a divergence risk if a new kwarg is added to KernelResultSet. F22 — auth_bridge._is_pat duplicates TokenFederationProvider-peek logic from telemetryLow severity — flagged by maintainability _is_pat peeks through TokenFederationProvider to find a wrapped AccessTokenAuthProvider. The same pattern lives at src/databricks/sql/telemetry/telemetry_client.py:94-101: if isinstance(auth_provider, TokenFederationProvider):
return TelemetryHelper.get_auth_mechanism(auth_provider.external_provider)
if isinstance(auth_provider, AccessTokenAuthProvider):
return AuthMech.PATRecent commit cbd6a883 ("Telemetry: unwrap TokenFederationProvider…") shows the team explicitly cares about this surface. Recommend: extract auth/utils.py::unwrap_to_pat_provider(auth_provider) -> Optional[AccessTokenAuthProvider] shared by both call sites. Keeps them in lock-step when OAuth-through-kernel ships. Not a blocker. F25 — No control-char / CRLF sanitization on extracted bearer tokenLow severity — flagged by security _extract_bearer_token returns auth[len("Bearer "):] verbatim. AccessTokenAuthProvider.__init__ doesn't validate either. User-supplied access_token= strings containing \r\n, \0, or other control chars flow straight through to the kernel. If the kernel layer ever places this back into an HTTP header without scrubbing, it's a header-injection sink. Recommend: defense-in-depth — reject tokens matching [\x00-\x1f\x7f] at the connector boundary. Cheap. Summary: 17 of 27 findings posted as inline comments. 10 failed due to line-number validation (not in diff context). All critical/high/medium findings are covered — either by inline comment or in this summary. |
Sorry, something went wrong.
…ve review fixes
Major change: route the kernel backend through a new ``use_kernel=True``
connection kwarg instead of repurposing ``use_sea=True``. ``use_sea=True``
once again routes to the native pure-Python SEA backend (no behaviour
change); ``use_kernel=True`` routes to the Rust kernel via PyO3. The
two flags are mutually exclusive.
This addresses the largest reviewer concern from the multi-agent
review: silently hijacking a documented public flag broke OAuth /
federation / parameter-binding callers on ``use_sea=True`` who had no
opt-out. With the new flag, the kernel backend is fully opt-in and
existing ``use_sea=True`` users continue to get the native SEA backend
they signed up for.
Other substantive fixes:
- session.py: restore ``SeaDatabricksClient`` import + routing. Reject
``use_kernel=True`` + ``use_sea=True`` together with a clear
``ValueError``.
- client.py (kernel ``Cursor.columns``): update docstring to flag the
``catalog_name=None`` divergence — kernel requires a catalog,
Thrift / native SEA do not (F13).
- conftest.py: drop the collection-time ``pytest_collection_modifyitems``
hook that was skipping ``extra_params={"use_sea": True}`` cases. With
``use_sea=True`` back on the native SEA backend, those cases run as
they did before this PR (F8).
- kernel/client.py: ``get_tables`` now applies the ``table_types``
filter client-side using ``ResultSetFilter._filter_arrow_table``
(the same helper the native SEA backend uses), wrapped in a tiny
``_StaticArrowHandle`` that flows the filtered table back through
the normal ``KernelResultSet`` path. Replaces the previous
"log a warning and return unfiltered" behaviour (F4).
- kernel/client.py: guard ``_async_handles`` with ``threading.RLock``
so concurrent cursors on the same connection don't race on
submit / close / close-session (F15).
- kernel/result_set.py: ``KernelResultSet.close()`` now drops the
entry from ``backend._async_handles`` so async-submitted statements
don't leave stale references behind (F5).
- kernel/{__init__,client,auth_bridge}.py, tests/e2e/test_kernel_backend.py:
update docstrings, error messages, and the e2e fixture to refer to
``use_kernel=True`` instead of ``use_sea=True``.
- client.py (``Connection`` docstring): document the new
``use_kernel`` kwarg + its Phase-1 limitations.
New tests:
- tests/unit/test_kernel_client.py (38 cases): cover the 14-entry
``_CODE_TO_EXCEPTION`` table, ``_reraise_kernel_error`` attribute
forwarding, the 6-entry ``_STATE_TO_COMMAND_STATE`` table, the
no-open-session guards on every method, ``open_session`` double-open,
``parameters`` / ``query_tags`` rejection, ``get_columns``'
catalog-required check, ``cancel_command`` / ``close_command``
no-handle tolerance, ``get_query_state`` sync-path SUCCEEDED, the
Failed-state re-raise, the synthetic-command-id UUID shape, and
``close_session`` cleanup even when per-handle close errors fire.
Uses a fake ``databricks_sql_kernel`` module installed into
``sys.modules`` so the test runs with no Rust extension dependency
(F9).
77/77 kernel unit tests pass.
Co-authored-by: Isaac
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
- src/databricks/sql/backend/kernel/result_set.py: fix the 3 mypy errors at L237/239/241 by casting ``self.backend`` to ``KernelDatabricksClient`` (the base ``DatabricksClient`` doesn't declare ``_async_handles`` / ``_async_handles_lock``). Folds in gopalldb's nit (3249904284) — replace the explicit ``acquire()/try/finally/release()`` with a ``with`` block to match the rest of the file. - tests/e2e/test_kernel_backend.py: harden the module-level skip so the suite doesn't run when the kernel wheel is absent in CI. The unit suite installs a fake ``databricks_sql_kernel`` ``ModuleType`` into ``sys.modules`` so the connector's import-time ``import databricks_sql_kernel`` succeeds without the Rust extension; that fake leaks across into the same pytest session and ``pytest.importorskip`` happily returns it. A real wheel exposes ``__file__`` (compiled extension on disk); the fake does not. Skip the module when ``__file__`` is missing. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
m1 — install hint (comment 3249904266): The ``databricks-sql-kernel`` wheel is not yet published on PyPI; ``pip install databricks-sql-kernel`` either finds nothing or pulls a squatted package. Drop the misleading hint from ``ImportError`` and from the ``use_kernel`` docstring on ``databricks.sql.connect``; point users at the ``maturin develop --release`` dev path until the wheel ships. m4a — auth_bridge ValueError → ProgrammingError (comment 3249904276): Two sites in ``_extract_bearer_token`` / ``kernel_auth_kwargs`` were raising bare ``ValueError`` for caller-misuse cases (control chars in the token, PAT provider that produced no Authorization header). The rest of the kernel-backend error surface uses PEP 249 exception types — code paths that catch ``DatabaseError`` / ``ProgrammingError`` would miss these. Convert to ``ProgrammingError`` and update the unit test. m4b — description null_ok (comment 3249904282): ``description_from_arrow_schema`` was hardcoding the 7th tuple element to ``None`` even though ``pyarrow.Field.nullable`` is available. PEP 249 §Cursor.description defines ``null_ok`` as "True if NULL values are allowed"; callers branching on it would have lost useful information the kernel already provides. Now emits ``field.nullable``; added a unit test covering both nullable and non-nullable fields; updated the two existing tests that asserted the old all-``None`` shape. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Addresses gopalldb's four major / two minor remaining review comments. The shared error-mapping primitives move to a new ``_errors.py`` module so both ``client.py`` and ``result_set.py`` can use them without ``result_set.py`` importing from ``client.py``. M1 — async handle leak in get_execution_result (3249904251): ``ResultStream`` from ``await_result()`` is wrapped in ``KernelResultSet``; the underlying ``ExecutedAsyncStatement`` has no further role once the stream is in hand. Close it immediately, drop the entry from ``_async_handles``, and add the guid to ``_closed_commands``. A failed ``async_exec.close()`` is logged but doesn't break the result-set return — the kernel's Drop impl reaps server-side state. M2 — PyO3 native exceptions wrapped as OperationalError (3249904255): Added ``kernel_call(what)`` context manager (in the new ``_errors.py``). ``KernelError`` flows through ``reraise_kernel_error`` as before; anything else (``TypeError`` / ``OverflowError`` / ``ValueError`` from PyO3 argument conversion, extension-internal errors) is wrapped in ``OperationalError`` so DB-API callers only ever see PEP 249 exception types. Applied at every PyO3 call site: ``open_session``, ``execute_command``, ``cancel_command``, ``close_command``, ``get_query_state``, ``get_execution_result``, ``get_catalogs`` / ``get_schemas`` / ``get_tables`` / ``get_columns``, plus ``fetch_next_batch`` / ``arrow_schema`` in ``KernelResultSet``. M3 — KernelError during result-set construction (3249904259): ``KernelResultSet.__init__`` calls ``kernel_handle.arrow_schema()`` which can itself raise. Every call to ``_make_result_set`` is now inside a ``kernel_call`` scope so the schema-fetch error becomes a mapped PEP 249 exception instead of leaking raw ``KernelError``. m3 — get_query_state of a closed async command (3249904273): Added ``_closed_commands: Set[str]`` (guarded by the existing ``_async_handles_lock``). ``close_command`` records the guid; ``close_session`` records every swept guid; ``get_execution_result`` records its own command after closing the async_exec. ``get_query_state`` now returns ``CommandState.CLOSED`` instead of falling through to ``SUCCEEDED`` for these. m2 — unit test for get_tables table_types client-side filter (3249904269): Added ``test_get_tables_with_table_types_filters_rows`` and ``test_get_tables_without_table_types_returns_full_stream`` in ``tests/unit/test_kernel_client.py``. The first feeds a fake stream with mixed ``TABLE`` / ``VIEW`` rows and asserts only ``TABLE`` survives; the second confirms the no-filter path bypasses the drain-and-rewrap and returns all rows unchanged. Plus new tests for every change above: - test_pyo3_native_exception_wrapped_as_operational_error (M2) - test_pyo3_native_exception_wrapped_for_metadata_calls (M2) - test_kernel_error_during_result_set_construction_is_mapped (M3) - test_get_execution_result_closes_async_exec_and_drops_tracking (M1) - test_get_execution_result_does_not_raise_on_async_exec_close_failure (M1) - test_get_query_state_returns_closed_after_close_command (m3) - test_close_session_marks_swept_handles_as_closed (m3) 87/87 kernel unit tests pass (added 9). Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
|
Review round-trip — addressed gopalldb's 9 unresolved comments + CI-greening pass. Three commits since the last review pass:
M4 (GIL release in fetch_next_batch) verified — the kernel-side PyO3 binding routes every kernel call through py.detach(...) macros in pyo3/src/macros.rs, so the GIL is released across the network read. No connector-side change needed; left the verification trail in the comment thread. |
Sorry, something went wrong.
…plicit try/except
Using a ``with`` block for error translation is bad form: ``with``
conventionally signals resource lifecycle (locks, files), so
``with kernel_call("X"):`` hides the fact that the block raises a
mapped exception. Replace with explicit ``try/except Exception as
exc: raise _wrap_kernel_exception("X", exc) from exc`` at every
PyO3 call site.
What changed:
- ``_errors.py``: drop the ``kernel_call`` context manager; export
``wrap_kernel_exception(what, exc)`` — a pure function that maps
a raw exception to a PEP 249 one (KernelError → mapped class via
``reraise_kernel_error``; existing Error → passthrough; anything
else → OperationalError).
- ``client.py``: replace 12 ``with _kernel_call(...):`` blocks with
inline try/except calling the helper.
- ``result_set.py``: same for the 3 sites (arrow_schema on
construct, fetch_next_batch in _pull_one_batch, fetch_next_batch
in _drain).
Behaviour is unchanged — same KernelError → PEP 249 mapping, same
non-KernelError → OperationalError wrapping. Just spelled in a way
that makes control flow visible at the call site and keeps
tracebacks one frame shorter (no ``__exit__`` frame).
87/87 kernel unit tests still pass.
Co-authored-by: Isaac
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Fixes ``check-linting`` CI: one long line in ``execute_command``'s async-submit branch needed to wrap (the ``CommandId.from_sea_statement_id`` call). Pure formatting; no behaviour change. Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
|
P0 — Critical (must fix before merge)
P1 — Important (should fix)
|
Sorry, something went wrong.
Local kernel-package fixes from the follow-up review pass on PR #787 (#787 (comment)). The two cross-cutting / pre-existing issues (P0 #1 async leak in shared ``Cursor.close()``, P1 #8/#9 fetch-after-close not raising ``InterfaceError``) are tracked separately as #791 and #792 — they affect Thrift and SEA equally and are out of scope for this PR. P1 #2 — call ``backend.close_command`` from ``KernelResultSet.close()``: The override previously bypassed the base ``ResultSet.close()`` entirely. Honour the contract by invoking ``backend.close_command(self.command_id)`` after the per-handle close + ``_async_handles`` pop. Our own ``close_command`` is tolerant of already-popped guids (no-op), so this is safe even though the per-handle close above already released server-side state. Doesn't go through ``super().close()`` directly because the base path warns when ``self.results`` is ``None`` (which it is for kernel result sets) — replicate the meaningful part of the base contract without the noisy warning. P1 #3 — case-insensitive ``Bearer`` prefix in auth_bridge: RFC 6750 §2.1 says the Authorization scheme is case-insensitive. Match leniently in case a federation proxy or future provider normalises the casing differently — failing closed would surface as a confusing ``ProgrammingError`` from the bridge. P1 #4 — drop redundant ``__cause__`` set in ``reraise_kernel_error``: ``raise wrap_kernel_exception(...) from exc`` already sets ``__cause__`` at the call site; the manual assignment in ``reraise_kernel_error`` was redundant. Updated the test that asserted on it; added ``test_kernel_error_chains_through_wrap`` to cover the end-to-end chain. P1 #5 — ``get_tables`` filter looks up TABLE_TYPE by name: Replaced ``schema.field(5).name`` (positional) with the literal ``"TABLE_TYPE"`` plus a missing-column guard. A future kernel reshape of ``SHOW TABLES`` now surfaces an explicit ``OperationalError`` instead of silently filtering the wrong column. The case-sensitive contract is now documented in the surrounding comment (matches SEA + warehouse). P1 #6 — ``KernelResultSet.close()`` guards on ``connection.open``: ``__del__``-driven close arriving after the parent connection is already closed previously issued a kernel call into a disposed session. Skip the kernel call entirely in that case; still mark the result set ``CLOSED`` locally so ``__del__`` is idempotent. P1 #7 — defer ``kernel_auth_kwargs`` to ``open_session``: ``KernelDatabricksClient.__init__`` previously called ``kernel_auth_kwargs(auth_provider)`` and stored the bearer token on ``self._auth_kwargs`` indefinitely. If ``open_session`` never ran (test paths, error paths, lazy retries) the token stayed resident on the connector object. Build the kwargs locally inside ``open_session`` now — local variable, GC-eligible the moment ``open_session`` returns. Also tightened the install-hint comment in ``pyproject.toml`` to match the rest of the codebase (the wheel isn't on PyPI; only the ``maturin develop`` path is supported today). 88/88 kernel unit tests pass (added 1). Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
|
Thanks @gopalldb — addressed the 6 in-scope items in 0e1a250b; filed the 2 cross-cutting ones as separate issues since they're pre-existing on Thrift/SEA too. Fixed in this PR (0e1a250b):
88/88 kernel unit tests pass (added 1). Filed as separate issues — pre-existing on Thrift/SEA, out of scope for this PR:
Isolation sweep done: every kernel change is gated on use_kernel=True (default False); session.py only imports the kernel package inside the if self.use_kernel: branch; client.py changes are docstring-only; pyproject.toml is comment-only; conftest.py and existing parametrized e2e tests are untouched. Thrift and SEA flows see zero behaviour change vs. main. (Edited: review-finding numbers like P1 #N are now wrapped in code spans to prevent GitHub auto-linking them to unrelated PRs.) |
Sorry, something went wrong.
gopalldb
left a comment
There was a problem hiding this comment.
Some more minor feedbacks, can merge after addressing them
Sorry, something went wrong.
|
P1 findings
Each is ~5 lines; the gap is worth closing. P2 findings
|
Sorry, something went wrong.
Review comment: #787 (comment) P1.1 — get_query_state handles non-BaseException failure: Previously ``raise _reraise_kernel_error(failure) from failure`` would explode with ``TypeError: exception causes must derive from BaseException`` if the kernel's ``status()`` ever returned a ``failure`` that wasn't a real ``KernelError`` (struct, dict — kernel API drift). Now route through ``_wrap_kernel_exception``, which isinstance-checks and falls through to ``OperationalError`` for non-PEP-249-shaped values. New unit test ``test_get_query_state_handles_non_baseexception_failure``. P1.2 — regression tests for prior fixes: - ``test_bearer_prefix_is_case_insensitive`` (P1 #3 from the earlier review): parametrised over "Bearer "/"bearer "/ "BEARER "/"BeArEr ". RFC 6750 §2.1 compliance was added but not covered by a test. - ``test_close_skips_kernel_call_when_connection_already_closed`` (P1 #6 from the earlier review): exercises the ``connection.open is False`` branch in ``KernelResultSet.close()`` — asserts neither the kernel handle's close nor backend.close_command fire, but the result set still ends in ``CLOSED`` so ``__del__`` is idempotent. - ``test_token_with_control_chars_or_whitespace_rejected`` (pre-existing security guard): parametrised over NUL / CR / LF / DEL / space / tab — the regex previously missed space (0x20). Covered + extended. P2.1 — wrap session_id extraction in open_session: ``SessionId.from_sea_session_id(self._kernel_session.session_id)`` was outside the ``try/except _wrap_kernel_exception`` scope. A raw PyO3 attribute-conversion error on the ``self._kernel_session.session_id`` access could escape unwrapped. Now wrapped. P2.2 — drop redundant _async_handles.pop in result-set close: After the M1 fix (``get_execution_result`` pops the guid before constructing the result set), the pop in ``KernelResultSet.close`` is dead code — every call misses. Sync-execute and metadata paths never registered in ``_async_handles`` to begin with. Drop the per-close pop; rewrote the surrounding comment so ``backend.close_command`` is now the single bookkeeping seam. P2.3 — control-char regex includes whitespace: Extended ``[\x00-\x1f\x7f]`` → ``[\x00-\x20\x7f]`` and renamed ``_CONTROL_CHAR_RE`` → ``_TOKEN_REJECT_RE``. RFC 6750 forbids whitespace within the credential token itself; a token like ``"Bearer doubled-space-token"`` previously slipped past the injection guard. Test parametrised above. type_mapping reuse — use SqlType constants: Replaced literal type strings ("bigint", "string", …) in ``_arrow_type_to_dbapi_string`` with the ``SqlType`` constants from ``databricks.sql.backend.sea.utils.conversion`` — same single source of truth the SEA backend already uses, so the kernel and SEA backends emit byte-identical type-code strings. The Arrow → SqlType lookup itself stays local to the kernel (SEA receives type-text from the server and normalises it; the kernel receives Arrow schemas directly), but the names are now shared. 100/100 kernel unit tests pass (added 12). Co-authored-by: Isaac Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
|
Thanks @gopalldb — addressed all P1 and the in-scope P2 items in a05781f6. P1.1 (get_query_state non-BaseException failure) — Routed through _wrap_kernel_exception instead of _reraise_kernel_error + from failure. The wrap helper isinstance-checks failure; if the kernel ever surfaces a struct/dict instead of a real KernelError, callers now see OperationalError(repr(failure)) instead of TypeError: exception causes must derive from BaseException. New test test_get_query_state_handles_non_baseexception_failure covers it. P1.2 (missing regression tests) — Added all three:
P2.1 (SessionId.from_sea_session_id outside try/except) — Wrapped. A raw PyO3 attribute-conversion error on self._kernel_session.session_id now surfaces as a mapped PEP 249 exception. P2.2 (redundant _async_handles.pop in KernelResultSet.close()) — Dropped. After the M1 fix, get_execution_result pops the guid before constructing the result set, so the pop in close was dead code. Sync-execute and metadata paths never registered there to begin with. backend.close_command is now the single bookkeeping seam on close; rewrote the surrounding comment to match. P2.3 (whitespace token slipping the injection guard) — Extended [\x00-\x1f\x7f] → [\x00-\x20\x7f] and renamed _CONTROL_CHAR_RE → _TOKEN_REJECT_RE. RFC 6750 forbids whitespace within the credential token itself; tokens like "Bearer doubled-space-token" previously slipped past. type_mapping reuse (inline comment on type_mapping.py:28) — Replaced the literal type strings with the canonical SqlType constants from databricks.sql.backend.sea.utils.conversion. Same single source of truth the SEA backend already uses, so both backends emit byte-identical type-code strings. The Arrow → SqlType lookup itself stays local because SEA and kernel ingest types from different surfaces (SEA: server-emitted text; kernel: Arrow schema). Detailed reply on the inline thread. Deferred:
100/100 kernel unit tests pass (added 12 this round). |
Sorry, something went wrong.
Summary
Phase 2 of the PySQL × kernel integration plan (design doc). Adds a new opt-in use_kernel=True connection flag that routes through a new backend/kernel/ module delegating to the Rust kernel via the databricks_sql_kernel PyO3 extension (kernel PR #13).
This replaces the previous ADBC POC branches (backend/adbc/ and backend/adbc_dm/ on adbc-rust-backend-via-dm, which were never merged) with a clean port that uses the kernel's v0 Databricks-native API directly instead of layering through ADBC.
Flag semantics
use_kernel=True and use_sea=True are mutually exclusive — passing both raises ValueError. Existing use_sea=True callers are unaffected.
What use_kernel=True does
New module layout
Error mapping
KernelError.code → PEP 249 exception class, in a single table in client.py. Structured fields (sql_state, error_code, query_id, http_status, retryable, vendor_code) are copied onto the re-raised exception so callers can branch on err.code / err.sql_state directly. Live e2e verified: bad SQL on use_kernel=True surfaces as DatabaseError(code='SqlError', sql_state='42P01').
Packaging
Without the kernel wheel, use_kernel=True raises:
Local dev: cd databricks-sql-kernel/pyo3 && maturin develop --release into the connector's venv. (The [kernel] extra is intentionally not declared in pyproject.toml yet — databricks-sql-kernel isn't on PyPI, and declaring an unpublished dep breaks poetry lock for every CI job. The extra will land once the wheel is on PyPI.)
⚠️ Known gaps — acknowledged follow-ups
Code review feedback addressed in this revision
Multi-reviewer (architecture, security, ops, performance, test, maintainability, agent-compat, language, devil's advocate) review surfaced several issues; the highest-impact ones are addressed in commits 37fa5446 (mechanical) and 24e9a5c2 (substantive):
Test plan
This pull request and its description were written by Isaac.