diff --git a/dataconnect/client.py b/dataconnect/client.py index 76d172a..c510b44 100644 --- a/dataconnect/client.py +++ b/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)) diff --git a/dataconnect/transport/arrow_flight/transport.py b/dataconnect/transport/arrow_flight/transport.py index 801d90b..4458010 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/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())) diff --git a/tests/test_logging_headers.py b/tests/test_logging_headers.py new file mode 100644 index 0000000..1b73358 --- /dev/null +++ b/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"