Skip to content

Commit 8d6db43

Browse files
vuanhphungIsaac
andauthored
fix(kernel): forward telemetry_flush_interval_ms to the kernel (#1026)
* fix(kernel): forward telemetry_flush_interval_ms to the kernel The kernel binding accepts telemetry_flush_interval_ms, but the connector never forwarded it, so kernel connections always used the kernel's 5s default timer even when the caller asked for a longer interval. Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * fix(kernel): reject invalid telemetry_flush_interval_ms in connect() PyO3 accepts True as 1ms and reports negative or non-int values with a conversion error that does not name the option. Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> --------- Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com>
1 parent 307c48d commit 8d6db43

5 files changed

Lines changed: 59 additions & 4 deletions

File tree

‎CONNECTION_PARAMETERS.md‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -177,15 +177,17 @@ to change without notice.
177177

178178
Thrift telemetry is owned by the Python driver and normally gated by the
179179
`enableTelemetryForPythonDriver` server feature flag. Kernel telemetry is owned
180-
by the Rust kernel: explicit `enable_telemetry`, `telemetry_batch_size`, and
181-
`_telemetry_circuit_breaker_enabled` values are forwarded when supported by the
182-
installed kernel binding. `force_enable_telemetry` remains Python-driver-only.
180+
by the Rust kernel: explicit `enable_telemetry`, `telemetry_batch_size`,
181+
`telemetry_flush_interval_ms`, and `_telemetry_circuit_breaker_enabled` values
182+
are forwarded when supported by the installed kernel binding.
183+
`force_enable_telemetry` remains Python-driver-only.
183184

184185
| Option | Type | Thrift | Kernel | Default Value | Note |
185186
| ---------------------------------- | ------ | :----: | :----: | ------------- | -------------------------------------------------------- |
186187
| `enable_telemetry` | `bool` | ✅ | ✅ | `True` (Thrift); kernel-defined when unset | Thrift is also gated by the server flag. An explicit Kernel value is forwarded without that Python-side gate. |
187188
| `force_enable_telemetry` | `bool` | ✅ | ❌ | `False` | Bypass the server-side feature flag on Thrift. Not forwarded to Kernel. |
188189
| `telemetry_batch_size` | `int` | ✅ | ✅ | `100` | Events buffered before a flush. |
190+
| `telemetry_flush_interval_ms` | `int` (ms) | ❌ | ✅ | kernel-defined when unset | Kernel-only periodic flush interval; must be a positive `int`, else `connect()` raises `ValueError`. Thrift always flushes every 300s. |
189191
| `_telemetry_circuit_breaker_enabled`| `bool`| ✅ | ✅ | effective `False` (Thrift); kernel-defined when unset | The `ClientContext` signature defaults to `True`, but `connect()` passes `None` when this option is omitted and `bool(None)` makes the effective Thrift default `False`. An explicit value is forwarded to Kernel. |
190192

191193
> **Telemetry _events_ differ by backend.** Because the kernel owns result
@@ -222,6 +224,7 @@ installed kernel binding. `force_enable_telemetry` remains Python-driver-only.
222224
`token_url`.
223225
2. Custom `oauth_scopes`.
224226
3. Kernel U2M encrypted token storage (`oauth_token_cache_enabled`).
227+
4. Telemetry flush interval (`telemetry_flush_interval_ms`).
225228

226229
### Behavioral divergences to watch
227230

‎src/databricks/sql/backend/kernel/client.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,17 @@ def _kernel_geospatial_kwargs(value: bool) -> Dict[str, bool]:
200200
return {"enable_geospatial_support": value}
201201

202202

203+
def _validate_telemetry_flush_interval_ms(value: Any) -> None:
204+
# PyO3 would accept True as 1ms and report other bad values without
205+
# naming the option, so reject them before the kernel sees them.
206+
if value is not None and (
207+
isinstance(value, bool) or not isinstance(value, int) or value <= 0
208+
):
209+
raise ValueError(
210+
f"telemetry_flush_interval_ms must be a positive int; got {value!r}"
211+
)
212+
213+
203214
def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]:
204215
"""Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``.
205216
@@ -227,6 +238,10 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]:
227238
candidates["telemetry_enabled"] = bool(options["enable_telemetry"])
228239
if options.get("telemetry_batch_size") is not None:
229240
candidates["telemetry_batch_size"] = options["telemetry_batch_size"]
241+
if options.get("telemetry_flush_interval_ms") is not None:
242+
candidates["telemetry_flush_interval_ms"] = options[
243+
"telemetry_flush_interval_ms"
244+
]
230245
if options.get("telemetry_circuit_breaker_enabled") is not None:
231246
candidates["telemetry_circuit_breaker_enabled"] = options[
232247
"telemetry_circuit_breaker_enabled"
@@ -302,6 +317,9 @@ def __init__(
302317
# Kernel telemetry phase 7 adds binding/runtime identity and
303318
# telemetry config kwargs directly to ``databricks_sql_kernel.Session``.
304319
self._telemetry_options = kwargs.get("telemetry_options") or {}
320+
_validate_telemetry_flush_interval_ms(
321+
self._telemetry_options.get("telemetry_flush_interval_ms")
322+
)
305323
self._catalog = catalog
306324
self._schema = schema
307325
self._escape_catalog_name = kwargs.get("escape_catalog_name", False)

‎src/databricks/sql/session.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,11 @@ def _create_backend(
299299
"telemetry_batch_size": kwargs.get(
300300
"telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE
301301
),
302+
# Kernel-only: Thrift's exporter flushes on a fixed interval.
303+
# When unset the kernel keeps its own default.
304+
"telemetry_flush_interval_ms": kwargs.get(
305+
"telemetry_flush_interval_ms"
306+
),
302307
# Preserve the caller's explicit circuit-breaker choice. When
303308
# unset, leave it as None so the kernel owns the decision --
304309
# mirroring the enable_telemetry handling above rather than

‎tests/unit/test_kernel_client.py‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,6 +518,7 @@ def fake_session(**kw):
518518
telemetry_options={
519519
"enable_telemetry": True,
520520
"telemetry_batch_size": 17,
521+
"telemetry_flush_interval_ms": 300_000,
521522
"telemetry_circuit_breaker_enabled": False,
522523
},
523524
)
@@ -537,9 +538,30 @@ def fake_session(**kw):
537538
assert captured["process_name"] is None
538539
assert captured["telemetry_enabled"] is True
539540
assert captured["telemetry_batch_size"] == 17
541+
assert captured["telemetry_flush_interval_ms"] == 300_000
540542
assert captured["telemetry_circuit_breaker_enabled"] is False
541543

542544

545+
@pytest.mark.parametrize("value", [0, -1, "300000", 1.5, True])
546+
def test_invalid_telemetry_flush_interval_is_rejected(value):
547+
with pytest.raises(ValueError, match="telemetry_flush_interval_ms"):
548+
kernel_client.KernelDatabricksClient(
549+
server_hostname="example.cloud.databricks.com",
550+
http_path="/sql/1.0/warehouses/abc",
551+
auth_provider=AccessTokenAuthProvider("dapi-test"),
552+
ssl_options=None,
553+
telemetry_options={"telemetry_flush_interval_ms": value},
554+
)
555+
556+
557+
def test_unset_telemetry_flush_interval_is_not_forwarded(monkeypatch):
558+
monkeypatch.setattr(kernel_client, "_kernel_session_accepts_kwarg", lambda _: True)
559+
kwargs = kernel_client._kernel_telemetry_kwargs(
560+
{"enable_telemetry": True, "telemetry_flush_interval_ms": None}
561+
)
562+
assert "telemetry_flush_interval_ms" not in kwargs
563+
564+
543565
def test_open_session_omits_optional_kwargs_kernel_does_not_accept(monkeypatch):
544566
"""Optional kwargs must NOT be forwarded to a kernel ``Session`` whose
545567
fixed constructor doesn't declare them.
@@ -612,7 +634,11 @@ def fake_session_without_optional_kwargs(
612634
# The kwargs the client builds must be filtered to what fake_session
613635
# accepts, so open_session succeeds instead of raising TypeError.
614636
kwargs = kernel_client._kernel_telemetry_kwargs(
615-
{"enable_telemetry": True, "telemetry_batch_size": 17}
637+
{
638+
"enable_telemetry": True,
639+
"telemetry_batch_size": 17,
640+
"telemetry_flush_interval_ms": 300_000,
641+
}
616642
)
617643
assert kwargs == {}, f"expected no unsupported phase-7 kwargs, got {kwargs}"
618644

‎tests/unit/test_session.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -697,13 +697,15 @@ def test_telemetry_kwargs_threaded_into_kernel_client(self):
697697
enable_telemetry=True,
698698
force_enable_telemetry=False,
699699
telemetry_batch_size=17,
700+
telemetry_flush_interval_ms=300_000,
700701
_telemetry_circuit_breaker_enabled=False,
701702
)
702703
try:
703704
_, kwargs = mock_kernel_client.call_args
704705
opts = kwargs["telemetry_options"]
705706
assert opts["enable_telemetry"] is True
706707
assert opts["telemetry_batch_size"] == 17
708+
assert opts["telemetry_flush_interval_ms"] == 300_000
707709
assert opts["telemetry_circuit_breaker_enabled"] is False
708710
finally:
709711
conn.close()
@@ -741,6 +743,7 @@ def test_telemetry_enabled_defaults_none_for_kernel_client(self):
741743
_, kwargs = mock_kernel_client.call_args
742744
opts = kwargs["telemetry_options"]
743745
assert opts["enable_telemetry"] is None
746+
assert opts["telemetry_flush_interval_ms"] is None
744747
finally:
745748
conn.close()
746749

0 commit comments

Comments
 (0)