diff --git a/src/impulse_query_engine/__init__.py b/src/impulse_query_engine/__init__.py index e9a99f08..831da88d 100644 --- a/src/impulse_query_engine/__init__.py +++ b/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 diff --git a/src/impulse_query_engine/telemetry.py b/src/impulse_query_engine/telemetry.py index 002c229b..b486902d 100644 --- a/src/impulse_query_engine/telemetry.py +++ b/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 + # ``/`` 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 diff --git a/tests/impulse_query_engine/unit/telemetry_test.py b/tests/impulse_query_engine/unit/telemetry_test.py index d15201c8..5dbaf450 100644 --- a/tests/impulse_query_engine/unit/telemetry_test.py +++ b/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")