Skip to content

Commit e987a66

Browse files
[MCC-1487238] Python Lib - Add Headers to each request for logging (#36)
* feat(logging): MCC-1487238 - inject logging headers into flight transport - Thread user_uuid from DataConnectClient.connect down to ArrowFlightTransport - Dynamically resolve client IP address via UDP socket with local loopback fallback - Dynamically read python-sdk-version with safe fallback on PackageNotFoundError - Force sdk-type header to 'Python' - Add comprehensive happy path and defensive edge-case unit tests * chore(lint): fix ruff formatting and import ordering issues - Fix import order sorting using ruff rules - Remove trailing whitespaces and redundant empty lines * style: remove commented-out code and redundant comments - Clean up unused or commented-out lines of code - Remove obsolete inline documentation to improve readability
1 parent 191f425 commit e987a66

3 files changed

Lines changed: 146 additions & 1 deletion

File tree

dataconnect/client.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,13 +33,14 @@ def connect(
3333
port: int = _DEFAULT_PORT,
3434
use_tls: bool = True,
3535
token: str = "",
36+
user_uuid: str = "",
3637
) -> DataConnectClient:
3738

3839
# Import is deferred so pyarrow.flight is only loaded when this factory
3940
# is called — callers injecting a custom transport are unaffected.
4041
from dataconnect.transport.arrow_flight.transport import ArrowFlightTransport
4142

42-
transport = ArrowFlightTransport(host=host, port=port, use_tls=use_tls, token=token)
43+
transport = ArrowFlightTransport(host=host, port=port, use_tls=use_tls, token=token, user_uuid=user_uuid)
4344

4445
return cls(DefaultDataConnectService(transport))
4546

dataconnect/transport/arrow_flight/transport.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,10 @@
1111
import dataclasses
1212
import json
1313
import platform
14+
import socket
1415
import subprocess
1516
from datetime import UTC, datetime
17+
from importlib.metadata import version
1618

1719
import pyarrow as pa
1820
import pyarrow.flight as flight
@@ -75,6 +77,7 @@ def __init__(
7577
port: int,
7678
use_tls: bool,
7779
token: str = "",
80+
user_uuid: str = "",
7881
) -> None:
7982
"""Create a new Arrow Flight transport.
8083
@@ -94,6 +97,26 @@ def __init__(
9497
except Exception as exc:
9598
raise parse_dataconnect_error(exc) from exc
9699

100+
if user_uuid:
101+
self._call_headers.append((b"user-uuid", str(user_uuid).encode("utf-8")))
102+
103+
try:
104+
sdk_version = version("dataconnect-library-python")
105+
except Exception:
106+
sdk_version = "0.1.0"
107+
self._call_headers.append((b"python-sdk-version", sdk_version.encode("utf-8")))
108+
109+
self._call_headers.append((b"sdk-type", b"Python"))
110+
111+
try:
112+
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
113+
s.connect(("8.8.8.8", 80))
114+
client_ip = s.getsockname()[0]
115+
s.close()
116+
except Exception:
117+
client_ip = "127.0.0.1"
118+
self._call_headers.append((b"client-ip", client_ip.encode("utf-8")))
119+
97120
if token:
98121
self._call_headers.append((b"authorization", f"Bearer {token}".encode()))
99122

tests/test_logging_headers.py

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
"""Unit tests for ArrowFlightTransport request headers injection."""
2+
3+
from __future__ import annotations
4+
5+
import importlib.metadata
6+
import socket
7+
from unittest.mock import MagicMock, patch
8+
9+
import pytest
10+
11+
from dataconnect.client import DataConnectClient
12+
13+
# ---------------------------------------------------------------------------
14+
# Happy-path
15+
# ---------------------------------------------------------------------------
16+
17+
18+
@patch("pyarrow.flight.FlightClient")
19+
def test_connect_injects_all_required_logging_headers(mock_flight_client: MagicMock) -> None:
20+
test_uuid = "enodia-user-uuid-12345"
21+
test_token = "some-bearer-token"
22+
mocked_sdk_version = "v1.2.3-test"
23+
24+
with patch("dataconnect.transport.arrow_flight.transport.version", return_value=mocked_sdk_version):
25+
client = DataConnectClient.connect(
26+
host="localhost",
27+
port=8888,
28+
use_tls=False,
29+
token=test_token,
30+
user_uuid=test_uuid,
31+
)
32+
33+
transport = client._service._transport
34+
headers_dict = dict(transport._call_headers)
35+
36+
# AC-01: User UUID
37+
assert b"user-uuid" in headers_dict
38+
assert headers_dict[b"user-uuid"] == b"enodia-user-uuid-12345"
39+
40+
# AC-02 & AC-03: SDK Version & SDK Type consistency
41+
assert b"python-sdk-version" in headers_dict
42+
assert headers_dict[b"python-sdk-version"] == b"v1.2.3-test"
43+
assert b"sdk-type" in headers_dict
44+
assert headers_dict[b"sdk-type"] == b"Python"
45+
46+
# AC-04: Client IP Address validation
47+
assert b"client-ip" in headers_dict
48+
try:
49+
socket.inet_aton(headers_dict[b"client-ip"].decode("utf-8"))
50+
ip_is_valid = True
51+
except OSError:
52+
ip_is_valid = False
53+
assert ip_is_valid
54+
55+
# Check that existing auth token logic is preserved
56+
assert b"authorization" in headers_dict
57+
assert headers_dict[b"authorization"] == f"Bearer {test_token}".encode()
58+
59+
60+
@pytest.mark.parametrize("simulated_version", ["0.5.0", "2.0.0-alpha"])
61+
@patch("pyarrow.flight.FlightClient")
62+
def test_sdk_type_remains_python_across_different_sdk_versions(
63+
mock_flight_client: MagicMock, simulated_version: str
64+
) -> None:
65+
with patch("dataconnect.transport.arrow_flight.transport.version", return_value=simulated_version):
66+
client = DataConnectClient.connect(
67+
host="localhost",
68+
port=8888,
69+
use_tls=False,
70+
user_uuid="dummy-uuid",
71+
)
72+
73+
transport = client._service._transport
74+
headers_dict = dict(transport._call_headers)
75+
76+
assert headers_dict[b"python-sdk-version"] == simulated_version.encode("utf-8")
77+
assert headers_dict[b"sdk-type"] == b"Python"
78+
79+
80+
# ---------------------------------------------------------------------------
81+
# Edge cases & Error handling
82+
# ---------------------------------------------------------------------------
83+
84+
85+
@patch("pyarrow.flight.FlightClient")
86+
def test_client_ip_falls_back_to_localhost_on_socket_error(mock_flight_client: MagicMock) -> None:
87+
"""AC-04: Verifică dacă IP-ul face fallback pe 127.0.0.1 când mașina e offline."""
88+
with (
89+
patch("socket.socket.connect", side_effect=OSError("No network")),
90+
patch("dataconnect.transport.arrow_flight.transport.version", return_value="1.0.0"),
91+
):
92+
client = DataConnectClient.connect(
93+
host="localhost",
94+
port=8888,
95+
use_tls=False,
96+
user_uuid="dummy-uuid",
97+
)
98+
99+
transport = client._service._transport
100+
headers_dict = dict(transport._call_headers)
101+
102+
assert b"client-ip" in headers_dict
103+
assert headers_dict[b"client-ip"] == b"127.0.0.1"
104+
105+
106+
@patch("pyarrow.flight.FlightClient")
107+
def test_sdk_version_falls_back_on_package_not_found_error(mock_flight_client: MagicMock) -> None:
108+
"""AC-02: Verifică dacă versiunea face fallback pe 0.1.0 când metadatele pachetului lipsesc."""
109+
with patch("importlib.metadata.version", side_effect=importlib.metadata.PackageNotFoundError):
110+
client = DataConnectClient.connect(
111+
host="localhost",
112+
port=8888,
113+
use_tls=False,
114+
user_uuid="dummy-uuid",
115+
)
116+
117+
transport = client._service._transport
118+
headers_dict = dict(transport._call_headers)
119+
120+
assert b"python-sdk-version" in headers_dict
121+
assert headers_dict[b"python-sdk-version"] == b"0.1.0"

0 commit comments

Comments
 (0)