diff --git a/src/flink_gateway/cursor.py b/src/flink_gateway/cursor.py index 816ba0c..a9b1d40 100644 --- a/src/flink_gateway/cursor.py +++ b/src/flink_gateway/cursor.py @@ -284,6 +284,8 @@ def _iter_rows( continue # Fetch next page. + if not token: + return response = client.fetch_results(session, op_handle, token, "") if response.result_type == ResultType.EOS: return diff --git a/tests/test_dbapi.py b/tests/test_dbapi.py index 24c9a16..ef82634 100644 --- a/tests/test_dbapi.py +++ b/tests/test_dbapi.py @@ -292,6 +292,37 @@ def test_pagination(self): rows = cur.fetchall() assert rows == [(1,), (2,), (3,), (4,)] + def test_single_page_eos_no_extra_fetch(self): + client = _mock_client() + cols = _columns(("id", "INTEGER", False)) + client.execute_statement.return_value = "op-1" + client.fetch_results.side_effect = [ + FetchResultsResponse( + result_type=ResultType.EOS, + result_kind=ResultKind.SUCCESS_WITH_CONTENT, + is_query_result=True, + next_result_uri="", + results=ResultSet( + columns=cols, + data=[ + RowData(kind="INSERT", fields=[1]), + RowData(kind="INSERT", fields=[2]), + ], + ), + ), + RuntimeError( + "fetch_results called with empty token — would be 404 on server" + ), + ] + + conn = _make_connection(client) + cur = conn.cursor() + cur.execute("SELECT id FROM t") + rows = list(cur) + + assert rows == [(1,), (2,)] + assert client.fetch_results.call_count == 1 + def test_exec_non_query(self): """Non-query statements should not produce rows.""" client = _mock_client()