feat: emit the X-Wherobots-Client attribution header (WBC-819)

DAG-submitted work was indistinguishable from any other unattributed
traffic, so Airflow usage could not be measured or debugged. Both hooks
now identify themselves on the shared, advisory `X-Wherobots-Client`
chain with the origin hop `client=airflow;ver=<provider version>`.

The REST hook talks to the platform directly and sends the hop on every
call. The SQL hook goes through the Python DB-API driver, so it passes
the hop down via the driver's `extra_headers`; the driver appends its
own hop on the right, yielding `client=airflow;..., client=dbapi;...`
and keeping Airflow as the origin. `extra_headers` postdates the
`wherobots-python-dbapi>=0.28.0` floor, so it is probed rather than
assumed and an older driver degrades to today's behaviour.

The version falls back to `unknown` when the distribution metadata is
missing, and hop values are stripped of the `,` and `;` delimiters so
the grammar can never be broken.

Author: ClayMav

airflow_providers_wherobots/client_attribution.py

@@ -0,0 +1,61 @@
+"""
+Client attribution via the shared ``X-Wherobots-Client`` header.
+
+``X-Wherobots-Client`` is an ordered, append-only, comma-separated list of
+hops modelled on ``X-Forwarded-For``: the leftmost hop is the ORIGIN client
+and every component appends its own hop on the right. It lets Wherobots
+services attribute a request to the client it came from.
+
+This provider is an origin client: an Airflow DAG is where the request enters
+the Wherobots client ecosystem, so it emits ``client=airflow;ver=<version>``
+as the leftmost hop. Where the provider goes through another Wherobots client
+(the Python DB-API driver, for the SQL hook), that client appends its own hop
+to the right, producing e.g.
+``client=airflow;ver=1.7.0, client=dbapi;ver=0.28.1``.
+
+The header is advisory: it is client-asserted and informational only, and must
+never influence authentication or authorization.
+"""
+
+from importlib import metadata
+from typing import Dict, Final
+
+from airflow_providers_wherobots.hooks.base import PACKAGE_NAME
+
+# Canonical name of the shared, cross-service client-chain header.
+WHEROBOTS_CLIENT_HEADER: Final[str] = "X-Wherobots-Client"
+
+# Canonical, stable vocabulary token for this client. Renaming it splits its
+# history in the platform's attribution analytics, so it must not change.
+CLIENT_TOKEN: Final[str] = "airflow"
+
+# Sentinel used when the installed distribution's version can't be resolved
+# (e.g. the provider is imported from a source tree that was never installed).
+UNKNOWN_VERSION: Final[str] = "unknown"
+
+# Commas separate hops and semicolons separate a hop's fields, so neither may
+# appear inside a value.
+_DELIMITERS = str.maketrans({",": "_", ";": "_"})
+
+
+def _resolve_provider_version() -> str:
+    """Return the installed provider version, or ``unknown`` if unavailable."""
+    try:
+        return metadata.version(PACKAGE_NAME)
+    except metadata.PackageNotFoundError:
+        return UNKNOWN_VERSION
+
+
+# Resolved once at import: `importlib.metadata.version` scans the installed
+# package database on each call, and the version can't change within a process.
+PROVIDER_VERSION: Final[str] = _resolve_provider_version()
+
+# This provider's single hop, e.g. `client=airflow;ver=1.7.0`.
+CLIENT_HOP: Final[str] = (
+    f"client={CLIENT_TOKEN};ver={PROVIDER_VERSION.translate(_DELIMITERS)}"
+)
+
+
+def client_attribution_header() -> Dict[str, str]:
+    """Return the ``X-Wherobots-Client`` header carrying this provider's hop."""
+    return {WHEROBOTS_CLIENT_HEADER: CLIENT_HOP}

airflow_providers_wherobots/hooks/rest_api.py

@@ -16,6 +16,7 @@
 from requests.auth import AuthBase
 from wherobots.db import Region
 
+from airflow_providers_wherobots.client_attribution import client_attribution_header
 from airflow_providers_wherobots.hooks.base import (
     DEFAULT_CONN_ID,
     PACKAGE_NAME,
@@ -85,6 +86,15 @@ def user_agent_header(self):
         )
         return {"User-Agent": header_value}
 
+    @cached_property
+    def default_headers(self) -> Dict[str, str]:
+        """Headers sent on every REST API request.
+
+        The hook talks to the platform directly, so it is the origin hop of the
+        advisory ``X-Wherobots-Client`` chain.
+        """
+        return {**self.user_agent_header, **client_attribution_header()}
+
     def _api_call(
         self,
         method: str,
@@ -100,7 +110,7 @@ def _api_call(
             json=payload,
             auth=auth,
             params=params,
-            headers=self.user_agent_header,
+            headers=self.default_headers,
         )
         try:
             resp.raise_for_status()

airflow_providers_wherobots/hooks/sql.py

@@ -2,7 +2,8 @@
 Hook for Wherobots' Spatial SQL API interface.
 """
 
-from typing import Optional, Union
+import inspect
+from typing import Any, Dict, Final, Optional, Union
 
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from wherobots.db import Connection as WDBConnection, connect
@@ -15,8 +16,24 @@
 from wherobots.db.runtime import Runtime
 from wherobots.db.session_type import SessionType
 
+from airflow_providers_wherobots.client_attribution import client_attribution_header
 from airflow_providers_wherobots.hooks.base import DEFAULT_CONN_ID
 
+# This hook reaches the platform through the Python DB-API driver rather than
+# directly. The driver appends its own `client=dbapi;ver=...` hop to any
+# inbound `X-Wherobots-Client` chain, so handing it our hop via `extra_headers`
+# keeps Airflow as the leftmost (origin) hop instead of letting the driver look
+# like the origin.
+#
+# `extra_headers` is newer than this provider's `wherobots-python-dbapi>=0.28.0`
+# floor, so it is probed rather than assumed: against an older driver the
+# provider still works and simply loses the Airflow hop. Delete this probe and
+# pass the argument unconditionally once the floor is raised to a release that
+# has it.
+_CONNECT_SUPPORTS_EXTRA_HEADERS: Final[bool] = (
+    "extra_headers" in inspect.signature(connect).parameters
+)
+
 
 class WherobotsSqlHook(DbApiHook):  # type: ignore[misc]
     conn_name_attr = "wherobots_conn_id"
@@ -50,6 +67,11 @@ def _create_or_get_sql_session(
         self,
         runtime: Optional[Union[str, Runtime]] = None,
     ) -> WDBConnection:
+        attribution: Dict[str, Any] = (
+            {"extra_headers": client_attribution_header()}
+            if _CONNECT_SUPPORTS_EXTRA_HEADERS
+            else {}
+        )
         return connect(
             host=self._conn.host,
             api_key=self._conn.password,
@@ -60,6 +82,7 @@ def _create_or_get_sql_session(
             read_timeout=self.read_timeout,
             force_new=self.force_new,
             session_type=self.session_type,
+            **attribution,
         )
 
     def get_conn(self) -> WDBConnection:

tests/unit_tests/hooks/test_rest_api.py

@@ -4,6 +4,7 @@
 
 import json
 from http import HTTPStatus
+from importlib import metadata
 
 import airflow
 import requests
@@ -13,10 +14,15 @@
 from responses import matchers
 from wherobots.db import Runtime, Region
 
+from airflow_providers_wherobots.client_attribution import (
+    CLIENT_HOP,
+    WHEROBOTS_CLIENT_HEADER,
+)
 from airflow_providers_wherobots.hooks.rest_api import (
     WherobotsAuth,
     WherobotsRestAPIHook,
 )
+from airflow_providers_wherobots.hooks.base import PACKAGE_NAME
 from airflow_providers_wherobots.wherobots.models import (
     Run,
     LogsResponse,
@@ -195,6 +201,31 @@ def test_cancel_run(self, test_default_conn) -> None:
         with WherobotsRestAPIHook() as hook:
             hook.cancel_run(run_id=test_run.ext_id)
 
+    @responses.activate
+    def test_api_call_sends_client_attribution_header(self, test_default_conn) -> None:
+        """Every REST call carries this provider's origin hop."""
+        url = f"https://{test_default_conn.host}/test"
+        responses.add(
+            responses.GET,
+            url,
+            json={},
+            status=HTTPStatus.OK,
+            match=[matchers.header_matcher({WHEROBOTS_CLIENT_HEADER: CLIENT_HOP})],
+        )
+        with WherobotsRestAPIHook() as hook:
+            hook._api_call("GET", "/test")
+
+        sent = responses.calls[0].request.headers[WHEROBOTS_CLIENT_HEADER]
+        assert sent == f"client=airflow;ver={metadata.version(PACKAGE_NAME)}"
+
+    def test_default_headers_keep_the_user_agent(self, test_default_conn) -> None:
+        """Client attribution is additive: the User-Agent is still sent."""
+        with WherobotsRestAPIHook() as hook:
+            assert hook.default_headers == {
+                **hook.user_agent_header,
+                WHEROBOTS_CLIENT_HEADER: CLIENT_HOP,
+            }
+
     def test_user_agent(self, test_default_conn, mocker: MockerFixture) -> None:
         """
         Test the user_agent_header property

tests/unit_tests/hooks/test_sql.py

@@ -6,10 +6,15 @@
 from unittest.mock import MagicMock
 
 from airflow.models import Connection
+from pytest_mock import MockerFixture
 from wherobots.db.region import Region
 from wherobots.db.runtime import Runtime
 from wherobots.db.session_type import SessionType
 
+from airflow_providers_wherobots.client_attribution import (
+    CLIENT_HOP,
+    WHEROBOTS_CLIENT_HEADER,
+)
 from airflow_providers_wherobots.hooks.sql import WherobotsSqlHook
 
 
@@ -37,3 +42,35 @@ def test_get_conn(self, mock_connect: MagicMock, test_default_conn: Connection):
             force_new=True,
             session_type=SessionType.SINGLE,
         )
+
+    @mock.patch("airflow_providers_wherobots.hooks.sql.connect")
+    def test_get_conn_passes_client_attribution_to_the_driver(
+        self,
+        mock_connect: MagicMock,
+        test_default_conn: Connection,
+        mocker: MockerFixture,
+    ):
+        """Our hop goes to the driver, which appends its own hop to the right."""
+        mocker.patch(
+            "airflow_providers_wherobots.hooks.sql._CONNECT_SUPPORTS_EXTRA_HEADERS",
+            True,
+        )
+        WherobotsSqlHook().get_conn()
+        assert mock_connect.call_args.kwargs["extra_headers"] == {
+            WHEROBOTS_CLIENT_HEADER: CLIENT_HOP
+        }
+
+    @mock.patch("airflow_providers_wherobots.hooks.sql.connect")
+    def test_get_conn_omits_client_attribution_on_older_drivers(
+        self,
+        mock_connect: MagicMock,
+        test_default_conn: Connection,
+        mocker: MockerFixture,
+    ):
+        """A driver without `extra_headers` still connects, just unattributed."""
+        mocker.patch(
+            "airflow_providers_wherobots.hooks.sql._CONNECT_SUPPORTS_EXTRA_HEADERS",
+            False,
+        )
+        WherobotsSqlHook().get_conn()
+        assert "extra_headers" not in mock_connect.call_args.kwargs

tests/unit_tests/test_client_attribution.py

@@ -0,0 +1,75 @@
+"""
+Test the shared ``X-Wherobots-Client`` client-attribution hop.
+"""
+
+import importlib
+import re
+from importlib import metadata
+
+from pytest_mock import MockerFixture
+
+from airflow_providers_wherobots import client_attribution
+from airflow_providers_wherobots.client_attribution import (
+    CLIENT_HOP,
+    UNKNOWN_VERSION,
+    WHEROBOTS_CLIENT_HEADER,
+    client_attribution_header,
+)
+from airflow_providers_wherobots.hooks.base import PACKAGE_NAME
+
+# One hop: `client=<token>` plus zero or more `;key=value` params, where no
+# value may contain the `,` / `;` delimiters.
+HOP_PATTERN = re.compile(r"^client=[^,;]+(?:;[^,;=]+=[^,;]*)*$")
+
+
+def test_header_name_and_hop_format() -> None:
+    """The hook emits exactly one canonical header carrying one well-formed hop."""
+    header = client_attribution_header()
+    assert list(header) == [WHEROBOTS_CLIENT_HEADER]
+    assert WHEROBOTS_CLIENT_HEADER == "X-Wherobots-Client"
+
+    hop = header[WHEROBOTS_CLIENT_HEADER]
+    assert HOP_PATTERN.match(hop)
+    assert hop == f"client=airflow;ver={metadata.version(PACKAGE_NAME)}"
+
+
+def test_hop_is_within_the_length_bound() -> None:
+    """The chain is bounded to 512 bytes, and tokens to 64 characters."""
+    assert len(CLIENT_HOP.encode("utf-8")) <= 512
+    assert len("airflow") <= 64
+
+
+def test_unresolvable_version_falls_back_to_unknown(mocker: MockerFixture) -> None:
+    """A provider that isn't installed still emits a well-formed hop."""
+    mocker.patch.object(
+        client_attribution.metadata,
+        "version",
+        side_effect=metadata.PackageNotFoundError(PACKAGE_NAME),
+    )
+    assert client_attribution._resolve_provider_version() == UNKNOWN_VERSION
+
+    # The hop is built once at import, so reload the module under the patch to
+    # observe the header a never-installed provider would actually send.
+    try:
+        reloaded = importlib.reload(client_attribution)
+        assert reloaded.client_attribution_header() == {
+            WHEROBOTS_CLIENT_HEADER: "client=airflow;ver=unknown"
+        }
+        assert HOP_PATTERN.match(reloaded.CLIENT_HOP)
+    finally:
+        mocker.stopall()
+        importlib.reload(client_attribution)
+
+
+def test_delimiters_are_stripped_from_the_version(mocker: MockerFixture) -> None:
+    """A version can never smuggle a hop or field separator into the grammar."""
+    mocker.patch.object(
+        client_attribution.metadata, "version", return_value="1.0;0,dev"
+    )
+    try:
+        reloaded = importlib.reload(client_attribution)
+        assert reloaded.CLIENT_HOP == "client=airflow;ver=1.0_0_dev"
+        assert HOP_PATTERN.match(reloaded.CLIENT_HOP)
+    finally:
+        mocker.stopall()
+        importlib.reload(client_attribution)