Forward kernel telemetry circuit breaker config

Author: jay-xiao446

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

@@ -43,6 +43,7 @@
     SessionId,
 )
 from databricks.sql.exc import (
+    Error,
     InterfaceError,
     NotSupportedError,
     ProgrammingError,
@@ -71,6 +72,8 @@
 # per-request skip-and-warn.
 _KERNEL_MANAGED_HEADERS = frozenset({"authorization", "x-databricks-org-id"})
 
+TIMESTAMP_AS_STRING_CONFIG = "spark.thriftserver.arrowBasedRowSet.timestampAsString"
+
 # Leading verbs of SQL volume/staging statements. Detected by the
 # leading token (case-insensitive) so the kernel backend can fail loud
 # on staging ops it can't service — see ``execute_command``.
@@ -107,6 +110,18 @@ def _strip_leading_sql_comments(sql: str) -> str:
     return sql[i:]
 
 
+def _check_session_configuration(session_configuration: Dict[str, str]) -> None:
+    # This client expects timestampAsString to be false, so do not allow overrides.
+    if session_configuration.get(TIMESTAMP_AS_STRING_CONFIG, "false").lower() != "false":
+        raise Error(
+            "Invalid session configuration: {} cannot be changed "
+            "while using the Databricks SQL connector, it must be false not {}".format(
+                TIMESTAMP_AS_STRING_CONFIG,
+                session_configuration[TIMESTAMP_AS_STRING_CONFIG],
+            )
+        )
+
+
 def _is_not_found(exc: BaseException) -> bool:
     """True iff ``exc`` is a kernel ``NotFound`` error (HTTP 404 /
     ``STATEMENT_NOT_FOUND``).
@@ -362,6 +377,7 @@ def open_session(
         session_conf: Optional[Dict[str, str]] = None
         if session_configuration:
             session_conf = {k: str(v) for k, v in session_configuration.items()}
+            _check_session_configuration(session_conf)
         # The kwarg builds run INSIDE the try so the ``finally`` scrub
         # below always fires — including when ``kernel_auth_kwargs``
         # itself raises mid-build (e.g. an OAuth token-exchange failure

tests/unit/test_kernel_client.py

@@ -89,6 +89,7 @@ def __init__(
 from databricks.sql.backend.types import CommandId, CommandState
 from databricks.sql.exc import (
     DatabaseError,
+    Error,
     InterfaceError,
     NotSupportedError,
     OperationalError,
@@ -306,6 +307,24 @@ def test_open_session_rejects_double_open(monkeypatch):
         c.open_session(session_configuration=None, catalog=None, schema=None)
 
 
+def test_open_session_rejects_timestamp_as_string_true(monkeypatch):
+    session_ctor = MagicMock()
+    monkeypatch.setattr(kernel_client._kernel, "Session", session_ctor)
+
+    c = _make_client()
+
+    with pytest.raises(Error, match="timestampAsString cannot be changed"):
+        c.open_session(
+            session_configuration={
+                "spark.thriftserver.arrowBasedRowSet.timestampAsString": True
+            },
+            catalog=None,
+            schema=None,
+        )
+
+    session_ctor.assert_not_called()
+
+
 @pytest.mark.parametrize(
     "kwargs, expected_flag",
     [