Acknowledge reception of data in TrinoResult

Author: mdesmet

trino/client.py

@@ -594,28 +594,26 @@ class TrinoResult(object):
 
     def __init__(self, query, rows=None):
         self._query = query
-        self._rows = rows or []
+        # Initial rows from the first POST request
+        self._rows = rows
         self._rownumber = 0
 
     @property
     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 row
-        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 row
 
+            self._rows = next_rows
+
     @property
     def response_headers(self):
         return self._query.response_headers
@@ -641,7 +639,7 @@ def __init__(
         self._request = request
         self._update_type = None
         self._sql = sql
-        self._result = TrinoResult(self)
+        self._result: Optional[TrinoResult] = None
         self._response_headers = None
         self._experimental_python_types = experimental_python_types
         self._row_mapper: Optional[RowMapper] = None

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)

trino/sqlalchemy/dialect.py

@@ -231,7 +231,7 @@ def get_view_definition(self, connection: Connection, view_name: str, schema: st
         """
         ).strip()
         res = connection.execute(sql.text(query), schema=schema, view=view_name)
-        return res.scalar()
+        return res.scalar_one_or_none()
 
     def get_indexes(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]:
         if not self.has_table(connection, table_name, schema):
@@ -284,7 +284,7 @@ def get_table_comment(self, connection: Connection, table_name: str, schema: str
                 sql.text(query),
                 catalog_name=catalog_name, schema_name=schema_name, table_name=table_name
             )
-            return dict(text=res.scalar())
+            return dict(text=res.scalar_one_or_none())
         except error.TrinoQueryError as e:
             if e.error_name in (
                 error.PERMISSION_DENIED,
@@ -326,7 +326,7 @@ def _get_server_version_info(self, connection: Connection) -> Any:
         query = "SELECT version()"
         try:
             res = connection.execute(sql.text(query))
-            version = res.scalar()
+            version = res.scalar_one()
             return tuple([version])
         except exc.ProgrammingError as e:
             logger.debug(f"Failed to get server version: {e.orig.message}")