fix: Multistatement fix idx (#118)

* fixes for multi-statement queries

* extend unit tests

* add integration tests

Author: stepansergeevitch

src/firebolt/async_db/cursor.py

@@ -222,6 +222,7 @@ def _pop_next_set(self) -> Optional[bool]:
         self._rowcount, self._descriptions, self._rows = self._row_sets[
             self._next_set_idx
         ]
+        self._idx = 0
         self._next_set_idx += 1
         return True
 
@@ -464,6 +465,11 @@ async def fetchall(self) -> List[List[ColType]]:
             return super().fetchall()
         """Fetch all remaining rows of a query result"""
 
+    @wraps(BaseCursor.nextset)
+    async def nextset(self) -> None:
+        async with self._async_query_lock.reader:
+            return super().nextset()
+
     # Iteration support
     @check_not_closed
     @check_query_executed

src/firebolt/db/cursor.py

@@ -75,6 +75,11 @@ def fetchall(self) -> List[List[ColType]]:
         with self._query_lock.gen_rlock():
             return super().fetchall()
 
+    @wraps(AsyncBaseCursor.nextset)
+    def nextset(self) -> None:
+        with self._query_lock.gen_rlock(), self._idx_lock:
+            return super().nextset()
+
     # Iteration support
     @check_not_closed
     @check_query_executed

tests/integration/dbapi/async/test_queries_async.py

@@ -260,14 +260,15 @@ async def test_multi_statement_query(connection: Connection) -> None:
         assert (
             await c.execute(
                 "INSERT INTO test_tb_multi_statement values (1, 'a'), (2, 'b');"
-                "SELECT * FROM test_tb_multi_statement"
+                "SELECT * FROM test_tb_multi_statement;"
+                "SELECT * FROM test_tb_multi_statement WHERE i <= 1"
             )
             == -1
         ), "Invalid row count returned for insert"
         assert c.rowcount == -1, "Invalid row count"
         assert c.description is None, "Invalid description"
 
-        assert c.nextset()
+        assert await c.nextset()
 
         assert c.rowcount == 2, "Invalid select row count"
         assert_deep_eq(
@@ -285,4 +286,22 @@ async def test_multi_statement_query(connection: Connection) -> None:
             "Invalid data in table after parameterized insert",
         )
 
-        assert c.nextset() is None
+        assert await c.nextset()
+
+        assert c.rowcount == 1, "Invalid select row count"
+        assert_deep_eq(
+            c.description,
+            [
+                Column("i", int, None, None, None, None, None),
+                Column("s", str, None, None, None, None, None),
+            ],
+            "Invalid select query description",
+        )
+
+        assert_deep_eq(
+            await c.fetchall(),
+            [[1, "a"]],
+            "Invalid data in table after parameterized insert",
+        )
+
+        assert await c.nextset() is None

tests/integration/dbapi/sync/test_queries.py

@@ -251,7 +251,8 @@ def test_multi_statement_query(connection: Connection) -> None:
         assert (
             c.execute(
                 "INSERT INTO test_tb_multi_statement values (1, 'a'), (2, 'b');"
-                "SELECT * FROM test_tb_multi_statement"
+                "SELECT * FROM test_tb_multi_statement;"
+                "SELECT * FROM test_tb_multi_statement WHERE i <= 1"
             )
             == -1
         ), "Invalid row count returned for insert"
@@ -276,4 +277,22 @@ def test_multi_statement_query(connection: Connection) -> None:
             "Invalid data in table after parameterized insert",
         )
 
+        assert c.nextset()
+
+        assert c.rowcount == 1, "Invalid select row count"
+        assert_deep_eq(
+            c.description,
+            [
+                Column("i", int, None, None, None, None, None),
+                Column("s", str, None, None, None, None, None),
+            ],
+            "Invalid select query description",
+        )
+
+        assert_deep_eq(
+            c.fetchall(),
+            [[1, "a"]],
+            "Invalid data in table after parameterized insert",
+        )
+
         assert c.nextset() is None

tests/unit/async_db/test_cursor.py

@@ -63,8 +63,9 @@ async def test_closed_cursor(cursor: Cursor):
         ("fetchone", ()),
         ("fetchmany", ()),
         ("fetchall", ()),
+        ("nextset", ()),
     )
-    methods = ("setinputsizes", "setoutputsize", "nextset")
+    methods = ("setinputsizes", "setoutputsize")
 
     cursor.close()
 
@@ -439,8 +440,11 @@ async def test_cursor_multi_statement(
     httpx_mock.add_callback(auth_callback, url=auth_url)
     httpx_mock.add_callback(query_callback, url=query_url)
     httpx_mock.add_callback(insert_query_callback, url=query_url)
+    httpx_mock.add_callback(query_callback, url=query_url)
 
-    rc = await cursor.execute("select * from t; insert into t values (1, 2)")
+    rc = await cursor.execute(
+        "select * from t; insert into t values (1, 2); select * from t"
+    )
     assert rc == len(python_query_data), "Invalid row count returned"
     assert cursor.rowcount == len(python_query_data), "Invalid cursor row count"
     for i, (desc, exp) in enumerate(zip(cursor.description, python_query_description)):
@@ -451,12 +455,23 @@ async def test_cursor_multi_statement(
             await cursor.fetchone() == python_query_data[i]
         ), f"Invalid data row at position {i}"
 
-    assert cursor.nextset()
+    assert await cursor.nextset()
     assert cursor.rowcount == -1, "Invalid cursor row count"
     assert cursor.description is None, "Invalid cursor description"
     with raises(DataError) as exc_info:
         await cursor.fetchall()
 
     assert str(exc_info.value) == "no rows to fetch", "Invalid error message"
 
-    assert cursor.nextset() is None
+    assert await cursor.nextset()
+
+    assert cursor.rowcount == len(python_query_data), "Invalid cursor row count"
+    for i, (desc, exp) in enumerate(zip(cursor.description, python_query_description)):
+        assert desc == exp, f"Invalid column description at position {i}"
+
+    for i in range(cursor.rowcount):
+        assert (
+            await cursor.fetchone() == python_query_data[i]
+        ), f"Invalid data row at position {i}"
+
+    assert await cursor.nextset() is None

tests/unit/db/test_cursor.py

@@ -60,7 +60,7 @@ def test_closed_cursor(cursor: Cursor):
         ("fetchall", ()),
         ("setinputsizes", (cursor, [0])),
         ("setoutputsize", (cursor, 0)),
-        ("nextset", (cursor, [])),
+        ("nextset", ()),
     )
 
     cursor.close()
@@ -71,6 +71,7 @@ def test_closed_cursor(cursor: Cursor):
 
     for method, args in methods:
         with raises(CursorClosedError):
+            print(method, args)
             getattr(cursor, method)(*args)
 
     with raises(CursorClosedError):
@@ -386,8 +387,9 @@ def test_cursor_multi_statement(
     httpx_mock.add_callback(auth_callback, url=auth_url)
     httpx_mock.add_callback(query_callback, url=query_url)
     httpx_mock.add_callback(insert_query_callback, url=query_url)
+    httpx_mock.add_callback(query_callback, url=query_url)
 
-    rc = cursor.execute("select * from t; insert into t values (1, 2)")
+    rc = cursor.execute("select * from t; insert into t values (1, 2); select * from t")
     assert rc == len(python_query_data), "Invalid row count returned"
     assert cursor.rowcount == len(python_query_data), "Invalid cursor row count"
     for i, (desc, exp) in enumerate(zip(cursor.description, python_query_description)):
@@ -406,4 +408,15 @@ def test_cursor_multi_statement(
 
     assert str(exc_info.value) == "no rows to fetch", "Invalid error message"
 
+    assert cursor.nextset()
+
+    assert cursor.rowcount == len(python_query_data), "Invalid cursor row count"
+    for i, (desc, exp) in enumerate(zip(cursor.description, python_query_description)):
+        assert desc == exp, f"Invalid column description at position {i}"
+
+    for i in range(cursor.rowcount):
+        assert (
+            cursor.fetchone() == python_query_data[i]
+        ), f"Invalid data row at position {i}"
+
     assert cursor.nextset() is None