[MCC-1487238] Python Lib - Add Headers to each request for logging

dataconnect/client.py

@@ -33,13 +33,14 @@ def connect(
         port: int = _DEFAULT_PORT,
         use_tls: bool = True,
         token: str = "",
+        user_uuid: str = "",
     ) -> DataConnectClient:
 
         # Import is deferred so pyarrow.flight is only loaded when this factory
         # is called — callers injecting a custom transport are unaffected.
         from dataconnect.transport.arrow_flight.transport import ArrowFlightTransport
 
-        transport = ArrowFlightTransport(host=host, port=port, use_tls=use_tls, token=token)
+        transport = ArrowFlightTransport(host=host, port=port, use_tls=use_tls, token=token, user_uuid=user_uuid)
 
         return cls(DefaultDataConnectService(transport))
 

dataconnect/transport/arrow_flight/transport.py

@@ -11,8 +11,10 @@
 import dataclasses
 import json
 import platform
+import socket
 import subprocess
 from datetime import UTC, datetime
+from importlib.metadata import version
 
 import pyarrow as pa
 import pyarrow.flight as flight
@@ -75,6 +77,7 @@ def __init__(
         port: int,
         use_tls: bool,
         token: str = "",
+        user_uuid: str = "",
     ) -> None:
         """Create a new Arrow Flight transport.
 
@@ -94,6 +97,26 @@ def __init__(
         except Exception as exc:
             raise parse_dataconnect_error(exc) from exc
 
+        if user_uuid:
+            self._call_headers.append((b"user-uuid", str(user_uuid).encode("utf-8")))
+
+        try:
+            sdk_version = version("dataconnect-library-python")
+        except Exception:
+            sdk_version = "0.1.0"
+        self._call_headers.append((b"python-sdk-version", sdk_version.encode("utf-8")))
+
+        self._call_headers.append((b"sdk-type", b"Python"))
+
+        try:
+            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            s.connect(("8.8.8.8", 80))
+            client_ip = s.getsockname()[0]
+            s.close()
+        except Exception:
+            client_ip = "127.0.0.1"
+        self._call_headers.append((b"client-ip", client_ip.encode("utf-8")))
+
         if token:
             self._call_headers.append((b"authorization", f"Bearer {token}".encode()))
 

tests/test_logging_headers.py

@@ -0,0 +1,121 @@
+"""Unit tests for ArrowFlightTransport request headers injection."""
+
+from __future__ import annotations
+
+import importlib.metadata
+import socket
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from dataconnect.client import DataConnectClient
+
+# ---------------------------------------------------------------------------
+# Happy-path
+# ---------------------------------------------------------------------------
+
+
+@patch("pyarrow.flight.FlightClient")
+def test_connect_injects_all_required_logging_headers(mock_flight_client: MagicMock) -> None:
+    test_uuid = "enodia-user-uuid-12345"
+    test_token = "some-bearer-token"
+    mocked_sdk_version = "v1.2.3-test"
+
+    with patch("dataconnect.transport.arrow_flight.transport.version", return_value=mocked_sdk_version):
+        client = DataConnectClient.connect(
+            host="localhost",
+            port=8888,
+            use_tls=False,
+            token=test_token,
+            user_uuid=test_uuid,
+        )
+
+    transport = client._service._transport
+    headers_dict = dict(transport._call_headers)
+
+    # AC-01: User UUID
+    assert b"user-uuid" in headers_dict
+    assert headers_dict[b"user-uuid"] == b"enodia-user-uuid-12345"
+
+    # AC-02 & AC-03: SDK Version & SDK Type consistency
+    assert b"python-sdk-version" in headers_dict
+    assert headers_dict[b"python-sdk-version"] == b"v1.2.3-test"
+    assert b"sdk-type" in headers_dict
+    assert headers_dict[b"sdk-type"] == b"Python"
+
+    # AC-04: Client IP Address validation
+    assert b"client-ip" in headers_dict
+    try:
+        socket.inet_aton(headers_dict[b"client-ip"].decode("utf-8"))
+        ip_is_valid = True
+    except OSError:
+        ip_is_valid = False
+    assert ip_is_valid
+
+    # Check that existing auth token logic is preserved
+    assert b"authorization" in headers_dict
+    assert headers_dict[b"authorization"] == f"Bearer {test_token}".encode()
+
+
+@pytest.mark.parametrize("simulated_version", ["0.5.0", "2.0.0-alpha"])
+@patch("pyarrow.flight.FlightClient")
+def test_sdk_type_remains_python_across_different_sdk_versions(
+    mock_flight_client: MagicMock, simulated_version: str
+) -> None:
+    with patch("dataconnect.transport.arrow_flight.transport.version", return_value=simulated_version):
+        client = DataConnectClient.connect(
+            host="localhost",
+            port=8888,
+            use_tls=False,
+            user_uuid="dummy-uuid",
+        )
+
+    transport = client._service._transport
+    headers_dict = dict(transport._call_headers)
+
+    assert headers_dict[b"python-sdk-version"] == simulated_version.encode("utf-8")
+    assert headers_dict[b"sdk-type"] == b"Python"
+
+
+# ---------------------------------------------------------------------------
+# Edge cases & Error handling
+# ---------------------------------------------------------------------------
+
+
+@patch("pyarrow.flight.FlightClient")
+def test_client_ip_falls_back_to_localhost_on_socket_error(mock_flight_client: MagicMock) -> None:
+    """AC-04: Verifică dacă IP-ul face fallback pe 127.0.0.1 când mașina e offline."""
+    with (
+        patch("socket.socket.connect", side_effect=OSError("No network")),
+        patch("dataconnect.transport.arrow_flight.transport.version", return_value="1.0.0"),
+    ):
+        client = DataConnectClient.connect(
+            host="localhost",
+            port=8888,
+            use_tls=False,
+            user_uuid="dummy-uuid",
+        )
+
+    transport = client._service._transport
+    headers_dict = dict(transport._call_headers)
+
+    assert b"client-ip" in headers_dict
+    assert headers_dict[b"client-ip"] == b"127.0.0.1"
+
+
+@patch("pyarrow.flight.FlightClient")
+def test_sdk_version_falls_back_on_package_not_found_error(mock_flight_client: MagicMock) -> None:
+    """AC-02: Verifică dacă versiunea face fallback pe 0.1.0 când metadatele pachetului lipsesc."""
+    with patch("importlib.metadata.version", side_effect=importlib.metadata.PackageNotFoundError):
+        client = DataConnectClient.connect(
+            host="localhost",
+            port=8888,
+            use_tls=False,
+            user_uuid="dummy-uuid",
+        )
+
+    transport = client._service._transport
+    headers_dict = dict(transport._call_headers)
+
+    assert b"python-sdk-version" in headers_dict
+    assert headers_dict[b"python-sdk-version"] == b"0.1.0"