From 81a627e8b86184e390b24026390207507829c356 Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Fri, 16 Oct 2026 10:24:15 +0200 Subject: [PATCH 1/3] fix(client): detect dead connections with gRPC keepalive Co-authored-by: Cursor --- src/altertable_flightsql/client.py | 11 +++ tests/test_client.py | 21 +++++ tests/test_keepalive.py | 139 +++++++++++++++++++++++++++++ 3 files changed, 171 insertions(+) create mode 100644 tests/test_keepalive.py diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index cfdb818..043cd18 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -18,6 +18,8 @@ from altertable_flightsql.generated import arrow_flight_sql_pb2 as sql_pb2 GRPC_MAX_MESSAGE_LENGTH = 512 * 1024 * 1024 +GRPC_KEEPALIVE_TIME_MS = 30_000 +GRPC_KEEPALIVE_TIMEOUT_MS = 20_000 def _pack_command(cmd) -> bytes: @@ -167,6 +169,15 @@ def __init__( generic_options=[ ("grpc.max_send_message_length", GRPC_MAX_MESSAGE_LENGTH), ("grpc.max_receive_message_length", GRPC_MAX_MESSAGE_LENGTH), + ("grpc.keepalive_time_ms", GRPC_KEEPALIVE_TIME_MS), + ("grpc.keepalive_timeout_ms", GRPC_KEEPALIVE_TIMEOUT_MS), + # The gRPC bundled with pyarrow times keepalive pings out after ping_timeout_ms + # (default 60s), not keepalive_timeout_ms. + ("grpc.http2.ping_timeout_ms", GRPC_KEEPALIVE_TIMEOUT_MS), + ("grpc.keepalive_permit_without_calls", 1), + # The client sends no data while a statement runs, and gRPC otherwise stops + # pinging after two pings without data. + ("grpc.http2.max_pings_without_data", 0), ], ) self._client.authenticate_basic_token(self._username, self._password) diff --git a/tests/test_client.py b/tests/test_client.py index 328f9e0..8776ad9 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -154,6 +154,27 @@ def test_constructor_sends_compute_size(monkeypatch): assert request.session_options["catalog"].string_value == "c" +def test_constructor_enables_grpc_keepalive(monkeypatch): + flight_client = FakeFlightClient() + flight_client.authenticate_basic_token = lambda *a: None + flight_client_kwargs = {} + + def build_flight_client(*args, **kwargs): + flight_client_kwargs.update(kwargs) + return flight_client + + monkeypatch.setattr(flight, "FlightClient", build_flight_client) + + Client("u", "p") + + generic_options = flight_client_kwargs["generic_options"] + assert ("grpc.keepalive_time_ms", 30_000) in generic_options + assert ("grpc.keepalive_timeout_ms", 20_000) in generic_options + assert ("grpc.http2.ping_timeout_ms", 20_000) in generic_options + assert ("grpc.keepalive_permit_without_calls", 1) in generic_options + assert ("grpc.http2.max_pings_without_data", 0) in generic_options + + def test_close_closes_server_session_before_transport(): flight_client = FakeFlightClient() client = _client_backed_by(flight_client) diff --git a/tests/test_keepalive.py b/tests/test_keepalive.py new file mode 100644 index 0000000..b1c8927 --- /dev/null +++ b/tests/test_keepalive.py @@ -0,0 +1,139 @@ +import contextlib +import socket +import threading +import time + +import pyarrow.flight as flight +import pytest + +from altertable_flightsql import client as client_module +from altertable_flightsql.client import Client + +KEEPALIVE_MS = 1000 +# Long enough for the client to exceed gRPC's default budget of two pings without sending data. +HEALTHY_SECONDS = 5 + + +class BearerTokenServerMiddleware(flight.ServerMiddleware): + def sending_headers(self): + return {"authorization": "Bearer test-token"} + + +class BearerTokenServerMiddlewareFactory(flight.ServerMiddlewareFactory): + def start_call(self, info, headers): + return BearerTokenServerMiddleware() + + +class NoopAuthHandler(flight.ServerAuthHandler): + def authenticate(self, outgoing, incoming): + pass + + def is_valid(self, token): + return "" + + +class BlockingFlightServer(flight.FlightServerBase): + """Flight server whose GetFlightInfo blocks like a long-running statement.""" + + def __init__(self): + super().__init__( + "grpc://127.0.0.1:0", + auth_handler=NoopAuthHandler(), + middleware={"auth": BearerTokenServerMiddlewareFactory()}, + ) + self.statement_started = threading.Event() + self.release = threading.Event() + + def get_flight_info(self, context, descriptor): + self.statement_started.set() + self.release.wait(timeout=60) + raise flight.FlightCancelledError("released") + + +class FreezableProxy: + """TCP proxy that can silently stop forwarding bytes while keeping sockets open.""" + + def __init__(self, upstream_port: int): + self._upstream_port = upstream_port + self._frozen = threading.Event() + self._closed = threading.Event() + self._sockets = [] + self._listener = socket.create_server(("127.0.0.1", 0)) + self.port = self._listener.getsockname()[1] + threading.Thread(target=self._accept_loop, daemon=True).start() + + def freeze(self): + self._frozen.set() + + def close(self): + self._closed.set() + for sock in [self._listener, *self._sockets]: + with contextlib.suppress(OSError): + sock.close() + + def _accept_loop(self): + while not self._closed.is_set(): + try: + downstream, _ = self._listener.accept() + except OSError: + return + upstream = socket.create_connection(("127.0.0.1", self._upstream_port)) + self._sockets += [downstream, upstream] + threading.Thread(target=self._pump, args=(downstream, upstream), daemon=True).start() + threading.Thread(target=self._pump, args=(upstream, downstream), daemon=True).start() + + def _pump(self, source: socket.socket, destination: socket.socket): + while True: + try: + data = source.recv(65536) + if not data: + return + if not self._frozen.is_set(): + destination.sendall(data) + except OSError: + return + + +@pytest.fixture +def blocking_server(): + server = BlockingFlightServer() + yield server + server.release.set() + server.shutdown() + + +@pytest.fixture +def proxy(blocking_server): + proxy = FreezableProxy(blocking_server.port) + yield proxy + proxy.close() + + +def test_connection_cut_during_long_statement_fails_the_call(monkeypatch, blocking_server, proxy): + monkeypatch.setattr(client_module, "GRPC_KEEPALIVE_TIME_MS", KEEPALIVE_MS) + monkeypatch.setattr(client_module, "GRPC_KEEPALIVE_TIMEOUT_MS", KEEPALIVE_MS) + client = Client("user", "pass", host="127.0.0.1", port=proxy.port, tls=False) + errors = [] + + def run_statement(): + try: + client.query("SELECT long_running()") + except Exception as e: + errors.append(e) + + try: + statement = threading.Thread(target=run_statement, daemon=True) + statement.start() + assert blocking_server.statement_started.wait(timeout=10) + time.sleep(HEALTHY_SECONDS) + assert statement.is_alive(), f"statement failed on a healthy connection: {errors}" + + proxy.freeze() + statement.join(timeout=15) + + assert not statement.is_alive(), "statement hung on a dead connection" + assert len(errors) == 1 + assert isinstance(errors[0], flight.FlightUnavailableError) + finally: + with contextlib.suppress(Exception): + client.close(timeout_seconds=1) From 5e606c40264d41c9cdfce2260af68ddcc5313eb8 Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Fri, 16 Oct 2026 10:29:36 +0200 Subject: [PATCH 2/3] test: use a grpcio server that tolerates frequent pings in keepalive test Co-authored-by: Cursor --- pyproject.toml | 1 + src/altertable_flightsql/client.py | 4 +- tests/test_keepalive.py | 63 ++++++++++++++++++------------ 3 files changed, 40 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 244d82b..8937b2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ [project.optional-dependencies] dev = [ + "grpcio==1.80.0", "grpcio-tools==1.80.0", "pytest>=7.0.0", "pytest-cov>=3.0.0", diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index 043cd18..ff82ca2 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -171,8 +171,8 @@ def __init__( ("grpc.max_receive_message_length", GRPC_MAX_MESSAGE_LENGTH), ("grpc.keepalive_time_ms", GRPC_KEEPALIVE_TIME_MS), ("grpc.keepalive_timeout_ms", GRPC_KEEPALIVE_TIMEOUT_MS), - # The gRPC bundled with pyarrow times keepalive pings out after ping_timeout_ms - # (default 60s), not keepalive_timeout_ms. + # The gRPC bundled with recent pyarrow releases times keepalive pings out after + # ping_timeout_ms (default 60s), not keepalive_timeout_ms. ("grpc.http2.ping_timeout_ms", GRPC_KEEPALIVE_TIMEOUT_MS), ("grpc.keepalive_permit_without_calls", 1), # The client sends no data while a statement runs, and gRPC otherwise stops diff --git a/tests/test_keepalive.py b/tests/test_keepalive.py index b1c8927..1808414 100644 --- a/tests/test_keepalive.py +++ b/tests/test_keepalive.py @@ -2,7 +2,9 @@ import socket import threading import time +from concurrent import futures +import grpc import pyarrow.flight as flight import pytest @@ -14,40 +16,49 @@ HEALTHY_SECONDS = 5 -class BearerTokenServerMiddleware(flight.ServerMiddleware): - def sending_headers(self): - return {"authorization": "Bearer test-token"} +class BlockingFlightServer: + """Flight server whose GetFlightInfo blocks like a long-running statement. - -class BearerTokenServerMiddlewareFactory(flight.ServerMiddlewareFactory): - def start_call(self, info, headers): - return BearerTokenServerMiddleware() - - -class NoopAuthHandler(flight.ServerAuthHandler): - def authenticate(self, outgoing, incoming): - pass - - def is_valid(self, token): - return "" - - -class BlockingFlightServer(flight.FlightServerBase): - """Flight server whose GetFlightInfo blocks like a long-running statement.""" + Built on grpcio rather than pyarrow so it can accept frequent client pings like the + production server; some pyarrow servers reject them with a "too_many_pings" GOAWAY. + """ def __init__(self): - super().__init__( - "grpc://127.0.0.1:0", - auth_handler=NoopAuthHandler(), - middleware={"auth": BearerTokenServerMiddlewareFactory()}, - ) self.statement_started = threading.Event() self.release = threading.Event() + self._server = grpc.server( + futures.ThreadPoolExecutor(max_workers=4), + options=[("grpc.http2.max_ping_strikes", 0)], + ) + self._server.add_generic_rpc_handlers( + ( + grpc.method_handlers_generic_handler( + "arrow.flight.protocol.FlightService", + { + "Handshake": grpc.stream_stream_rpc_method_handler(self._handshake), + "GetFlightInfo": grpc.unary_unary_rpc_method_handler( + self._get_flight_info + ), + }, + ), + ) + ) + self.port = self._server.add_insecure_port("127.0.0.1:0") + self._server.start() + + def shutdown(self): + self._server.stop(grace=None) + + def _handshake(self, requests, context): + for _ in requests: + pass + context.send_initial_metadata((("authorization", "Bearer test-token"),)) + return iter(()) - def get_flight_info(self, context, descriptor): + def _get_flight_info(self, request, context): self.statement_started.set() self.release.wait(timeout=60) - raise flight.FlightCancelledError("released") + context.abort(grpc.StatusCode.CANCELLED, "released") class FreezableProxy: From bd4da2f22f20aec7602f2339c5905512f932cf74 Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Fri, 16 Oct 2026 10:31:15 +0200 Subject: [PATCH 3/3] style: keep keepalive test stable across black versions Co-authored-by: Cursor --- tests/test_keepalive.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/tests/test_keepalive.py b/tests/test_keepalive.py index 1808414..a0fff10 100644 --- a/tests/test_keepalive.py +++ b/tests/test_keepalive.py @@ -30,18 +30,12 @@ def __init__(self): futures.ThreadPoolExecutor(max_workers=4), options=[("grpc.http2.max_ping_strikes", 0)], ) + handlers = { + "Handshake": grpc.stream_stream_rpc_method_handler(self._handshake), + "GetFlightInfo": grpc.unary_unary_rpc_method_handler(self._get_flight_info), + } self._server.add_generic_rpc_handlers( - ( - grpc.method_handlers_generic_handler( - "arrow.flight.protocol.FlightService", - { - "Handshake": grpc.stream_stream_rpc_method_handler(self._handshake), - "GetFlightInfo": grpc.unary_unary_rpc_method_handler( - self._get_flight_info - ), - }, - ), - ) + (grpc.method_handlers_generic_handler("arrow.flight.protocol.FlightService", handlers),) ) self.port = self._server.add_insecure_port("127.0.0.1:0") self._server.start()