Acknowledge reception of data in TrinoResult

Author: mdesmet

trino/client.py

@@ -594,7 +594,8 @@ class TrinoResult(object):
 
     def __init__(self, query, rows=None, experimental_python_types: bool = False):
         self._query = query
-        self._rows = rows or []
+        # Initial rows from the first POST request
+        self._rows = rows
         self._rownumber = 0
         self._experimental_python_types = experimental_python_types
 
@@ -603,20 +604,17 @@ def rownumber(self) -> int:
         return self._rownumber
 
     def __iter__(self):
-        # Initial fetch from the first POST request
-        for row in self._rows:
-            self._rownumber += 1
-            yield self._map_row(self._experimental_python_types, row, self._query.columns)
-        self._rows = None
-
-        # Subsequent fetches from GET requests until next_uri is empty.
-        while not self._query.finished:
-            rows = self._query.fetch()
-            for row in rows:
+        # A query only transitions to a FINISHED state when the results are fully consumed:
+        # The reception of the data is acknowledged by calling the next_uri before exposing the data through dbapi.
+        while not self._query.finished or self._rows is not None:
+            next_rows = self._query.fetch() if not self._query.finished else None
+            for row in self._rows:
                 self._rownumber += 1
                 logger.debug("row %s", row)
                 yield self._map_row(self._experimental_python_types, row, self._query.columns)
 
+            self._rows = next_rows
+
     @property
     def response_headers(self):
         return self._query.response_headers

trino/dbapi.py

@@ -322,7 +322,7 @@ def _prepare_statement(self, operation, statement_name):
             operation=operation
         )
 
-        # Send prepare statement. Copy the _request object to avoid poluting the
+        # Send prepare statement. Copy the _request object to avoid polluting the
         # one that is going to be used to execute the actual operation.
         query = trino.client.TrinoQuery(copy.deepcopy(self._request), sql=sql,
                                         experimental_python_types=self._experimental_pyton_types)