Multi resultset consistency

pymonetdb/sql/cursors.py

@@ -49,7 +49,7 @@ class Cursor(object):
     lastrowid: Optional[int]
     _unpack_int64: str
 
-    _next_result_sets: List[Tuple[str, int, List[Description], List[Tuple]]]
+    _next_result_sets: List[Tuple[Optional[str], int, Optional[List[Description]], List[Tuple]]]
 
     def __init__(self, connection: 'pymonetdb.sql.connections.Connection'):
         """This read-only attribute return a reference to the Connection
@@ -525,6 +525,11 @@ def _store_result(self, block, *, update_existing: bool):  # noqa: C901
                 self._rows = []
                 self.description = None
                 self.rowcount = -1
+                self._query_id = None
+                if not update_existing:
+                    self._next_result_sets.append(
+                        (self._query_id, self.rowcount, self.description, self._rows)
+                    )
 
             elif line.startswith(mapi.MSG_QUPDATE):
                 (affected, identity) = line[2:].split()[:2]
@@ -534,6 +539,10 @@ def _store_result(self, block, *, update_existing: bool):  # noqa: C901
                 self.rowcount = int(affected)
                 self.lastrowid = int(identity)
                 self._query_id = None
+                if not update_existing:
+                    self._next_result_sets.append(
+                        (self._query_id, self.rowcount, self.description, self._rows)
+                    )
 
             elif line.startswith(mapi.MSG_QTRANS):
                 self._offset = 0

tests/test_resultset.py

@@ -402,6 +402,48 @@ def access_pattern(self, n):
 
         self.assertFalse(self.cursor.nextset())
 
+    def test_multiple_with_dml_dql(self):
+        self.do_connect()
+
+        queries = 'SELECT 1;'
+        queries += 'CREATE TABLE IF NOT EXISTS bar (n int);'
+        queries += 'SELECT 10 UNION SELECT 20;'
+        queries += 'INSERT INTO bar values (10), (20);'
+        queries += "SELECT 'alice', 'bob';"
+        self.cursor.execute(queries)
+
+        self.assertEqual(self.cursor.rowcount, 1)
+        self.assertEqual(self.cursor.fetchall(), [(1,)])
+        self.assertTrue(self.cursor.description is not None)
+        self.assertTrue(self.cursor.nextset())
+
+        self.assertEqual(self.cursor.rowcount, -1)
+        try:
+            self.cursor.fetchall()
+        except Exception as e:
+            self.assertTrue(isinstance(e, pymonetdb.ProgrammingError))
+        self.assertTrue(self.cursor.description is None)
+        self.assertTrue(self.cursor.nextset())
+
+        self.assertEqual(self.cursor.rowcount, 2)
+        self.assertEqual(self.cursor.fetchall(), [(10,), (20,)])
+        self.assertTrue(self.cursor.description is not None)
+        self.assertTrue(self.cursor.nextset())
+
+        self.assertEqual(self.cursor.rowcount, 2)
+        try:
+            self.cursor.fetchall()
+        except Exception as e:
+            self.assertTrue(isinstance(e, pymonetdb.ProgrammingError))
+        self.assertTrue(self.cursor.description is None)
+        self.assertTrue(self.cursor.nextset())
+
+        self.assertEqual(self.cursor.rowcount, 1)
+        self.assertEqual(self.cursor.fetchall(), [('alice', 'bob')])
+        self.assertTrue(self.cursor.description is not None)
+
+        self.assertFalse(self.cursor.nextset())
+
     def test_huge(self):
         self.skip_unless_have_sqltype('hugeint')
         max_value = (1 << 127) - 1