From ba3926b2aca58deb02c887ce3923094cb11b0840 Mon Sep 17 00:00:00 2001 From: Gaurav Londhe Date: Thu, 10 Sep 2026 18:06:24 +0530 Subject: [PATCH 1/2] refactor: Remove client IP address retrieval from ArrowFlightTransport and related tests --- .../transport/arrow_flight/transport.py | 10 ------ tests/test_logging_headers.py | 36 +------------------ 2 files changed, 1 insertion(+), 45 deletions(-) diff --git a/dataconnect/transport/arrow_flight/transport.py b/dataconnect/transport/arrow_flight/transport.py index a869c73..55e69be 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/dataconnect/transport/arrow_flight/transport.py @@ -11,7 +11,6 @@ import dataclasses import json import platform -import socket import subprocess from datetime import UTC, datetime from importlib.metadata import version @@ -133,15 +132,6 @@ def __init__( client_info_value = f"Python_SDK;{sdk_version};" self._call_headers.append((b"x-client-dataconnect", client_info_value.encode("utf-8"))) - 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"x-client-public-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 index 8fc6f0e..55a42e1 100644 --- a/tests/test_logging_headers.py +++ b/tests/test_logging_headers.py @@ -3,7 +3,6 @@ from __future__ import annotations import importlib.metadata -import socket from unittest.mock import MagicMock, patch import pytest @@ -37,14 +36,7 @@ def test_connect_injects_all_required_logging_headers(mock_flight_client: MagicM assert b"x-client-dataconnect" in headers_dict assert headers_dict[b"x-client-dataconnect"] == b"Python_SDK;v1.2.3-test;" - # AC-04: Client IP Address validation - assert b"x-client-public-ip" in headers_dict - try: - socket.inet_aton(headers_dict[b"x-client-public-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 @@ -71,32 +63,6 @@ def test_sdk_type_remains_python_across_different_sdk_versions( assert headers_dict[b"x-client-dataconnect"] == expected_client_info -# --------------------------------------------------------------------------- -# 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: Check if the IP falls back to 127.0.0.1 when the machine is 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"x-client-public-ip" in headers_dict - assert headers_dict[b"x-client-public-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: Checks if the version falls back to 0.1.0 when package metadata is missing..""" From 88f37ff9c44703dbf9f9d17132291395f16fe02c Mon Sep 17 00:00:00 2001 From: Gaurav Londhe Date: Thu, 10 Sep 2026 18:44:31 +0530 Subject: [PATCH 2/2] refactor: Remove unnecessary blank lines in test_connect_injects_all_required_logging_headers --- tests/test_logging_headers.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_logging_headers.py b/tests/test_logging_headers.py index 55a42e1..3a4ff79 100644 --- a/tests/test_logging_headers.py +++ b/tests/test_logging_headers.py @@ -36,8 +36,6 @@ def test_connect_injects_all_required_logging_headers(mock_flight_client: MagicM assert b"x-client-dataconnect" in headers_dict assert headers_dict[b"x-client-dataconnect"] == b"Python_SDK;v1.2.3-test;" - - # Check that existing auth token logic is preserved assert b"authorization" in headers_dict assert headers_dict[b"authorization"] == f"Bearer {test_token}".encode()