Return row count if available in status

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -1268,6 +1268,32 @@ def test_describe_table_query(run_trino):
     ]
 
 
+def test_rowcount_select(trino_connection):
+    cur = trino_connection.cursor()
+    cur.execute("SELECT 1 as a")
+    cur.fetchall()
+    assert cur.rowcount == -1
+
+
+def test_rowcount_create_table(trino_connection):
+    with _TestTable(trino_connection, "memory.default.test_rowcount_create_table", "(a varchar)") as (_, cur):
+        assert cur.rowcount == -1
+
+
+def test_rowcount_create_table_as_select(trino_connection):
+    with _TestTable(
+        trino_connection,
+        "memory.default.test_rowcount_ctas", "AS SELECT 1 a UNION ALL SELECT 2"
+    ) as (_, cur):
+        assert cur.rowcount == 2
+
+
+def test_rowcount_insert(trino_connection):
+    with _TestTable(trino_connection, "memory.default.test_rowcount_ctas", "(a VARCHAR)") as (table, cur):
+        cur.execute(f"INSERT INTO {table.table_name} (a) VALUES ('test')")
+        assert cur.rowcount == 1
+
+
 def assert_cursor_description(cur, trino_type, size=None, precision=None, scale=None):
     assert cur.description[0][1] == trino_type
     assert cur.description[0][2] is None

trino/client.py

@@ -299,6 +299,7 @@ class TrinoStatus:
     info_uri: str
     next_uri: Optional[str]
     update_type: Optional[str]
+    update_count: Optional[int]
     rows: List[Any]
     columns: List[Any]
 
@@ -666,6 +667,7 @@ def process(self, http_response) -> TrinoStatus:
             info_uri=response["infoUri"],
             next_uri=self._next_uri,
             update_type=response.get("updateType"),
+            update_count=response.get("updateCount"),
             rows=response.get("data", []),
             columns=response.get("columns"),
         )
@@ -743,6 +745,7 @@ def __init__(
         self._cancelled = False
         self._request = request
         self._update_type = None
+        self._update_count = None
         self._sql = sql
         self._result: Optional[TrinoResult] = None
         self._legacy_primitive_types = legacy_primitive_types
@@ -765,6 +768,10 @@ def stats(self):
     def update_type(self):
         return self._update_type
 
+    @property
+    def update_count(self):
+        return self._update_count
+
     @property
     def warnings(self):
         return self._warnings
@@ -809,6 +816,7 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
     def _update_state(self, status):
         self._stats.update(status.stats)
         self._update_type = status.update_type
+        self._update_count = status.update_count
         if not self._row_mapper and status.columns:
             self._row_mapper = RowMapperFactory().create(columns=status.columns,
                                                          legacy_primitive_types=self._legacy_primitive_types)

trino/dbapi.py

@@ -315,13 +315,20 @@ def description(self) -> List[ColumnDescription]:
 
     @property
     def rowcount(self):
-        """Not supported.
+        """The rowcount will be returned for INSERT, UPDATE, DELETE, MERGE
+        and CTAS statements based on `update_count` returned by the Trino
+        API.
 
-        Trino cannot reliablity determine the number of rows returned by an
-        operation. For example, the result of a SELECT query is streamed and
-        the number of rows is only knowns when all rows have been retrieved.
-        """
+        If the rowcount can't be determined, -1 will be returned.
+
+        Trino cannot reliably determine the number of rows returned for DQL
+        queries. For example, the result of a SELECT query is streamed and
+        the number of rows is only known when all rows have been retrieved.
 
+        See https://peps.python.org/pep-0249/#rowcount
+        """
+        if self._query is not None and self._query.update_count is not None:
+            return self._query.update_count
         return -1
 
     @property