feat(kernel): wire statement query tags + http_headers/user_agent on use_kernel

Consume the kernel surface from kernel PR (query tags + custom HTTP
headers) on the use_kernel path. Bumps KERNEL_REV.

Query tags (statement-level):
- execute_command no longer raises NotSupportedError for query_tags;
  it calls stmt.set_query_tags(query_tags) after set_sql. The connector
  already passes Dict[str, Optional[str]], which the kernel accepts
  (None value -> bare key in the SEA query_tags conf).

http_headers + user_agent_entry:
- The kernel client now forwards http_headers to the kernel Session as
  the `http_headers` kwarg (was accept-and-ignore). session.py already
  passes all_headers, which carries the connector's composed User-Agent
  (PyDatabricksSqlConnector/x (entry)) + caller headers + SPOG org-id.
- The kernel applies them per request: its own Authorization / org-id
  win; a caller User-Agent is APPENDED to the kernel base UA (the base
  carries the DatabricksJDBCDriverOSS token that gates the SEA result
  disposition, so it's never replaced). So user_agent_entry is honored
  end-to-end via the existing http_headers forwarding — no separate
  kwarg needed.

Tests:
- unit: query_tags forwarded to set_query_tags (was: rejection test);
  http_headers forwarded to the kernel Session (and omitted when empty).
- e2e (test_kernel_backend.py): a query_tagged query and a connection
  with user_agent_entry + a custom http_header both round-trip green
  against a dogfood warehouse. The UA case specifically guards the
  append behavior (replacing the base UA would 400 on the result
  disposition).

KERNEL_REV -> c2053f68b75fef4a29425096dc6bbafb774d8b83 (kernel PR #119
branch HEAD; re-pin to the squash-merge SHA once #119 lands).

194 kernel unit tests pass; black + mypy clean; 3 e2e pass live.

Co-authored-by: Isaac
Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

KERNEL_REV

@@ -1 +1 @@
-b4d88220cdfad8dba1cfa89892269342ae26feeb
+c2053f68b75fef4a29425096dc6bbafb774d8b83

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

@@ -91,13 +91,19 @@ def __init__(
     ):
         # ``ssl_options`` is translated to the kernel's ``tls_*``
         # Session kwargs in ``open_session`` (custom CA, verify
-        # toggles, mTLS client cert/key). ``http_headers`` /
-        # ``http_client`` / ``port`` are still accept-and-ignore — the
-        # kernel manages its own HTTP stack.
+        # toggles, mTLS client cert/key). ``http_headers`` is forwarded
+        # to the kernel as custom request headers (it carries the
+        # connector's composed ``User-Agent`` + any caller headers + the
+        # SPOG ``x-databricks-org-id``). ``http_client`` / ``port`` are
+        # still accept-and-ignore — the kernel manages its own HTTP
+        # stack.
         self._server_hostname = server_hostname
         self._http_path = http_path
         self._auth_provider = auth_provider
         self._ssl_options = ssl_options
+        # Caller / connector HTTP headers (list of (name, value) pairs).
+        # Forwarded to the kernel Session in ``open_session``.
+        self._http_headers = http_headers or []
         # Raw auth-relevant connect() kwargs (auth_type,
         # oauth_client_id/secret, redirect port, credentials_provider).
         # The kernel auth bridge needs these to build OAuth kwargs — the
@@ -187,6 +193,15 @@ def open_session(
         # Translate the connector's ``_retry_*`` kwargs into the kernel's
         # ``retry_*`` Session kwargs. Empty when retry is left at defaults.
         retry_kwargs = _kernel_retry_kwargs(self._retry_options)
+        # Forward caller / connector HTTP headers. The kernel applies
+        # them on every request (its own Authorization / org-id win); a
+        # caller ``User-Agent`` is appended to the kernel's base UA. Only
+        # pass the kwarg when there's something to send.
+        http_headers_kwargs: Dict[str, Any] = {}
+        if self._http_headers:
+            http_headers_kwargs["http_headers"] = [
+                (str(k), str(v)) for k, v in self._http_headers
+            ]
         try:
             self._kernel_session = _kernel.Session(
                 host=self._server_hostname,
@@ -208,6 +223,7 @@ def open_session(
                 **auth_kwargs,
                 **tls_kwargs,
                 **retry_kwargs,
+                **http_headers_kwargs,
             )
         except Exception as exc:
             raise _wrap_kernel_exception("open_session", exc) from exc
@@ -304,10 +320,6 @@ def execute_command(
     ) -> Union["ResultSet", None]:
         if self._kernel_session is None:
             raise InterfaceError("Cannot execute_command without an open session.")
-        if query_tags:
-            raise NotSupportedError(
-                "Statement-level query_tags are not yet supported on the kernel backend."
-            )
 
         try:
             stmt = self._kernel_session.statement()
@@ -321,6 +333,13 @@ def execute_command(
         try:
             try:
                 stmt.set_sql(operation)
+                if query_tags:
+                    # Per-statement query tags. The kernel serialises the
+                    # dict (None value -> bare key) into the SEA
+                    # `query_tags` statement conf. ``query_tags`` is
+                    # already ``Dict[str, Optional[str]]`` from the
+                    # connector, which the kernel accepts directly.
+                    stmt.set_query_tags(query_tags)
                 if parameters:
                     bind_tspark_params(stmt, parameters)
                 if async_op:

tests/e2e/test_kernel_backend.py

@@ -332,3 +332,32 @@ def test_parameterized_query_decimal(conn):
         rows = cur.fetchall()
         # Server echoes back as decimal.Decimal.
         assert str(rows[0][0]) == "-123.45"
+
+
+def test_query_tags_round_trip(kernel_conn_params):
+    """Per-statement query_tags are forwarded to the kernel and accepted
+    by the server. Smoke-level: a malformed query_tags conf would fail
+    the execute. (query.history ingestion lag makes a sync tag-readback
+    assertion infeasible.)"""
+    with sql.connect(**kernel_conn_params) as c:
+        with c.cursor() as cur:
+            cur.execute(
+                "SELECT 1 AS n",
+                query_tags={"team": "platform", "production": None},
+            )
+            assert cur.fetchall()[0][0] == 1
+
+
+def test_user_agent_entry_and_http_headers_round_trip(kernel_conn_params):
+    """A connection with user_agent_entry (folded into the connector's
+    User-Agent, then appended to the kernel base UA) and a custom HTTP
+    header opens and queries cleanly. Replacing the kernel base UA would
+    break the SEA result disposition (HTTP 400); appending preserves it
+    — this exercises that end-to-end."""
+    params = dict(kernel_conn_params)
+    params["user_agent_entry"] = "kernel-e2e-app"
+    params["http_headers"] = [("X-Kernel-E2E", "yes")]
+    with sql.connect(**params) as c:
+        with c.cursor() as cur:
+            cur.execute("SELECT 1 AS n")
+            assert cur.fetchall()[0][0] == 1

tests/unit/test_kernel_client.py

@@ -332,26 +332,43 @@ def test_execute_command_forwards_parameters_to_bind_param():
     assert stmt.execute.called
 
 
-def test_execute_command_rejects_query_tags():
+def test_execute_command_forwards_query_tags():
+    """Statement-level query_tags are forwarded to the kernel statement
+    via set_query_tags (the kernel serialises them into the SEA
+    query_tags conf). Previously rejected with NotSupportedError; now
+    wired (kernel PR adding Statement.set_query_tags)."""
     c = _make_client()
     c._kernel_session = MagicMock()
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
-    with pytest.raises(NotSupportedError, match="query_tags"):
-        c.execute_command(
-            operation="SELECT 1",
-            session_id=MagicMock(),
-            max_rows=1,
-            max_bytes=1,
-            lz4_compression=False,
-            cursor=cursor,
-            use_cloud_fetch=False,
-            parameters=[],
-            async_op=False,
-            enforce_embedded_schema_correctness=False,
-            query_tags={"team": "x"},
-        )
+
+    stmt = MagicMock()
+    stmt.set_sql = MagicMock()
+    stmt.set_query_tags = MagicMock()
+    stmt.execute.return_value = MagicMock(
+        statement_id="stmt-id",
+        arrow_schema=MagicMock(return_value=pa.schema([("x", pa.int64())])),
+    )
+    c._kernel_session.statement.return_value = stmt
+
+    tags = {"team": "platform", "production": None}
+    c.execute_command(
+        operation="SELECT 1",
+        session_id=MagicMock(),
+        max_rows=1,
+        max_bytes=1,
+        lz4_compression=False,
+        cursor=cursor,
+        use_cloud_fetch=False,
+        parameters=[],
+        async_op=False,
+        enforce_embedded_schema_correctness=False,
+        query_tags=tags,
+    )
+
+    stmt.set_query_tags.assert_called_once_with(tags)
+    assert stmt.execute.called
 
 
 def test_get_columns_accepts_none_catalog():
@@ -1015,3 +1032,48 @@ def test_retry_delay_default_has_no_mapping(self):
         # recognised key here — it has no kernel equivalent.
         out = kernel_client._kernel_retry_kwargs({"retry_delay_default": 5.0})
         assert out == {}
+
+
+class TestKernelHttpHeadersForwarding:
+    """http_headers (the connector's caller headers + composed
+    User-Agent + SPOG org-id) are forwarded to the kernel Session as the
+    ``http_headers`` kwarg. The kernel applies them per request (its own
+    Authorization / org-id win; a caller User-Agent is appended to the
+    kernel base UA)."""
+
+    def _open_capturing(self, monkeypatch, http_headers):
+        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)
+        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,
+            http_headers=http_headers,
+        )
+        c.open_session(session_configuration=None, catalog=None, schema=None)
+        return captured
+
+    def test_http_headers_forwarded_to_kernel_session(self, monkeypatch):
+        headers = [
+            ("User-Agent", "PyDatabricksSqlConnector/4.0 (myentry)"),
+            ("X-Custom", "v1"),
+        ]
+        captured = self._open_capturing(monkeypatch, headers)
+        assert captured.get("http_headers") == [
+            ("User-Agent", "PyDatabricksSqlConnector/4.0 (myentry)"),
+            ("X-Custom", "v1"),
+        ]
+
+    def test_no_http_headers_omits_kwarg(self, monkeypatch):
+        # Empty/none headers → the kwarg isn't passed at all (kernel
+        # keeps its defaults).
+        captured = self._open_capturing(monkeypatch, [])
+        assert "http_headers" not in captured