diff --git a/src/flink_gateway/cursor.py b/src/flink_gateway/cursor.py index d7e4b42..07a8d9a 100644 --- a/src/flink_gateway/cursor.py +++ b/src/flink_gateway/cursor.py @@ -2,8 +2,9 @@ from __future__ import annotations +import enum import time -from typing import TYPE_CHECKING, Any, Iterator, Sequence +from typing import TYPE_CHECKING, Any, Iterator, Sequence, Union from flink_gateway.exceptions import InterfaceError, ProgrammingError, TimeoutError from flink_gateway.models import ( @@ -52,6 +53,8 @@ def __init__( self._columns: list[ColumnInfo] = [] self._query_timeout = query_timeout self._idle_timeout = idle_timeout + self._effective_query_timeout: float | None = query_timeout + self._effective_idle_timeout: float | None = idle_timeout self._query_deadline: float = float("inf") # ── PEP 249 attributes ───────────────────────────────────────── @@ -95,16 +98,37 @@ def _check_open(self) -> None: # ── Execution ────────────────────────────────────────────────── + class _Sentinel(enum.Enum): + """Sentinel value to distinguish 'not provided' from ``None``.""" + + MISSING = enum.auto() + + _MISSING: _Sentinel = _Sentinel.MISSING + def execute( self, operation: str, parameters: Sequence[Any] | None = None, + *, + query_timeout: Union[float, None, _Sentinel] = _MISSING, + idle_timeout: Union[float, None, _Sentinel] = _MISSING, ) -> Cursor: """Execute a SQL statement. Args: operation: The SQL string. parameters: Not yet supported. + query_timeout: Maximum seconds this query may run before raising + :class:`~flink_gateway.TimeoutError`. Overrides the + connection-level default for this execution only. + Pass ``None`` to explicitly disable the timeout. + Omit (or leave as default) to use the connection-level value. + idle_timeout: Maximum seconds to wait between rows during + streaming iteration before raising + :class:`~flink_gateway.TimeoutError`. Overrides the + connection-level default for this execution only. + Pass ``None`` to explicitly disable the timeout. + Omit (or leave as default) to use the connection-level value. Returns: self (for chaining). @@ -114,15 +138,27 @@ def execute( # Close any previous operation. self._close_current_operation() + # Resolve effective timeouts: per-execute overrides > connection defaults. + eff_qt: float | None = ( + self._query_timeout + if isinstance(query_timeout, self._Sentinel) + else query_timeout + ) + eff_it: float | None = ( + self._idle_timeout + if isinstance(idle_timeout, self._Sentinel) + else idle_timeout + ) + self._effective_query_timeout = eff_qt + self._effective_idle_timeout = eff_it + # Reset state. self._description = None 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") + time.monotonic() + eff_qt if eff_qt is not None else float("inf") ) if parameters is not None: @@ -245,7 +281,9 @@ def _fetch_until_ready(self, op_handle: str) -> FetchResultsResponse: if result.result_type != ResultType.NOT_READY: return result if time.monotonic() > self._query_deadline: - raise TimeoutError(f"query timed out after {self._query_timeout}s") + raise TimeoutError( + f"query timed out after {self._effective_query_timeout}s" + ) time.sleep(_POLL_INTERVAL) def _build_description(self) -> None: @@ -280,9 +318,10 @@ def _iter_rows( client = self._connection.client session = self._connection.session_handle + effective_idle_timeout = self._effective_idle_timeout idle_deadline: float = ( - time.monotonic() + self._idle_timeout - if self._idle_timeout is not None + time.monotonic() + effective_idle_timeout + if effective_idle_timeout is not None else float("inf") ) @@ -299,7 +338,7 @@ def _iter_rows( if time.monotonic() >= self._query_deadline: raise TimeoutError( f"streaming iteration stopped: query_timeout of " - f"{self._query_timeout}s exceeded" + f"{self._effective_query_timeout}s exceeded" ) continue @@ -319,12 +358,12 @@ def _iter_rows( if now >= self._query_deadline: raise TimeoutError( f"streaming iteration stopped: query_timeout of " - f"{self._query_timeout}s exceeded" + f"{self._effective_query_timeout}s exceeded" ) if now >= idle_deadline: raise TimeoutError( f"streaming iteration stopped: idle_timeout of " - f"{self._idle_timeout}s exceeded" + f"{effective_idle_timeout}s exceeded" ) sleep_dur = min( backoff, self._query_deadline - now, idle_deadline - now @@ -333,8 +372,8 @@ def _iter_rows( 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 + if effective_idle_timeout is not None: + idle_deadline = time.monotonic() + effective_idle_timeout def _decode_row(self, row: RowData) -> tuple[Any, ...]: """Decode a single RowData into a Python tuple.""" diff --git a/tests/test_dbapi.py b/tests/test_dbapi.py index 645673e..e95412e 100644 --- a/tests/test_dbapi.py +++ b/tests/test_dbapi.py @@ -887,3 +887,131 @@ def test_no_timeout_defaults_unchanged(self): cur = conn.cursor() cur.execute("SELECT v FROM t") assert cur.fetchall() == [(1,), (2,), (3,), (4,)] + + +class TestPerExecuteTimeouts: + """Tests for per-execute query_timeout and idle_timeout overrides.""" + + def test_execute_query_timeout_overrides_connection_default(self): + """Per-execute query_timeout should override connection-level default.""" + 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"), + ] + + # Connection has query_timeout=10.0, but execute overrides with 1.0 + # 0.0 → execute(): _query_deadline = 0.0 + 1.0 = 1.0 + # 0.0 → post-yield row 1 (ok) + # 0.5 → empty-page 1 check (ok) + # 1.5 → empty-page 2 check → RAISE + monotonic_seq = iter([0.0, 0.0, 0.5, 1.5]) + + conn = _make_connection(client, query_timeout=10.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", query_timeout=1.0) + with pytest.raises(TimeoutError, match="query_timeout"): + cur.fetchall() + + def test_execute_idle_timeout_overrides_connection_default(self): + """Per-execute idle_timeout should override connection-level default.""" + 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"), + ] + + # Connection has idle_timeout=60.0, but execute overrides with 0.3 + # 0.0 → _iter_rows: idle_deadline = 0.3 + # 0.0 → post-yield row 1 → ok + # 0.25 → empty-page 1 check → ok (0.25 < 0.3) + # 0.35 → empty-page 2 check → RAISE (0.35 >= 0.3) + monotonic_seq = iter([0.0, 0.0, 0.25, 0.35]) + + conn = _make_connection(client, idle_timeout=60.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", idle_timeout=0.3) + with pytest.raises(TimeoutError, match="idle_timeout"): + cur.fetchall() + + def test_execute_none_disables_connection_timeout(self): + """Passing None to execute() should disable the connection-level timeout.""" + 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(), + ] + + # Connection has aggressive timeouts, but execute disables them. + conn = _make_connection(client, query_timeout=0.001, idle_timeout=0.001) + cur = conn.cursor() + cur.execute("SELECT v FROM t", query_timeout=None, idle_timeout=None) + rows = cur.fetchall() + assert rows == [(1,), (2,), (3,), (4,)] + + def test_omitted_execute_timeout_uses_connection_default(self): + """When execute() omits timeout kwargs, connection defaults are used.""" + 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"), + ] + + # Connection has query_timeout=1.0; execute() does NOT override. + monotonic_seq = iter([0.0, 0.0, 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") # no override + with pytest.raises(TimeoutError, match="query_timeout"): + cur.fetchall() + + def test_per_execute_timeout_does_not_persist_across_executions(self): + """Per-execute timeout should not bleed into the next execute() call.""" + client = _mock_client() + cols = _columns(("v", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + + # First execute: no connection timeout, per-execute query_timeout=1.0 + # Second execute: no per-execute override → should use connection default (None) + client.fetch_results.side_effect = [ + # First query — single page, no next_uri + _ready_result(cols, [[1]]), + # Second query — single page, no next_uri + _ready_result(cols, [[2], [3]]), + ] + + conn = _make_connection(client) # no connection-level timeout + cur = conn.cursor() + + cur.execute("SELECT 1", query_timeout=1.0) + assert cur.fetchall() == [(1,)] + + # Second execute without per-execute override falls back to None + cur.execute("SELECT 2") + rows = cur.fetchall() + assert rows == [(2,), (3,)]