From d654a74f87034cfb2ec394c9d9b24759a8267e8b Mon Sep 17 00:00:00 2001 From: Florian Valeye Date: Fri, 24 Jul 2026 11:46:12 +0200 Subject: [PATCH 1/5] fix: close Flight sessions before transport --- src/altertable_flightsql/client.py | 23 ++++++-- tests/test_client.py | 86 +++++++++++++++++++++++++++++- 2 files changed, 104 insertions(+), 5 deletions(-) diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index 2adad99..70d5767 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -549,9 +549,26 @@ def _end_transaction(self, transaction: "Transaction", commit: bool) -> None: action = flight.Action("EndTransaction", _pack_command(request)) list(self._client.do_action(action)) - def close(self) -> None: - """Close the client connection.""" - self._client.close() + def close(self, timeout_seconds: float = 10.0) -> None: + """Close the server session and client transport. Idempotent.""" + if getattr(self, "_closed", False): + return + self._closed = True + request = flight_pb2.CloseSessionRequest() + action = flight.Action("CloseSession", request.SerializeToString()) + options = flight.FlightCallOptions(timeout=timeout_seconds) + try: + result = next(iter(self._client.do_action(action, options)), None) + if result is not None: + close_result = flight_pb2.CloseSessionResult.FromString(bytes(result.body)) + if close_result.status not in ( + flight_pb2.CloseSessionResult.CLOSED, + flight_pb2.CloseSessionResult.CLOSING, + ): + status = flight_pb2.CloseSessionResult.Status.Name(close_result.status) + raise RuntimeError(f"Server did not close Flight session: {status}") + finally: + self._client.close() def __enter__(self) -> "Client": """Context manager entry.""" diff --git a/tests/test_client.py b/tests/test_client.py index fae3f3e..ba2e5e8 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,3 +1,7 @@ +from types import SimpleNamespace + +import pyarrow.flight as flight +import pytest from google.protobuf import any_pb2 from altertable_flightsql.client import Client @@ -7,10 +11,29 @@ class FakeFlightClient: def __init__(self): self.actions = [] + self.options = [] + self.events = [] + self.closed = False + close_result = flight_pb2.CloseSessionResult(status=flight_pb2.CloseSessionResult.CLOSED) + self.action_results = [SimpleNamespace(body=close_result.SerializeToString())] - def do_action(self, action): + def do_action(self, action, options=None): + if self.closed: + raise RuntimeError("FlightClient is closed") self.actions.append(action) - return [] + self.options.append(options) + self.events.append(("action", action.type)) + return self.action_results + + def close(self): + self.closed = True + self.events.append(("close", None)) + + +class FailingCloseSessionFlightClient(FakeFlightClient): + def do_action(self, action, options=None): + super().do_action(action, options) + raise RuntimeError("close session failed") def _action_body_bytes(action) -> bytes: @@ -38,3 +61,62 @@ def test_set_options_serializes_flight_session_request_without_any(): assert action.type == "SetSessionOptions" assert _action_body_bytes(action) == request.SerializeToString() assert _action_body_bytes(action) != wrapped_request.SerializeToString() + + +def test_close_closes_server_session_before_transport(): + flight_client = FakeFlightClient() + client = Client.__new__(Client) + client._client = flight_client + + client.close() + + action = flight_client.actions[0] + request = flight_pb2.CloseSessionRequest() + + assert flight_client.events == [("action", "CloseSession"), ("close", None)] + assert _action_body_bytes(action) == request.SerializeToString() + + +def test_close_closes_transport_when_server_session_close_fails(): + flight_client = FailingCloseSessionFlightClient() + client = Client.__new__(Client) + client._client = flight_client + + with pytest.raises(RuntimeError, match="close session failed"): + client.close() + + assert flight_client.events == [("action", "CloseSession"), ("close", None)] + + +def test_close_rejects_unclosed_server_session(): + flight_client = FakeFlightClient() + close_result = flight_pb2.CloseSessionResult(status=flight_pb2.CloseSessionResult.NOT_CLOSEABLE) + flight_client.action_results = [SimpleNamespace(body=close_result.SerializeToString())] + client = Client.__new__(Client) + client._client = flight_client + + with pytest.raises(RuntimeError, match="NOT_CLOSEABLE"): + client.close() + + assert flight_client.events == [("action", "CloseSession"), ("close", None)] + + +def test_close_is_idempotent(): + flight_client = FakeFlightClient() + client = Client.__new__(Client) + client._client = flight_client + + client.close() + client.close() + + assert flight_client.events == [("action", "CloseSession"), ("close", None)] + + +def test_close_passes_bounded_timeout_to_do_action(): + flight_client = FakeFlightClient() + client = Client.__new__(Client) + client._client = flight_client + + client.close() + + assert isinstance(flight_client.options[0], flight.FlightCallOptions) From 4344d765b551e80a9c9c85808cc68e0a04cb7420 Mon Sep 17 00:00:00 2001 From: Florian Valeye Date: Sat, 25 Jul 2026 15:42:38 +0200 Subject: [PATCH 2/5] build: pin grpcio-tools for reproducible proto codegen --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 55aef88..244d82b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ [project.optional-dependencies] dev = [ - "grpcio-tools>=1.73.0", + "grpcio-tools==1.80.0", "pytest>=7.0.0", "pytest-cov>=3.0.0", "black>=22.0.0", diff --git a/uv.lock b/uv.lock index aaedd47..d32c456 100644 --- a/uv.lock +++ b/uv.lock @@ -43,7 +43,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "black", marker = "extra == 'dev'", specifier = ">=22.0.0" }, - { name = "grpcio-tools", marker = "extra == 'dev'", specifier = ">=1.73.0" }, + { name = "grpcio-tools", marker = "extra == 'dev'", specifier = "==1.80.0" }, { name = "isort", marker = "extra == 'dev'", specifier = ">=5.10.0" }, { name = "pandas", marker = "extra == 'dev'", specifier = ">=1.3.0" }, { name = "protobuf", specifier = ">=5.26.0" }, From 4219a36a127e0da56b41bb7b776424fc780a0595 Mon Sep 17 00:00:00 2001 From: Florian Valeye Date: Sat, 25 Jul 2026 16:42:17 +0200 Subject: [PATCH 3/5] refactor: initialize the closed flag in the constructor --- src/altertable_flightsql/client.py | 3 ++- tests/test_client.py | 25 +++++++++++++------------ 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index 70d5767..e01deef 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -142,6 +142,7 @@ def __init__( self._password = password self._auto_commit = auto_commit self._transaction = None + self._closed = False auth_middleware = BearerAuthMiddlewareFactory() self._client = flight.FlightClient( @@ -551,7 +552,7 @@ def _end_transaction(self, transaction: "Transaction", commit: bool) -> None: def close(self, timeout_seconds: float = 10.0) -> None: """Close the server session and client transport. Idempotent.""" - if getattr(self, "_closed", False): + if self._closed: return self._closed = True request = flight_pb2.CloseSessionRequest() diff --git a/tests/test_client.py b/tests/test_client.py index ba2e5e8..8f8f8fe 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -36,6 +36,13 @@ def do_action(self, action, options=None): raise RuntimeError("close session failed") +def _client_backed_by(flight_client) -> Client: + client = Client.__new__(Client) + client._client = flight_client + client._closed = False + return client + + def _action_body_bytes(action) -> bytes: body = action.body if hasattr(body, "to_pybytes"): @@ -45,8 +52,7 @@ def _action_body_bytes(action) -> bytes: def test_set_options_serializes_flight_session_request_without_any(): flight_client = FakeFlightClient() - client = Client.__new__(Client) - client._client = flight_client + client = _client_backed_by(flight_client) session_options = { "catalog": flight_pb2.SessionOptionValue(string_value="test_catalog"), @@ -65,8 +71,7 @@ def test_set_options_serializes_flight_session_request_without_any(): def test_close_closes_server_session_before_transport(): flight_client = FakeFlightClient() - client = Client.__new__(Client) - client._client = flight_client + client = _client_backed_by(flight_client) client.close() @@ -79,8 +84,7 @@ def test_close_closes_server_session_before_transport(): def test_close_closes_transport_when_server_session_close_fails(): flight_client = FailingCloseSessionFlightClient() - client = Client.__new__(Client) - client._client = flight_client + client = _client_backed_by(flight_client) with pytest.raises(RuntimeError, match="close session failed"): client.close() @@ -92,8 +96,7 @@ def test_close_rejects_unclosed_server_session(): flight_client = FakeFlightClient() close_result = flight_pb2.CloseSessionResult(status=flight_pb2.CloseSessionResult.NOT_CLOSEABLE) flight_client.action_results = [SimpleNamespace(body=close_result.SerializeToString())] - client = Client.__new__(Client) - client._client = flight_client + client = _client_backed_by(flight_client) with pytest.raises(RuntimeError, match="NOT_CLOSEABLE"): client.close() @@ -103,8 +106,7 @@ def test_close_rejects_unclosed_server_session(): def test_close_is_idempotent(): flight_client = FakeFlightClient() - client = Client.__new__(Client) - client._client = flight_client + client = _client_backed_by(flight_client) client.close() client.close() @@ -114,8 +116,7 @@ def test_close_is_idempotent(): def test_close_passes_bounded_timeout_to_do_action(): flight_client = FakeFlightClient() - client = Client.__new__(Client) - client._client = flight_client + client = _client_backed_by(flight_client) client.close() From 4b27286048b5245cc413ff73ff121b95e3979550 Mon Sep 17 00:00:00 2001 From: Florian Valeye Date: Sat, 25 Jul 2026 20:42:38 +0200 Subject: [PATCH 4/5] fix: surface Flight errors raised after the first CloseSession result --- src/altertable_flightsql/client.py | 19 +++++----- tests/test_client.py | 56 ++++++++++++++++++++++++++++-- 2 files changed, 63 insertions(+), 12 deletions(-) diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index e01deef..dc50837 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -559,15 +559,16 @@ def close(self, timeout_seconds: float = 10.0) -> None: action = flight.Action("CloseSession", request.SerializeToString()) options = flight.FlightCallOptions(timeout=timeout_seconds) try: - result = next(iter(self._client.do_action(action, options)), None) - if result is not None: - close_result = flight_pb2.CloseSessionResult.FromString(bytes(result.body)) - if close_result.status not in ( - flight_pb2.CloseSessionResult.CLOSED, - flight_pb2.CloseSessionResult.CLOSING, - ): - status = flight_pb2.CloseSessionResult.Status.Name(close_result.status) - raise RuntimeError(f"Server did not close Flight session: {status}") + results = list(self._client.do_action(action, options)) + if not results: + raise RuntimeError("Server returned no CloseSessionResult") + close_result = flight_pb2.CloseSessionResult.FromString(bytes(results[0].body)) + if close_result.status not in ( + flight_pb2.CloseSessionResult.CLOSED, + flight_pb2.CloseSessionResult.CLOSING, + ): + status = flight_pb2.CloseSessionResult.Status.Name(close_result.status) + raise RuntimeError(f"Server did not close Flight session: {status}") finally: self._client.close() diff --git a/tests/test_client.py b/tests/test_client.py index 8f8f8fe..a3efd6e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -36,6 +36,25 @@ def do_action(self, action, options=None): raise RuntimeError("close session failed") +def _record_call_option_timeouts(monkeypatch) -> list: + """`FlightCallOptions.timeout` is unreadable on older PyArrow, so record construction.""" + timeouts = [] + build_options = flight.FlightCallOptions + + def recording_options(**kwargs): + timeouts.append(kwargs.get("timeout")) + return build_options(**kwargs) + + monkeypatch.setattr(flight, "FlightCallOptions", recording_options) + return timeouts + + +def _closed_result_then_error(): + close_result = flight_pb2.CloseSessionResult(status=flight_pb2.CloseSessionResult.CLOSED) + yield SimpleNamespace(body=close_result.SerializeToString()) + raise RuntimeError("server failed after the first result") + + def _client_backed_by(flight_client) -> Client: client = Client.__new__(Client) client._client = flight_client @@ -114,10 +133,41 @@ def test_close_is_idempotent(): assert flight_client.events == [("action", "CloseSession"), ("close", None)] -def test_close_passes_bounded_timeout_to_do_action(): +def test_close_passes_bounded_timeout_to_do_action(monkeypatch): + timeouts = _record_call_option_timeouts(monkeypatch) + client = _client_backed_by(FakeFlightClient()) + + client.close() + + assert timeouts == [10.0] + + +def test_close_passes_caller_timeout_to_do_action(monkeypatch): + timeouts = _record_call_option_timeouts(monkeypatch) + client = _client_backed_by(FakeFlightClient()) + + client.close(timeout_seconds=2.5) + + assert timeouts == [2.5] + + +def test_close_surfaces_server_error_raised_after_the_first_result(): + flight_client = FakeFlightClient() + flight_client.action_results = _closed_result_then_error() + client = _client_backed_by(flight_client) + + with pytest.raises(RuntimeError, match="server failed after the first result"): + client.close() + + assert flight_client.events == [("action", "CloseSession"), ("close", None)] + + +def test_close_rejects_empty_server_response(): flight_client = FakeFlightClient() + flight_client.action_results = [] client = _client_backed_by(flight_client) - client.close() + with pytest.raises(RuntimeError, match="no CloseSessionResult"): + client.close() - assert isinstance(flight_client.options[0], flight.FlightCallOptions) + assert flight_client.events == [("action", "CloseSession"), ("close", None)] From f5e265b80db40e9f5c82163cbbf2504d38e9757e Mon Sep 17 00:00:00 2001 From: Florian Valeye Date: Sat, 25 Jul 2026 21:09:38 +0200 Subject: [PATCH 5/5] fix: reject a Flight session the server is still closing --- src/altertable_flightsql/client.py | 5 +---- tests/test_client.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index dc50837..a5b61ac 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -563,10 +563,7 @@ def close(self, timeout_seconds: float = 10.0) -> None: if not results: raise RuntimeError("Server returned no CloseSessionResult") close_result = flight_pb2.CloseSessionResult.FromString(bytes(results[0].body)) - if close_result.status not in ( - flight_pb2.CloseSessionResult.CLOSED, - flight_pb2.CloseSessionResult.CLOSING, - ): + if close_result.status != flight_pb2.CloseSessionResult.CLOSED: status = flight_pb2.CloseSessionResult.Status.Name(close_result.status) raise RuntimeError(f"Server did not close Flight session: {status}") finally: diff --git a/tests/test_client.py b/tests/test_client.py index a3efd6e..b0ebd02 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -123,6 +123,18 @@ def test_close_rejects_unclosed_server_session(): assert flight_client.events == [("action", "CloseSession"), ("close", None)] +def test_close_rejects_a_session_the_server_is_still_closing(): + flight_client = FakeFlightClient() + close_result = flight_pb2.CloseSessionResult(status=flight_pb2.CloseSessionResult.CLOSING) + flight_client.action_results = [SimpleNamespace(body=close_result.SerializeToString())] + client = _client_backed_by(flight_client) + + with pytest.raises(RuntimeError, match="CLOSING"): + client.close() + + assert flight_client.events == [("action", "CloseSession"), ("close", None)] + + def test_close_is_idempotent(): flight_client = FakeFlightClient() client = _client_backed_by(flight_client)