Add databricks-impulse into SDK and Spark Connect channel user-agents

src/impulse_query_engine/__init__.py

@@ -7,3 +7,11 @@
     from pathlib import Path
 
     __version__ = (Path(__file__).resolve().parents[2] / "VERSION").read_text().strip()
+
+try:
+    import databricks.sdk.useragent as _ua
+
+    _ua.with_extra("databricks-impulse", __version__)
+    _ua.with_product("databricks-impulse", __version__)
+except Exception:  # noqa: BLE001 - telemetry must never break import
+    pass

src/impulse_query_engine/telemetry.py

@@ -1,12 +1,17 @@
 import functools
+import inspect
 import logging
 from collections.abc import Callable
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import DatabricksError
 
+from impulse_query_engine import __version__
+
 logger = logging.getLogger(__name__)
 
+PRODUCT_NAME = "databricks-impulse"
+
 
 def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
     """Trace telemetry via the Databricks User-Agent header.
@@ -35,6 +40,38 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
         logger.debug(f"Databricks workspace is not available: {e}")
 
 
+def tag_spark_connect_user_agent(spark, product: str, version: str) -> None:
+    """Best-effort: tag a Spark Connect session's per-request user-agent.
+
+    Parameters
+    ----------
+    spark : SparkSession
+        The Spark session used for query execution. Only Spark Connect sessions are
+        tagged; anything else is silently ignored.
+    product : str
+        Product identifier (e.g. ``"databricks-impulse"``).
+    version : str
+        Product version string.
+    """
+    # The Spark Connect connection-string parameter name for the user agent
+    # (``ChannelBuilder.PARAM_USER_AGENT``). Referenced by value to avoid importing
+    # ``pyspark.sql.connect.client.core``, which pulls in grpcio and is absent outside a
+    # Connect client anyway.
+    user_agent_key = "user_agent"
+    try:
+        builder = spark._client._builder  # SparkConnectClient -> ChannelBuilder
+        tag = f"{product}/{version}"
+        existing = builder._params.get(user_agent_key, "")
+        # Idempotent: skip if this product already tags the session (match on the
+        # ``<product>/`` prefix, not the full tag, so a version substring such as
+        # 0.6.1 vs 0.6.10 can't false-match and versions don't stack on re-tag).
+        if f"{product}/" not in existing:
+            builder._params[user_agent_key] = f"{tag} {existing}".strip()
+            logger.debug(f"Tagged Spark Connect user-agent with {tag}")
+    except Exception as e:  # noqa: BLE001 - classic session or internals changed
+        logger.debug(f"Skipped Spark Connect user-agent tagging: {e}")
+
+
 def telemetry_logger(key: str, value: str, workspace_client_attr: str = "ws") -> Callable:
     """Decorator that logs telemetry before executing the wrapped method.
 
@@ -52,6 +89,7 @@ def telemetry_logger(key: str, value: str, workspace_client_attr: str = "ws") ->
     """
 
     def decorator(func: Callable) -> Callable:
+        sig = inspect.signature(func)
 
         @functools.wraps(func)
         def wrapper(self, *args, **kwargs):
@@ -64,6 +102,16 @@ def wrapper(self, *args, **kwargs):
                     f"on {self.__class__.__name__}. "
                     f"Make sure your class has the specified workspace client attribute."
                 )
+            # If the wrapped method received a Spark session, tag it so external
+            # Spark Connect compute is attributable (best-effort; no-op otherwise).
+            # Guarded like every other telemetry step here: a binding hiccup must
+            # never break the wrapped business call.
+            try:
+                spark = sig.bind_partial(self, *args, **kwargs).arguments.get("spark")
+            except TypeError:
+                spark = None
+            if spark is not None:
+                tag_spark_connect_user_agent(spark, PRODUCT_NAME, __version__)
             return func(self, *args, **kwargs)
 
         return wrapper

tests/impulse_query_engine/unit/telemetry_test.py

@@ -1,10 +1,19 @@
+import importlib
+from types import SimpleNamespace
 from unittest.mock import MagicMock, create_autospec, patch
 
 import pytest
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import DatabricksError
 
-from impulse_query_engine.telemetry import log_telemetry, telemetry_logger, verify_workspace_client
+import databricks.sdk.useragent as ua
+import impulse_query_engine
+from impulse_query_engine.telemetry import (
+    log_telemetry,
+    tag_spark_connect_user_agent,
+    telemetry_logger,
+    verify_workspace_client,
+)
 
 
 class TestLogTelemetry:
@@ -77,6 +86,43 @@ def toPandas(self):
         mock_log.assert_called_once_with(ws, "query", "to_pandas")
         assert result == "done"
 
+    def test_decorator_tags_spark_connect_session_when_spark_arg_present(self):
+        builder = SimpleNamespace(_params={"user_agent": "databricks-session"})
+
+        class QueryBuilder:
+            def __init__(self, ws):
+                self.ws = ws
+
+            @telemetry_logger("query", "solve")
+            def solve(self, spark):
+                return "result"
+
+        ws = create_autospec(WorkspaceClient)
+        spark = SimpleNamespace(_client=SimpleNamespace(_builder=builder))
+
+        with patch("impulse_query_engine.telemetry.log_telemetry"):
+            QueryBuilder(ws).solve(spark)
+
+        assert builder._params["user_agent"].startswith("databricks-impulse/")
+
+    def test_decorator_does_not_tag_when_no_spark_arg(self):
+        # Methods without a ``spark`` parameter must not raise or attempt tagging.
+        class QueryBuilder:
+            def __init__(self, ws):
+                self.ws = ws
+
+            @telemetry_logger("query", "solve")
+            def solve(self):
+                return "result"
+
+        ws = create_autospec(WorkspaceClient)
+        with (
+            patch("impulse_query_engine.telemetry.log_telemetry"),
+            patch("impulse_query_engine.telemetry.tag_spark_connect_user_agent") as mock_tag,
+        ):
+            QueryBuilder(ws).solve()
+        mock_tag.assert_not_called()
+
     def test_decorator_preserves_function_metadata(self):
         class QueryBuilder:
             def __init__(self):
@@ -125,3 +171,50 @@ def test_raises_databricks_error_when_workspace_unreachable(self):
 
         with pytest.raises(DatabricksError):
             verify_workspace_client(ws, "mda", "0.0.4")
+
+
+class TestGlobalUserAgentRegistration:
+    """Importing the package registers impulse in the SDK's process-global user-agent."""
+
+    def test_import_registers_product_and_extra(self):
+        # Re-run the registration in-body (rather than relying on collection-time
+        # global state) so the assertion is independent of suite ordering.
+        importlib.reload(impulse_query_engine)
+        assert ua.product() == ("databricks-impulse", impulse_query_engine.__version__)
+        assert ("databricks-impulse", impulse_query_engine.__version__) in ua._extra
+
+    def test_registration_failure_does_not_break_import(self):
+        # A failure inside the registration block must never propagate out of import.
+        with patch.object(ua, "with_product", side_effect=ValueError("boom")):
+            importlib.reload(impulse_query_engine)  # must not raise
+        # Restore a clean registration for any subsequent assertions in the session.
+        importlib.reload(impulse_query_engine)
+        assert ua.product() == ("databricks-impulse", impulse_query_engine.__version__)
+
+
+class TestTagSparkConnectUserAgent:
+    @staticmethod
+    def _fake_connect_session(user_agent="databricks-session"):
+        builder = SimpleNamespace(_params={"user_agent": user_agent})
+        return SimpleNamespace(_client=SimpleNamespace(_builder=builder)), builder
+
+    def test_prepends_tag_to_connect_session(self):
+        spark, builder = self._fake_connect_session()
+        tag_spark_connect_user_agent(spark, "databricks-impulse", "0.6.1")
+        assert builder._params["user_agent"] == "databricks-impulse/0.6.1 databricks-session"
+
+    def test_is_idempotent(self):
+        spark, builder = self._fake_connect_session()
+        tag_spark_connect_user_agent(spark, "databricks-impulse", "0.6.1")
+        tag_spark_connect_user_agent(spark, "databricks-impulse", "0.6.1")
+        # Tag appears exactly once; the existing value is preserved.
+        assert builder._params["user_agent"] == "databricks-impulse/0.6.1 databricks-session"
+
+    def test_handles_empty_existing_user_agent(self):
+        spark, builder = self._fake_connect_session(user_agent="")
+        tag_spark_connect_user_agent(spark, "databricks-impulse", "0.6.1")
+        assert builder._params["user_agent"] == "databricks-impulse/0.6.1"
+
+    def test_classic_session_is_a_silent_no_op(self):
+        # A classic SparkSession has no ``_client`` — must not raise.
+        tag_spark_connect_user_agent(object(), "databricks-impulse", "0.6.1")