DI-5716 execute timeout

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."""

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,)]