Skip to content
Merged
Show file tree
Hide file tree
Changes from 14 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions src/databricks/sql/backend/kernel/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
NotSupportedError,
ProgrammingError,
)
from databricks.sql.telemetry.telemetry_client import TelemetryHelper

if TYPE_CHECKING:
from databricks.sql.client import Cursor
Expand Down Expand Up @@ -172,6 +173,36 @@ def _is_staging_statement(operation: str) -> bool:
return verb in _STAGING_VERBS


def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]:
Comment thread
jay-xiao446 marked this conversation as resolved.
"""Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``."""
system = TelemetryHelper.get_driver_system_configuration()
out: Dict[str, Any] = {
"driver_name": system.driver_name,
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
"driver_version": system.driver_version,
"runtime_name": system.runtime_name,
"runtime_version": system.runtime_version,
"runtime_vendor": system.runtime_vendor,
"os_name": system.os_name,
"os_version": system.os_version,
"os_arch": system.os_arch,
"client_app_name": system.client_app_name,
"locale_name": system.locale_name,
"char_set_encoding": system.char_set_encoding,
# The Python telemetry model does not currently track process
# name; omit it and let the kernel fill what it can derive.
"process_name": None,
}
if options.get("enable_telemetry") is not None:
out["telemetry_enabled"] = bool(options["enable_telemetry"])
if options.get("telemetry_batch_size") is not None:
out["telemetry_batch_size"] = options["telemetry_batch_size"]
if options.get("telemetry_circuit_breaker_enabled") is not None:
out["telemetry_circuit_breaker_enabled"] = options[
"telemetry_circuit_breaker_enabled"
]
return out


# ─── Client ─────────────────────────────────────────────────────────────────


Expand Down Expand Up @@ -226,6 +257,9 @@ def __init__(
self._retry_options = kwargs.get("retry_options") or {}
# The kernel binding owns type and range validation.
self._request_timeout_secs = kwargs.get("request_timeout_secs")
# Kernel telemetry phase 7 adds binding/runtime identity and
# telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
self._telemetry_options = kwargs.get("telemetry_options") or {}
self._catalog = catalog
self._schema = schema
# ``_use_arrow_native_complex_types`` is the connector-side
Expand Down Expand Up @@ -339,6 +373,7 @@ def open_session(
# Translate the connector's ``_retry_*`` kwargs into the
# kernel's ``retry_*`` kwargs. Empty when at defaults.
retry_kwargs = _kernel_retry_kwargs(self._retry_options)
telemetry_kwargs = _kernel_telemetry_kwargs(self._telemetry_options)
# Forward caller / connector HTTP headers. The kernel applies
# them on every request; a caller ``User-Agent`` is appended
# to the kernel's base UA. Only pass the kwarg when there's
Expand Down Expand Up @@ -382,6 +417,7 @@ def open_session(
**auth_kwargs,
**tls_kwargs,
**retry_kwargs,
**telemetry_kwargs,
Comment thread
jay-xiao446 marked this conversation as resolved.
**http_headers_kwargs,
)
except Exception as exc:
Expand Down
18 changes: 16 additions & 2 deletions src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,16 @@ def __init__(
``ImportError``. Supports PAT, OAuth M2M, and OAuth
U2M auth, and native (positional and named) parameter
binding. Mutually exclusive with ``use_sea``.

Telemetry note: on the kernel path telemetry is
strictly opt-in. An explicit ``enable_telemetry=True``
turns telemetry on unconditionally; it is NOT gated by
the server-side ``enableTelemetryForPythonDriver``
feature flag that applies on the Thrift/SEA paths. When
``enable_telemetry`` is unset the kernel owns the enable
decision. This is an intentional divergence from the
Thrift/SEA paths, where an explicit ``True`` can still
be suppressed by the feature flag.
:param use_hybrid_disposition: `bool`, optional (default is False)
Use the hybrid disposition instead of the inline disposition.
:param server_hostname: Databricks instance host name.
Expand Down Expand Up @@ -402,8 +412,12 @@ def read(self) -> Optional[OAuthToken]:
)
self.session.open()
except Exception as e:
# Respect user's telemetry preference even during connection failure
enable_telemetry = kwargs.get("enable_telemetry", True)
# Respect user's telemetry preference even during connection failure.
# For use_kernel connections the kernel owns telemetry, so suppress
# the wrapper-side failure log to avoid wrapper-vs-kernel duplication.
enable_telemetry = kwargs.get("enable_telemetry", True) and not kwargs.get(
Comment thread
jay-xiao446 marked this conversation as resolved.
"use_kernel", False
)
TelemetryClientFactory.connection_failure_log(
error_name="Exception",
error_message=str(e),
Expand Down
51 changes: 51 additions & 0 deletions src/databricks/sql/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
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
from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory

if TYPE_CHECKING:
from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
Expand Down Expand Up @@ -250,6 +251,55 @@ def _create_backend(
"_retry_stop_after_attempts_duration"
),
}
# Forward the binding/runtime identity and telemetry knobs
# added by kernel telemetry phase 7. Python-side telemetry
# still owns feature-flag evaluation and event export for the
# Thrift/SEA paths; the kernel path needs the same driver
# identity at Session construction time so kernel-owned
# telemetry can populate its system configuration.
kernel_telemetry_options = {
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
# Preserve the caller's explicit telemetry choice. When unset,
# leave it as None so the kernel owns the enable decision --
# this is an INTENTIONAL divergence from the Thrift/SEA path,
# not an oversight. There, an unset enable_telemetry defaults to
# True (client.py) but only actually emits when the
# `enableTelemetryForPythonDriver` server feature flag is on
# (telemetry_client.py). That feature-flag gate is Python-side
# and is bypassed on the kernel path (is_telemetry_enabled
# short-circuits to False for use_kernel), so forwarding the
# connector's True default here would force telemetry on without
# an equivalent gate. Passing None instead defers to the
# kernel's own default/gating, which is expected to mirror the
# feature-flag-gated wrapper behaviour; only an explicit caller
# opt-in/opt-out overrides it.
"enable_telemetry": kwargs.get("enable_telemetry"),
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
# Match the connector's default batch size (client.py forwards
# the same TelemetryClientFactory.DEFAULT_BATCH_SIZE fallback)
# so an unset telemetry_batch_size resolves to the same value
# on the kernel path as on the Thrift/SEA path, rather than
# letting the kernel silently pick its own internal default.
"telemetry_batch_size": kwargs.get(
"telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE
),
# Preserve the caller's explicit circuit-breaker choice. When
# unset, leave it as None so the kernel owns the decision --
# mirroring the enable_telemetry handling above rather than
# forcing a default. This is deliberately NOT defaulted to True:
# although ClientContext's signature default is True
# (auth/common.py:55), the Thrift/SEA path never reaches it --
# build_client_context (utils.py:1018) always passes
# _telemetry_circuit_breaker_enabled explicitly (None when the
# caller leaves it unset), and ClientContext coerces it with
# bool(None) -> False (auth/common.py:89). So the *effective*
# Thrift/SEA default when unset is False, not True; forwarding
# True here would turn the circuit breaker on for an
# unconfigured connection while Thrift/SEA leaves it off.
# Passing None instead defers to the kernel's own default;
# only an explicit caller value overrides it.
"telemetry_circuit_breaker_enabled": kwargs.get(
Comment thread
jay-xiao446 marked this conversation as resolved.
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
"_telemetry_circuit_breaker_enabled"
),
}
return KernelDatabricksClient(
server_hostname=server_hostname,
http_path=http_path,
Expand All @@ -263,6 +313,7 @@ def _create_backend(
auth_options=kernel_auth_options,
retry_options=kernel_retry_options,
request_timeout_secs=kwargs.get("_socket_timeout"),
telemetry_options=kernel_telemetry_options,
)

# These reference the lazily-resolved module attributes defined via
Expand Down
3 changes: 3 additions & 0 deletions src/databricks/sql/telemetry/telemetry_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,9 @@ def get_auth_flow(auth_provider):

@staticmethod
def is_telemetry_enabled(connection: "Connection") -> bool:
if getattr(connection.session, "use_kernel", False):
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
return False

# Fast path: force enabled - skip feature flag fetch entirely
if connection.force_enable_telemetry:
return True
Expand Down
60 changes: 60 additions & 0 deletions tests/unit/test_kernel_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -368,6 +368,66 @@ def fake_session(**kw):
assert captured["request_timeout_secs"] == timeout


def test_open_session_passes_phase_7_telemetry_kwargs_to_kernel(monkeypatch):
"""Kernel telemetry phase 7 added binding/runtime identity and
telemetry config kwargs to ``databricks_sql_kernel.Session``."""
captured = {}

def fake_session(**kw):
captured.update(kw)
sess = MagicMock()
sess.session_id = "sess-id"
return sess

monkeypatch.setattr(kernel_client._kernel, "Session", fake_session)
monkeypatch.setattr(
kernel_client.TelemetryHelper,
"get_driver_system_configuration",
lambda: types.SimpleNamespace(
driver_name="Databricks SQL Python Connector",
driver_version="1.2.3",
runtime_name="Python 3.12.0",
runtime_version="3.12.0",
runtime_vendor="CPython",
os_name="Linux",
os_version="6.1",
os_arch="x86_64",
client_app_name=None,
locale_name="en_US",
char_set_encoding="utf-8",
),
)

c = kernel_client.KernelDatabricksClient(
server_hostname="example.cloud.databricks.com",
http_path="/sql/1.0/warehouses/abc",
auth_provider=AccessTokenAuthProvider("dapi-test"),
ssl_options=None,
telemetry_options={
"enable_telemetry": True,
"telemetry_batch_size": 17,
"telemetry_circuit_breaker_enabled": False,
},
)
c.open_session(session_configuration=None, catalog=None, schema=None)

assert captured["driver_name"] == "Databricks SQL Python Connector"
assert captured["driver_version"] == "1.2.3"
assert captured["runtime_name"] == "Python 3.12.0"
assert captured["runtime_version"] == "3.12.0"
assert captured["runtime_vendor"] == "CPython"
assert captured["os_name"] == "Linux"
assert captured["os_version"] == "6.1"
assert captured["os_arch"] == "x86_64"
assert captured["client_app_name"] is None
assert captured["locale_name"] == "en_US"
assert captured["char_set_encoding"] == "utf-8"
assert captured["process_name"] is None
assert captured["telemetry_enabled"] is True
assert captured["telemetry_batch_size"] == 17
assert captured["telemetry_circuit_breaker_enabled"] is False


def test_execute_command_forwards_parameters_to_bind_param():
"""``execute_command(parameters=[...])`` routes each parameter
through ``bind_tspark_params`` onto the kernel statement before
Expand Down
86 changes: 86 additions & 0 deletions tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -570,6 +570,92 @@ def test_oauth_token_cache_enabled_threaded_into_kernel_auth_options(self):
conn.close()


class TestKernelTelemetryOptionsThreading:
"""The kernel path must forward telemetry options from connect()
into ``KernelDatabricksClient`` so phase-7 PyO3 Session kwargs can
be populated before the kernel opens its session."""

PACKAGE = "databricks.sql"

def test_telemetry_kwargs_threaded_into_kernel_client(self):
import sys
import types

pytest.importorskip(
"pyarrow",
reason="kernel client module imports pyarrow at load",
)

fake = types.ModuleType("databricks_sql_kernel")
fake.KernelError = type("KernelError", (Exception,), {})
fake.Session = MagicMock()

with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch(
"databricks.sql.backend.kernel.client.KernelDatabricksClient"
) as mock_kernel_client, patch(
"%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE
):
instance = mock_kernel_client.return_value
instance.open_session.return_value = SessionId(
BackendType.SEA, "sess-id", None
)

conn = databricks.sql.connect(
server_hostname="foo",
http_path="/sql/1.0/warehouses/abc",
use_kernel=True,
access_token="dapi-xyz",
enable_telemetry=True,
force_enable_telemetry=False,
telemetry_batch_size=17,
_telemetry_circuit_breaker_enabled=False,
)
try:
_, kwargs = mock_kernel_client.call_args
opts = kwargs["telemetry_options"]
assert opts["enable_telemetry"] is True
assert opts["telemetry_batch_size"] == 17
assert opts["telemetry_circuit_breaker_enabled"] is False
finally:
conn.close()

def test_telemetry_enabled_defaults_none_for_kernel_client(self):
import sys
import types

pytest.importorskip(
"pyarrow",
reason="kernel client module imports pyarrow at load",
)

fake = types.ModuleType("databricks_sql_kernel")
fake.KernelError = type("KernelError", (Exception,), {})
fake.Session = MagicMock()

with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch(
"databricks.sql.backend.kernel.client.KernelDatabricksClient"
) as mock_kernel_client, patch(
"%s.session.get_python_sql_connector_auth_provider" % self.PACKAGE
):
instance = mock_kernel_client.return_value
instance.open_session.return_value = SessionId(
BackendType.SEA, "sess-id", None
)

conn = databricks.sql.connect(
server_hostname="foo",
http_path="/sql/1.0/warehouses/abc",
use_kernel=True,
access_token="dapi-xyz",
)
try:
_, kwargs = mock_kernel_client.call_args
opts = kwargs["telemetry_options"]
assert opts["enable_telemetry"] is None
finally:
conn.close()


class TestKernelUserAgentForwarding:
"""user_agent_entry must reach the kernel on the use_kernel path —
session.py folds it into the composed User-Agent and includes it in
Expand Down
Loading
Loading