NOTASK add timeouts

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

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

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

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

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.