Enable autoscroll for fetchone

Author: vgvoleg

tests/test_cursor.py

@@ -218,3 +218,18 @@ def test_cursor_next_set(self, session_sync: ydb.QuerySession) -> None:
 
             nextset = cursor.nextset()
             assert not nextset
+
+    def test_cursor_autoscroll(self, session_sync: ydb.QuerySession) -> None:
+        with ydb_dbapi.Cursor(
+            session=session_sync, auto_scroll_result_sets=True
+        ) as cursor:
+            yql_text = "SELECT 1 as val; SELECT 2 as val; SELECT 3 as val;"
+            cursor.execute(query=yql_text)
+
+            for i in range(3):
+                res = cursor.fetchone()
+                assert res is not None
+                assert res[0] == i + 1
+
+            assert cursor.fetchone() is None
+            assert not cursor.nextset()

ydb_dbapi/cursors.py

@@ -127,16 +127,29 @@ def __init__(
         tx_context: ydb.QueryTxContext | None = None,
         table_path_prefix: str = "",
         autocommit: bool = True,
+        auto_scroll_result_sets: bool = False,
     ) -> None:
         self._session = session
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
         self._autocommit = autocommit
+        self._auto_scroll = auto_scroll_result_sets
 
         self._stream: Iterator | None = None
 
     def fetchone(self) -> tuple | None:
-        return self._fetchone_from_buffer()
+        row = self._fetchone_from_buffer()
+        if not self._auto_scroll:
+            return row
+
+        if row is None:
+            while self.nextset():
+                # We should skip empty result sets
+                row = self._fetchone_from_buffer()
+                if row is not None:
+                    return row
+
+        return row
 
     def fetchmany(self, size: int | None = None) -> list | None:
         return self._fetchmany_from_buffer(size)