From 3b522963f7c68a50b5374750b33546647d440142 Mon Sep 17 00:00:00 2001 From: Ilya Soin Date: Thu, 19 Mar 2026 16:32:38 +0200 Subject: [PATCH 1/2] NOTASK add timeouts --- src/flink_gateway/__init__.py | 2 + src/flink_gateway/connection.py | 26 +++- src/flink_gateway/cursor.py | 59 ++++++++-- src/flink_gateway/exceptions.py | 4 + tests/test_dbapi.py | 202 +++++++++++++++++++++++++++++++- 5 files changed, 277 insertions(+), 16 deletions(-) diff --git a/src/flink_gateway/__init__.py b/src/flink_gateway/__init__.py index d0efee5..1f52984 100644 --- a/src/flink_gateway/__init__.py +++ b/src/flink_gateway/__init__.py @@ -14,6 +14,7 @@ NotSupportedError, OperationalError, ProgrammingError, + TimeoutError, ) from flink_gateway.models import ( ColumnInfo, @@ -65,6 +66,7 @@ "InterfaceError", "DatabaseError", "OperationalError", + "TimeoutError", "FlinkSqlGatewayError", "ProgrammingError", "NotSupportedError", diff --git a/src/flink_gateway/connection.py b/src/flink_gateway/connection.py index e88421e..57500ee 100644 --- a/src/flink_gateway/connection.py +++ b/src/flink_gateway/connection.py @@ -23,10 +23,14 @@ def __init__( client: FlinkSqlGatewayClient, session_handle: str, *, + query_timeout: float | None = None, + idle_timeout: float | None = None, _owns_client: bool = True, ) -> None: self._client = client self._session_handle = session_handle + self._query_timeout = query_timeout + self._idle_timeout = idle_timeout self._owns_client = _owns_client self._closed = False @@ -79,7 +83,11 @@ def cursor(self) -> Cursor: """Create a new Cursor bound to this connection.""" if self._closed: raise ProgrammingError("connection is closed") - return Cursor(self) + return Cursor( + self, + query_timeout=self._query_timeout, + idle_timeout=self._idle_timeout, + ) def connect( @@ -88,6 +96,8 @@ def connect( properties: dict[str, str] | None = None, api_version: str = "v3", http_client: httpx.Client | None = None, + query_timeout: float | None = None, + idle_timeout: float | None = None, ) -> Connection: """Open a connection to a Flink SQL Gateway. @@ -99,6 +109,12 @@ def connect( api_version: REST API version (default ``"v3"``). http_client: Optional pre-configured :class:`httpx.Client` for custom SSL, timeouts, authentication, etc. + query_timeout: Maximum seconds a query may run + before raising :class:`~flink_gateway.TimeoutError`. + ``None`` means no limit. + idle_timeout: Maximum seconds to wait between rows during streaming + iteration before raising :class:`~flink_gateway.TimeoutError`. + ``None`` means no limit. Returns: A :class:`Connection` instance. @@ -113,4 +129,10 @@ def connect( except Exception: client.close() raise - return Connection(client, session_handle, _owns_client=True) + return Connection( + client, + session_handle, + query_timeout=query_timeout, + idle_timeout=idle_timeout, + _owns_client=True, + ) diff --git a/src/flink_gateway/cursor.py b/src/flink_gateway/cursor.py index a9b1d40..d7e4b42 100644 --- a/src/flink_gateway/cursor.py +++ b/src/flink_gateway/cursor.py @@ -5,7 +5,7 @@ import time from typing import TYPE_CHECKING, Any, Iterator, Sequence -from flink_gateway.exceptions import InterfaceError, OperationalError, ProgrammingError +from flink_gateway.exceptions import InterfaceError, ProgrammingError, TimeoutError from flink_gateway.models import ( ColumnInfo, ExecuteStatementRequest, @@ -22,7 +22,6 @@ _POLL_INTERVAL = 1.0 # seconds _RESULTS_MIN_BACKOFF = 0.1 # seconds _RESULTS_MAX_BACKOFF = 1.0 # seconds -_DEFAULT_QUERY_TIMEOUT = 300.0 # 5 minutes # PEP 249 description row: (name, type_code, display_size, internal_size, # precision, scale, null_ok) @@ -38,7 +37,11 @@ class Cursor: arraysize: int = 1 def __init__( - self, connection: Connection, *, query_timeout: float = _DEFAULT_QUERY_TIMEOUT + self, + connection: Connection, + *, + query_timeout: float | None = None, + idle_timeout: float | None = None, ) -> None: self._connection = connection self._closed = False @@ -48,6 +51,8 @@ def __init__( self._rows_iterator: Iterator[tuple[Any, ...]] | None = None self._columns: list[ColumnInfo] = [] self._query_timeout = query_timeout + self._idle_timeout = idle_timeout + self._query_deadline: float = float("inf") # ── PEP 249 attributes ───────────────────────────────────────── @@ -91,7 +96,9 @@ def _check_open(self) -> None: # ── Execution ────────────────────────────────────────────────── def execute( - self, operation: str, parameters: Sequence[Any] | None = None + self, + operation: str, + parameters: Sequence[Any] | None = None, ) -> Cursor: """Execute a SQL statement. @@ -112,6 +119,11 @@ def execute( self._rowcount = -1 self._rows_iterator = None self._columns = [] + self._query_deadline = ( + time.monotonic() + self._query_timeout + if self._query_timeout is not None + else float("inf") + ) if parameters is not None: raise InterfaceError("parameterized queries are not yet supported") @@ -228,15 +240,12 @@ def _fetch_until_ready(self, op_handle: str) -> FetchResultsResponse: """Poll ``fetch_results`` until the result is no longer NOT_READY.""" client = self._connection.client session = self._connection.session_handle - deadline = time.monotonic() + self._query_timeout while True: result = client.fetch_results(session, op_handle, "0", "json") if result.result_type != ResultType.NOT_READY: return result - if time.monotonic() > deadline: - raise OperationalError( - f"query timed out after {self._query_timeout}s waiting for results" - ) + if time.monotonic() > self._query_deadline: + raise TimeoutError(f"query timed out after {self._query_timeout}s") time.sleep(_POLL_INTERVAL) def _build_description(self) -> None: @@ -271,6 +280,12 @@ def _iter_rows( client = self._connection.client session = self._connection.session_handle + idle_deadline: float = ( + time.monotonic() + self._idle_timeout + if self._idle_timeout is not None + else float("inf") + ) + data = initial_data token = next_token pos = 0 @@ -281,6 +296,11 @@ def _iter_rows( row = data[pos] pos += 1 yield self._decode_row(row) + if time.monotonic() >= self._query_deadline: + raise TimeoutError( + f"streaming iteration stopped: query_timeout of " + f"{self._query_timeout}s exceeded" + ) continue # Fetch next page. @@ -295,11 +315,26 @@ def _iter_rows( pos = 0 if not data: - time.sleep(backoff) - if backoff < _RESULTS_MAX_BACKOFF: - backoff = min(backoff * 2, _RESULTS_MAX_BACKOFF) + now = time.monotonic() + if now >= self._query_deadline: + raise TimeoutError( + f"streaming iteration stopped: query_timeout of " + f"{self._query_timeout}s exceeded" + ) + if now >= idle_deadline: + raise TimeoutError( + f"streaming iteration stopped: idle_timeout of " + f"{self._idle_timeout}s exceeded" + ) + sleep_dur = min( + backoff, self._query_deadline - now, idle_deadline - now + ) + time.sleep(sleep_dur) + backoff = min(backoff * 2, _RESULTS_MAX_BACKOFF) else: backoff = _RESULTS_MIN_BACKOFF + if self._idle_timeout is not None: + idle_deadline = time.monotonic() + self._idle_timeout def _decode_row(self, row: RowData) -> tuple[Any, ...]: """Decode a single RowData into a Python tuple.""" diff --git a/src/flink_gateway/exceptions.py b/src/flink_gateway/exceptions.py index 97f4b7d..2bb284b 100644 --- a/src/flink_gateway/exceptions.py +++ b/src/flink_gateway/exceptions.py @@ -28,6 +28,10 @@ class FlinkSqlGatewayError(OperationalError): """Raised when a Flink SQL Gateway REST call fails.""" +class TimeoutError(OperationalError): + """Raised when a query or streaming iteration exceeds a configured timeout.""" + + class ProgrammingError(DatabaseError): """Exception for programming errors. diff --git a/tests/test_dbapi.py b/tests/test_dbapi.py index ef82634..645673e 100644 --- a/tests/test_dbapi.py +++ b/tests/test_dbapi.py @@ -15,6 +15,7 @@ import flink_gateway from flink_gateway import Connection, Cursor, NotSupportedError, ProgrammingError from flink_gateway.connection import connect +from flink_gateway.exceptions import TimeoutError from flink_gateway.models import ( ColumnInfo, FetchResultsResponse, @@ -44,10 +45,21 @@ def _mock_client() -> MagicMock: return client -def _make_connection(client: MagicMock | None = None) -> Connection: +def _make_connection( + client: MagicMock | None = None, + *, + query_timeout: float | None = None, + idle_timeout: float | None = None, +) -> Connection: """Create a Connection backed by a mock client.""" c = client or _mock_client() - return Connection(c, "test-session", _owns_client=False) + return Connection( + c, + "test-session", + query_timeout=query_timeout, + idle_timeout=idle_timeout, + _owns_client=False, + ) def _columns(*specs: tuple[str, str, bool]) -> list[ColumnInfo]: @@ -689,3 +701,189 @@ def test_connect_cleans_up_on_failure(self): connect("http://localhost:8083") mock_instance.close.assert_called_once() + + +# ── Streaming timeouts ───────────────────────────────────────────────── +def _empty_payload(cols, next_uri: str) -> FetchResultsResponse: + """An empty PAYLOAD page with a pagination token (server has no rows yet).""" + return FetchResultsResponse( + result_type=ResultType.PAYLOAD, + next_result_uri=next_uri, + results=ResultSet(columns=cols, data=[]), + ) + + +class TestStreamingTimeouts: + """Tests for query_timeout and idle_timeout on connect(). + All tests mock time.monotonic and time.sleep to avoid real waits. + """ + + def test_query_timeout_fires_during_empty_pages(self): + """query_timeout raises TimeoutError when empty pages keep arriving.""" + client = _mock_client() + cols = _columns(("id", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + # Initial page has data + token; subsequent pages are empty but keep token + # so iteration continues until the timeout fires. + client.fetch_results.side_effect = [ + _ready_result(cols, [[1]], next_uri="/v3/.../result/1"), + _empty_payload(cols, "/v3/.../result/2"), + _empty_payload(cols, "/v3/.../result/3"), + _empty_payload(cols, "/v3/.../result/4"), + ] + + # Monotonic call map (query_timeout=1.0, no idle_timeout): + # 0.0 → execute(): _query_deadline = 1.0 + # 0.0 → post-yield check after row 1 (ok) + # 0.5 → empty-page 1 timeout check (ok, sleep capped to 0.5) + # 0.5 → empty-page 2 timeout check (ok) + # 1.5 → empty-page 3 timeout check → RAISE + monotonic_seq = iter([0.0, 0.0, 0.5, 0.5, 1.5]) + + conn = _make_connection(client, query_timeout=1.0) + cur = conn.cursor() + with ( + patch("flink_gateway.cursor.time.sleep"), + patch("flink_gateway.cursor.time.monotonic", side_effect=monotonic_seq), + ): + cur.execute("SELECT id FROM t") + with pytest.raises(TimeoutError, match="query_timeout"): + cur.fetchall() + + def test_idle_timeout_fires_during_empty_pages(self): + """idle_timeout raises TimeoutError when no rows arrive for too long.""" + client = _mock_client() + cols = _columns(("id", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + client.fetch_results.side_effect = [ + _ready_result(cols, [[1]], next_uri="/v3/.../result/1"), + _empty_payload(cols, "/v3/.../result/2"), + _empty_payload(cols, "/v3/.../result/3"), + ] + + # Monotonic call map (idle_timeout=5.0, no query_timeout): + # 0.0 → _iter_rows: idle_deadline = 5.0 + # 0.0 → post-yield row 1 → ok + # 2.0 → empty-page 1 check → 2.0 < 5.0 → ok + # 6.0 → empty-page 2 check → 6.0 >= 5.0 → RAISE + monotonic_seq = iter([0.0, 0.0, 2.0, 6.0]) + + conn = _make_connection(client, idle_timeout=5.0) + cur = conn.cursor() + with ( + patch("flink_gateway.cursor.time.sleep"), + patch("flink_gateway.cursor.time.monotonic", side_effect=monotonic_seq), + ): + cur.execute("SELECT id FROM t") + with pytest.raises(TimeoutError, match="idle_timeout"): + cur.fetchall() + + def test_idle_deadline_resets_on_new_data(self): + """idle_deadline resets each time a non-empty page arrives; no timeout + raised.""" + client = _mock_client() + cols = _columns(("id", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + client.fetch_results.side_effect = [ + _ready_result(cols, [[1]], next_uri="/v3/.../result/1"), + _empty_payload(cols, "/v3/.../result/2"), + _ready_result( + cols, [[2]], next_uri="/v3/.../result/3" + ), # new data → reset idle + _eos_result(), + ] + + # Monotonic call map (idle_timeout=5.0, no query_timeout): + # 0.0 → _iter_rows: idle_deadline = 5.0 + # 0.0 → post-yield row 1 → ok + # 1.0 → empty-page check → 1.0 < 5.0 → ok + # 2.0 → non-empty page → idle_deadline reset to 7.0 + # 2.0 → post-yield row 2 → ok + monotonic_seq = iter([0.0, 0.0, 1.0, 2.0, 2.0]) + + conn = _make_connection(client, idle_timeout=5.0) + cur = conn.cursor() + with ( + patch("flink_gateway.cursor.time.sleep"), + patch("flink_gateway.cursor.time.monotonic", side_effect=monotonic_seq), + ): + cur.execute("SELECT id FROM t") + rows = cur.fetchall() + + assert rows == [(1,), (2,)] + + def test_query_timeout_fires_while_yielding_rows(self): + """query_timeout fires even when rows are arriving continuously.""" + client = _mock_client() + cols = _columns(("id", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + client.fetch_results.side_effect = [ + _ready_result(cols, [[1], [2], [3]], next_uri="/v3/.../result/1"), + _ready_result(cols, [[4], [5]], next_uri="/v3/.../result/2"), + ] + + # Monotonic call map (query_timeout=2.0, no idle_timeout): + # 0.0 → execute(): _query_deadline = 2.0 + # 0.5 → post-yield row 1 → ok + # 1.0 → post-yield row 2 → ok + # 2.5 → post-yield row 3 → >= 2.0 → RAISE + monotonic_seq = iter([0.0, 0.5, 1.0, 2.5]) + + conn = _make_connection(client, query_timeout=2.0) + cur = conn.cursor() + with ( + patch("flink_gateway.cursor.time.sleep"), + patch("flink_gateway.cursor.time.monotonic", side_effect=monotonic_seq), + ): + cur.execute("SELECT id FROM t") + with pytest.raises(TimeoutError, match="query_timeout"): + cur.fetchall() + + def test_sleep_is_capped_to_nearest_deadline(self): + """time.sleep is capped so it does not overshoot the idle deadline.""" + client = _mock_client() + cols = _columns(("id", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + client.fetch_results.side_effect = [ + _ready_result(cols, [[1]], next_uri="/v3/.../result/1"), + _empty_payload(cols, "/v3/.../result/2"), + _empty_payload(cols, "/v3/.../result/3"), + ] + + # Monotonic call map (idle_timeout=0.3, no query_timeout): + # 0.0 → _iter_rows: idle_deadline = 0.3 + # 0.0 → post-yield row 1 → ok + # 0.25 → empty-page 1: idle remaining = 0.05; backoff=0.1 → sleep(0.05) + # 0.35 → empty-page 2: 0.35 >= 0.3 → RAISE + monotonic_seq = iter([0.0, 0.0, 0.25, 0.35]) + + conn = _make_connection(client, idle_timeout=0.3) + cur = conn.cursor() + sleep_calls: list[float] = [] + + with ( + patch("flink_gateway.cursor.time.sleep", side_effect=sleep_calls.append), + patch("flink_gateway.cursor.time.monotonic", side_effect=monotonic_seq), + ): + cur.execute("SELECT id FROM t") + with pytest.raises(TimeoutError, match="idle_timeout"): + cur.fetchall() + + assert sleep_calls, "expected at least one sleep call" + assert sleep_calls[0] == pytest.approx(0.05) + + def test_no_timeout_defaults_unchanged(self): + """With no timeouts set, behavior is identical to the original.""" + client = _mock_client() + cols = _columns(("v", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + client.fetch_results.side_effect = [ + _ready_result(cols, [[1], [2]], next_uri="/v3/.../result/1"), + _ready_result(cols, [[3], [4]]), + _eos_result(), + ] + + conn = _make_connection(client) + cur = conn.cursor() + cur.execute("SELECT v FROM t") + assert cur.fetchall() == [(1,), (2,), (3,), (4,)] From 9e158cdf46f69c558663906a85e740bbdaed8c19 Mon Sep 17 00:00:00 2001 From: Ilya Soin Date: Thu, 19 Mar 2026 16:42:50 +0200 Subject: [PATCH 2/2] NOTASK update readme --- README.md | 35 ++++++++++++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 668036f..771f3da 100644 --- a/README.md +++ b/README.md @@ -77,10 +77,39 @@ import httpx client = httpx.Client(verify="/path/to/ca.crt", timeout=60.0) with connect("https://localhost:8083", http_client=client) as conn: ... +``` + +### Timeouts + +Both timeouts are `None` by default — no implicit interruptions. + +```python +from flink_gateway import connect, TimeoutError -# Set a query timeout (default: 300s) -with conn.cursor(query_timeout=60.0) as cur: - cur.execute("SELECT ...") +# query_timeout: hard wall-clock limit on the entire query (including streaming iteration). +# Raise TimeoutError once the deadline is exceeded, regardless of whether rows are arriving. +with connect("http://localhost:8083", query_timeout=60.0) as conn: + with conn.cursor() as cur: + cur.execute("SELECT * FROM orders") + for row in cur: + process(row) + +# idle_timeout: raise TimeoutError if no new rows arrive within the given window. +# Useful for streaming queries where silence means the source is done or stuck. +with connect("http://localhost:8083", idle_timeout=30.0) as conn: + with conn.cursor() as cur: + count = 0 + try: + cur.execute("SELECT * FROM orders") + for row in cur: + process(row) + count += 1 + except TimeoutError: + print(f"no new rows for 30s, processed {count} rows so far") + +# Both can be combined: stop as soon as either limit is hit. +with connect("http://localhost:8083", query_timeout=300.0, idle_timeout=30.0) as conn: + ... ``` ### Low-level REST access