feat(cursor): Add redshift_rowcount for SELECT rowcount support

Author: Brooke-white

redshift_connector/core.py

@@ -1241,6 +1241,7 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
 
         cursor._cached_rows.clear()
         cursor._row_count = -1
+        cursor._redshift_row_count = -1
 
         # Byte1('B') - Identifies the Bind command.
         # Int32 - Message length, including self.
@@ -1315,6 +1316,10 @@ def handle_NO_DATA(self: "Connection", msg, ps) -> None:
     # how many SQL commands was completed
     # not support for 'SELECT' and 'COPY' query
     def handle_COMMAND_COMPLETE(self: "Connection", data: bytes, cursor: Cursor) -> None:
+        """
+        Modifies the cursor object and prepared statement when receiving COMMAND COMPLETE b'C' message from Redshift
+        server.
+        """
         values: typing.List[bytes] = data[:-1].split(b" ")
         command = values[0]
         if command in self._commands_with_count:
@@ -1323,6 +1328,12 @@ def handle_COMMAND_COMPLETE(self: "Connection", data: bytes, cursor: Cursor) ->
                 cursor._row_count = row_count
             else:
                 cursor._row_count += row_count
+            cursor._redshift_row_count = cursor._row_count
+        elif command == b"SELECT":
+            # Redshift server does not support row count for SELECT statement
+            # so we derive this from the size of the rows associated with the
+            # cursor object
+            cursor._redshift_row_count = len(cursor._cached_rows)
 
         if command in (b"ALTER", b"CREATE"):
             for scache in self._caches.values():

redshift_connector/cursor.py

@@ -95,6 +95,7 @@ def __init__(self: "Cursor", connection: "Connection", paramstyle=None) -> None:
         self.arraysize: int = 1
         self.ps: typing.Optional[typing.Dict[str, typing.Any]] = None
         self._row_count: int = -1
+        self._redshift_row_count: int = -1
         self._cached_rows: deque = deque()
         if paramstyle is None:
             self.paramstyle: str = redshift_connector.paramstyle
@@ -116,8 +117,28 @@ def connection(self: "Cursor") -> typing.Optional["Connection"]:
 
     @property
     def rowcount(self: "Cursor") -> int:
+        """
+        This read-only attribute specifies the number of rows that the last .execute*() produced
+        (for DQL statements like SELECT) or affected (for DML statements like UPDATE or INSERT).
+
+        The attribute is -1 in case no .execute*() has been performed on the cursor or the rowcount of the last
+        operation is cannot be determined by the interface.
+        """
         return self._row_count
 
+    @property
+    def redshift_rowcount(self: "Cursor") -> int:
+        """
+        Native to ``redshift_connector``, this read-only attribute specifies the number of rows that the last .execute*() produced.
+
+        For DQL statements (like SELECT) the number of rows is derived by ``redshift_connector`` rather than
+        provided by the server. For DML statements (like UPDATE or INSERT) this value is provided by the server.
+
+        This property's behavior is subject to change inline with modifications made to query execution.
+        Use ``rowcount`` as an alternative to this property.
+        """
+        return self._redshift_row_count
+
     @typing.no_type_check
     @functools.lru_cache()
     def truncated_row_desc(self: "Cursor"):
@@ -232,11 +253,14 @@ def executemany(self: "Cursor", operation, param_sets) -> "Cursor":
         The Cursor object used for executing the specified database operation: :class:`Cursor`
         """
         rowcounts: typing.List[int] = []
+        redshift_rowcounts: typing.List[int] = []
         for parameters in param_sets:
             self.execute(operation, parameters)
             rowcounts.append(self._row_count)
+            redshift_rowcounts.append(self._redshift_row_count)
 
         self._row_count = -1 if -1 in rowcounts else sum(rowcounts)
+        self._redshift_row_count = -1 if -1 in redshift_rowcounts else sum(rowcounts)
         return self
 
     def insert_data_bulk(

test/integration/test_query.py

@@ -104,15 +104,21 @@ def test_insert_returning(db_table):
         # Test INSERT ... RETURNING with one row...
         cursor.execute("INSERT INTO t2 VALUES (%s, %s)", (row_id, "test1"))
 
-        assert cursor.rowcount == 1
+        assert 1 == cursor.rowcount
+        assert 1 == cursor.redshift_rowcount
 
         cursor.execute("SELECT data FROM t2 WHERE id = %s", (row_id,))
+        assert 1 == cursor.redshift_rowcount
+
         assert "test1" == cursor.fetchone()[0]
 
         # Test with multiple rows...
         cursor.execute("INSERT INTO t2 VALUES (2, 'test2'), (3, 'test3'), (4,'test4') ")
-        assert cursor.rowcount == 3
+        assert 3 == cursor.rowcount
+        assert 3 == cursor.redshift_rowcount
+
         cursor.execute("SELECT * FROM t2")
+        assert 4 == cursor.redshift_rowcount
         ids: typing.Tuple[typing.List[typing.Union[int, str], ...]] = cursor.fetchall()
         assert len(ids) == 4
 
@@ -134,12 +140,14 @@ def test_row_count(db_table):
 
         # Check row_count without doing any reading first...
         assert -1 == cursor.rowcount
+        assert expected_count == cursor.redshift_rowcount
 
         # Check rowcount after reading some rows, make sure it still
         # works...
         for i in range(expected_count // 2):
             cursor.fetchone()
         assert -1 == cursor.rowcount
+        assert expected_count == cursor.redshift_rowcount
 
     with db_table.cursor() as cursor:
         # Restart the cursor, read a few rows, and then check rowcount
@@ -148,10 +156,12 @@ def test_row_count(db_table):
         for i in range(expected_count // 3):
             cursor.fetchone()
         assert -1 == cursor.rowcount
+        assert expected_count == cursor.redshift_rowcount
 
         # Should be -1 for a command with no results
         cursor.execute("DROP TABLE t1")
         assert -1 == cursor.rowcount
+        assert -1 == cursor.redshift_rowcount
 
 
 def test_row_count_fetch(db_table):
@@ -160,6 +170,7 @@ def test_row_count_fetch(db_table):
         cursor.execute("SELECT * FROM t1")
         cursor.fetchall()
         assert -1 == cursor.rowcount
+        assert 1 == cursor.redshift_rowcount
 
 
 def test_row_count_update(db_table):
@@ -170,7 +181,8 @@ def test_row_count_update(db_table):
         cursor.execute("INSERT INTO t1 (f1, f2, f3) VALUES (%s, %s, %s)", (4, 1000, None))
         cursor.execute("INSERT INTO t1 (f1, f2, f3) VALUES (%s, %s, %s)", (5, 10000, None))
         cursor.execute("UPDATE t1 SET f3 = %s WHERE f2 > 101", ("Hello!",))
-        assert cursor.rowcount == 2
+        assert 2 == cursor.rowcount
+        assert 2 == cursor.redshift_rowcount
 
 
 def test_int_oid(cursor):
@@ -199,7 +211,9 @@ def test_transactions(db_table):
     with db_table.cursor() as cursor:
         cursor.execute("commit")
         cursor.execute("INSERT INTO t1 (f1, f2, f3) VALUES (%s, %s, %s)", (1, 1, "Zombie"))
-        assert cursor.rowcount == 1
+        assert 1 == cursor.rowcount
+        assert 1 == cursor.redshift_rowcount
+
         cursor.execute("rollback")
         cursor.execute("select * from t1")