diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 37736ed1a..063176cc3 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -177,15 +177,17 @@ to change without notice. Thrift telemetry is owned by the Python driver and normally gated by the `enableTelemetryForPythonDriver` server feature flag. Kernel telemetry is owned -by the Rust kernel: explicit `enable_telemetry`, `telemetry_batch_size`, and -`_telemetry_circuit_breaker_enabled` values are forwarded when supported by the -installed kernel binding. `force_enable_telemetry` remains Python-driver-only. +by the Rust kernel: explicit `enable_telemetry`, `telemetry_batch_size`, +`telemetry_flush_interval_ms`, and `_telemetry_circuit_breaker_enabled` values +are forwarded when supported by the installed kernel binding. +`force_enable_telemetry` remains Python-driver-only. | Option | Type | Thrift | Kernel | Default Value | Note | | ---------------------------------- | ------ | :----: | :----: | ------------- | -------------------------------------------------------- | | `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. | | `force_enable_telemetry` | `bool` | ✅ | ❌ | `False` | Bypass the server-side feature flag on Thrift. Not forwarded to Kernel. | | `telemetry_batch_size` | `int` | ✅ | ✅ | `100` | Events buffered before a flush. | +| `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. | | `_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. | > **Telemetry _events_ differ by backend.** Because the kernel owns result @@ -222,6 +224,7 @@ installed kernel binding. `force_enable_telemetry` remains Python-driver-only. `token_url`. 2. Custom `oauth_scopes`. 3. Kernel U2M encrypted token storage (`oauth_token_cache_enabled`). +4. Telemetry flush interval (`telemetry_flush_interval_ms`). ### Behavioral divergences to watch diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0ca7ba6a8..01889e03a 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -200,6 +200,17 @@ def _kernel_geospatial_kwargs(value: bool) -> Dict[str, bool]: return {"enable_geospatial_support": value} +def _validate_telemetry_flush_interval_ms(value: Any) -> None: + # PyO3 would accept True as 1ms and report other bad values without + # naming the option, so reject them before the kernel sees them. + if value is not None and ( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + ): + raise ValueError( + f"telemetry_flush_interval_ms must be a positive int; got {value!r}" + ) + + def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: """Build phase-7 telemetry/system kwargs for ``databricks_sql_kernel.Session``. @@ -227,6 +238,10 @@ def _kernel_telemetry_kwargs(options: Dict[str, Any]) -> Dict[str, Any]: candidates["telemetry_enabled"] = bool(options["enable_telemetry"]) if options.get("telemetry_batch_size") is not None: candidates["telemetry_batch_size"] = options["telemetry_batch_size"] + if options.get("telemetry_flush_interval_ms") is not None: + candidates["telemetry_flush_interval_ms"] = options[ + "telemetry_flush_interval_ms" + ] if options.get("telemetry_circuit_breaker_enabled") is not None: candidates["telemetry_circuit_breaker_enabled"] = options[ "telemetry_circuit_breaker_enabled" @@ -302,6 +317,9 @@ def __init__( # 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 {} + _validate_telemetry_flush_interval_ms( + self._telemetry_options.get("telemetry_flush_interval_ms") + ) self._catalog = catalog self._schema = schema self._escape_catalog_name = kwargs.get("escape_catalog_name", False) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index e2b27a980..04d603c21 100644 --- a/src/databricks/sql/session.py +++ b/src/databricks/sql/session.py @@ -299,6 +299,11 @@ def _create_backend( "telemetry_batch_size": kwargs.get( "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE ), + # Kernel-only: Thrift's exporter flushes on a fixed interval. + # When unset the kernel keeps its own default. + "telemetry_flush_interval_ms": kwargs.get( + "telemetry_flush_interval_ms" + ), # 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 diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index b285ce6e1..e0b8d99e0 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -518,6 +518,7 @@ def fake_session(**kw): telemetry_options={ "enable_telemetry": True, "telemetry_batch_size": 17, + "telemetry_flush_interval_ms": 300_000, "telemetry_circuit_breaker_enabled": False, }, ) @@ -537,9 +538,30 @@ def fake_session(**kw): assert captured["process_name"] is None assert captured["telemetry_enabled"] is True assert captured["telemetry_batch_size"] == 17 + assert captured["telemetry_flush_interval_ms"] == 300_000 assert captured["telemetry_circuit_breaker_enabled"] is False +@pytest.mark.parametrize("value", [0, -1, "300000", 1.5, True]) +def test_invalid_telemetry_flush_interval_is_rejected(value): + with pytest.raises(ValueError, match="telemetry_flush_interval_ms"): + 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={"telemetry_flush_interval_ms": value}, + ) + + +def test_unset_telemetry_flush_interval_is_not_forwarded(monkeypatch): + monkeypatch.setattr(kernel_client, "_kernel_session_accepts_kwarg", lambda _: True) + kwargs = kernel_client._kernel_telemetry_kwargs( + {"enable_telemetry": True, "telemetry_flush_interval_ms": None} + ) + assert "telemetry_flush_interval_ms" not in kwargs + + def test_open_session_omits_optional_kwargs_kernel_does_not_accept(monkeypatch): """Optional kwargs must NOT be forwarded to a kernel ``Session`` whose fixed constructor doesn't declare them. @@ -612,7 +634,11 @@ def fake_session_without_optional_kwargs( # The kwargs the client builds must be filtered to what fake_session # accepts, so open_session succeeds instead of raising TypeError. kwargs = kernel_client._kernel_telemetry_kwargs( - {"enable_telemetry": True, "telemetry_batch_size": 17} + { + "enable_telemetry": True, + "telemetry_batch_size": 17, + "telemetry_flush_interval_ms": 300_000, + } ) assert kwargs == {}, f"expected no unsupported phase-7 kwargs, got {kwargs}" diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 2ca944dc5..a2ab2479e 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -697,6 +697,7 @@ def test_telemetry_kwargs_threaded_into_kernel_client(self): enable_telemetry=True, force_enable_telemetry=False, telemetry_batch_size=17, + telemetry_flush_interval_ms=300_000, _telemetry_circuit_breaker_enabled=False, ) try: @@ -704,6 +705,7 @@ def test_telemetry_kwargs_threaded_into_kernel_client(self): opts = kwargs["telemetry_options"] assert opts["enable_telemetry"] is True assert opts["telemetry_batch_size"] == 17 + assert opts["telemetry_flush_interval_ms"] == 300_000 assert opts["telemetry_circuit_breaker_enabled"] is False finally: conn.close() @@ -741,6 +743,7 @@ def test_telemetry_enabled_defaults_none_for_kernel_client(self): _, kwargs = mock_kernel_client.call_args opts = kwargs["telemetry_options"] assert opts["enable_telemetry"] is None + assert opts["telemetry_flush_interval_ms"] is None finally: conn.close()