Set rowcount correctly after retrieving all rows through fetchall or fetchmany

ref pymssql#141

Author: bladams

pymssql.pyx

@@ -505,8 +505,8 @@ cdef class Cursor:
 
         try:
             return self.getrow()
-
         except StopIteration:
+            self._rownumber = self._source._conn.rows_affected
             return None
         except _mssql.MSSQLDatabaseException, e:
             raise OperationalError, e.args[0]
@@ -527,6 +527,7 @@ cdef class Cursor:
                 try:
                     rows.append(self.getrow())
                 except StopIteration:
+                    self._rownumber = self._source._conn.rows_affected
                     break
             return rows
         except _mssql.MSSQLDatabaseException, e:

tests/helpers.py

@@ -285,6 +285,24 @@ def test_select_rowcount(self):
         cur.fetchall()
         eq_(cur.rowcount, 5)
 
+    def test_fetchone_rowcount(self):
+        cur = self.execute('select * from test')
+        eq_(cur.rowcount, -1)
+
+        for _ in iter(cur.fetchone, None):
+            eq_(cur.rowcount, -1)
+
+        eq_(cur.rowcount, 5)
+
+    def test_fetchmany_rowcount(self):
+        cur = self.execute('select * from test')
+        eq_(cur.rowcount, -1)
+
+        for _ in iter(cur.fetchmany, []):
+            eq_(cur.rowcount, -1)
+
+        eq_(cur.rowcount, 5)
+
     def test_as_dict(self):
         # test for http://code.google.com/p/pymssql/issues/detail?id=92
         cur = self.conn.cursor(as_dict=True)