feat(kernel): honor _connection_uri and _port on the use_kernel path - #915
feat(kernel): honor _connection_uri and _port on the use_kernel path#915eric-wang-1990 wants to merge 10 commits into
Conversation
Sorry, something went wrong.
The kernel branch of Session._create_backend forwarded only server_hostname and http_path, so _connection_uri and _port were silently ignored on use_kernel=True (connection reached server_hostname/http_path with no error). Add _kernel_host_and_path(): decompose _connection_uri into the kernel host (scheme+authority) + http_path, and fold _port into the host authority. No kernel change needed — the kernel Session host accepts a fully-qualified https://host:port and its normalise_host preserves scheme and port. PECOBLR-4151. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
There was a problem hiding this comment.
Pull request overview
This PR fixes use_kernel=True connection override handling so _connection_uri and _port are honored when constructing the kernel backend endpoint, preventing silent connections to an unintended host/path.
Changes:
- Added _kernel_host_and_path() to resolve (host, http_path) for the kernel backend, honoring _connection_uri and _port.
- Updated Session._create_backend to pass the resolved host/path to KernelDatabricksClient.
- Added unit tests covering _connection_uri splitting, query preservation, precedence over _port, and port folding.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| tests/unit/test_session.py | Adds unit tests for kernel host/http_path override resolution. |
| src/databricks/sql/session.py | Implements _kernel_host_and_path() and wires it into the kernel backend creation path. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Sorry, something went wrong.
| # Ensure a scheme so urlsplit populates netloc rather than path; the | ||
| # Thrift backend defaults a scheme-less URI to https, so do the same. |
There was a problem hiding this comment.
Reworded the misleading inline comment in _kernel_host_and_path (session.py:47-49). The reviewer correctly noted that the Thrift backend passes _connection_uri through unmodified (uri = kwargs.get("_connection_uri")) rather than defaulting a scheme-less URI to https — that defaulting only applies to the server_hostname/http_path-built URI in the elif branch. The comment now justifies the https default on connector-transport-default grounds instead of a nonexistent Thrift behavior. Comment-only edit; no logic or tests affected.
Pushed a5e0a08 (bundled with 4 other thread(s)).
Sorry, something went wrong.
| parts = urlsplit(uri) | ||
| host = "{}://{}".format(parts.scheme, parts.netloc) | ||
| path = parts.path or http_path |
There was a problem hiding this comment.
Resolved with a code change. _kernel_host_and_path now raises a clear ValueError when _connection_uri has no host authority, and I added a test covering the malformed cases the reviewer flagged.
Pushed a5e0a08 (bundled with 4 other thread(s)).
Sorry, something went wrong.
| # Split off any scheme so we can inspect the authority; the kernel | ||
| # re-adds https:// when it is absent. Only append the port when the | ||
| # authority does not already carry one. | ||
| scheme_match = re.match(r"^(https?://)(.*)$", server_hostname) | ||
| scheme = scheme_match.group(1) if scheme_match else "" | ||
| authority = (scheme_match.group(2) if scheme_match else server_hostname).rstrip( | ||
| "/" | ||
| ) | ||
| if ":" not in authority: | ||
| authority = "{}:{}".format(authority, port) | ||
| return "{}{}".format(scheme, authority), http_path |
There was a problem hiding this comment.
Fixed as the reviewer suggested. Summary:
- src/databricks/sql/session.py: In the _port branch of _kernel_host_and_path, swapped the ":" not in host heuristic for urlsplit(probe).port is None. IPv6 authorities contain : even with no port (e.g. [::1]), so the old check wrongly skipped appending _port. Since urlsplit only populates port when a scheme is present, a temporary https:// is prepended for scheme-less hosts before probing.
- tests/unit/test_session.py: Added two regression tests — [::1] gets :8443 appended, and [::1]:7000 is left untouched.
Tests pass (46 passed, realkernel deselected).
Pushed a5e0a08 (bundled with 4 other thread(s)).
Sorry, something went wrong.
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the new _kernel_host_and_path faithfully mirrors the Thrift backend's _connection_uri/_port handling (thrift_backend.py:155-162), is a pure function with solid unit coverage, and leaves the common (no-override) path untouched. One low-severity edge case: SPOG org-id headers are derived from the original http_path, not the _connection_uri-resolved one.
Sorry, something went wrong.
server_hostname reaches the backend as a bare host on this path, so drop the defensive scheme peel/re-add: just append the port when the host has none, and let the kernel's normalise_host add the scheme. Removes the now-moot scheme-preservation test. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the connector-side remap in _kernel_host_and_path correctly mirrors the Thrift backend's _connection_uri/_port precedence (verified against thrift_backend.py:154-162), and the _port is not None check (no 443 default) correctly folds only an explicit port. One low-severity coverage gap: the pure-function tests don't guard the _create_backend → KernelDatabricksClient wiring that this PR actually fixes.
Sorry, something went wrong.
Addresses: - #3800053148 at src/databricks/sql/session.py:48 - #3800053178 at src/databricks/sql/session.py:52 - #3800053209 at src/databricks/sql/session.py:69 - #3800056285 at src/databricks/sql/session.py:245 - #3800209433 at src/databricks/sql/session.py:245 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Medium
Solid, well-tested change that correctly maps _connection_uri/_port onto the kernel host/path. One medium gap: the SPOG-header re-derivation guard (and self._spog_headers) is asymmetric — it won't inject x-databricks-org-id when the original path lacked workspace info but a _connection_uri override introduces one, leaving that direction of mis-routing unfixed and untested.
Sorry, something went wrong.
There was a problem hiding this comment.
Verdict: 1 Low
Solid, well-tested change — the _connection_uri/_port resolution, IPv6-aware port detection, and bidirectional SPOG org-id re-derivation are all correct and covered by pure-function and integration tests. One low-severity latent crash: a path-less, query-bearing _connection_uri combined with http_path=None dereferences None.split(...).
Sorry, something went wrong.
Addresses: - #3802164062 at src/databricks/sql/session.py:71 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — solid, thoroughly-tested change; one low-severity type-hint accuracy note. The _connection_uri/_port resolution correctly mirrors Thrift precedence, the SPOG org-id header is re-derived in both directions when the resolved path changes, caller-set headers are preserved, and edge cases (IPv6 literals, double-query, scheme defaulting, missing authority) are covered by pure-function and integration tests.
Sorry, something went wrong.
Addresses: - #3802212055 at src/databricks/sql/session.py:24 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a focused, well-tested fix that honors _connection_uri/_port on the kernel path. The _kernel_host_and_path URI split, port folding (incl. IPv6 handling), and the SPOG org-id header re-derivation from the resolved path (both directions, caller-header precedence preserved) are all correct and covered by pure-function + integration tests. One low-severity robustness note filed inline about an opaque ValueError on malformed server_hostname ports.
Sorry, something went wrong.
Addresses: - #3802255225 at src/databricks/sql/session.py:93 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good overall — the _connection_uri/_port resolution and the SPOG org-id re-derivation are carefully reasoned and thoroughly unit-tested (pure-function + kernel-client-threading + header-re-derivation cases). One low-severity edge case: a _connection_uri with a bare trailing slash silently discards the original warehouse path.
Sorry, something went wrong.
Addresses: - #3802314928 at src/databricks/sql/session.py:71 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
What
On use_kernel=True, Session._create_backend forwarded only server_hostname and http_path to KernelDatabricksClient, so _connection_uri and _port were silently ignored — a connection using either reached server_hostname/http_path instead, with no error (a successful connection to the wrong endpoint).
This adds _kernel_host_and_path(), which resolves the (host, http_path) the kernel Session should use:
Why no kernel change is needed
Verified against the kernel repo: the kernel Session host accepts a fully-qualified URL, and normalise_host (src/config.rs) only prepends https:// when the scheme is absent and trims a trailing slash — it preserves the scheme and never strips the port. So both overrides are expressible connector-side. (Kernel-genuine transport knobs — socket timeout, connection pool — remain tracked in PECOBLR-4150.)
Tests
TestKernelHostAndPathOverrides in tests/unit/test_session.py — pure-function coverage (no kernel wheel needed) for: URI split (authority/path/query), scheme defaulting, _connection_uri winning over _port, port folding, scheme preservation, and no-double-port.
Jira: PECOBLR-4151. Reference doc rows for _port / _connection_uri in #913 should flip to supported-on-kernel once this merges.
This pull request and its description were written by Isaac.
This PR was created with GitHub MCP.