Skip to content
Merged
Show file tree
Hide file tree
Changes from 10 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
38 changes: 38 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,38 @@ 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,
# Defaults to False by design: the kernel path deliberately diverges
# from the connector-wide True default (see session.py and client.py).
# Telemetry is off unless explicitly enabled on this backend.
"telemetry_enabled": bool(options.get("enable_telemetry", False)),
}
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 +259,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 +375,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 +419,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
8 changes: 6 additions & 2 deletions src/databricks/sql/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -402,8 +402,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
34 changes: 34 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,38 @@ 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.
# Intentionally defaults to False, diverging from the
# connector-wide True default on the Thrift/SEA path
# (client.py). The kernel path opts out of telemetry unless
# explicitly enabled; this is asserted by
# test_telemetry_enabled_defaults_false_for_kernel_client.
# Do not "fix" this back to True to match the other backends.
"enable_telemetry": kwargs.get("enable_telemetry", False),
Comment thread
jay-xiao446 marked this conversation as resolved.
Outdated
# 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
),
# Match the connector-wide True default (auth/common.py:55,
# where ClientContext defaults telemetry_circuit_breaker_enabled
# to True on the Thrift/SEA path). Forwarding an explicit default
# here keeps parity with telemetry_batch_size above, rather than
# passing None and letting the kernel silently pick its own
# internal default.
"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", True
),
}
return KernelDatabricksClient(
server_hostname=server_hostname,
http_path=http_path,
Expand All @@ -263,6 +296,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_false_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 False
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