implement autoscroll for all fetches

Author: vgvoleg

tests/conftest.py

@@ -151,18 +151,17 @@ async def session_pool(
 ) -> AsyncGenerator[ydb.aio.QuerySessionPool]:
     session_pool = ydb.aio.QuerySessionPool(driver)
     async with session_pool:
-        await session_pool.execute_with_retries(
-            """DROP TABLE IF EXISTS table"""
-        )
-        await session_pool.execute_with_retries(
-            """
-            CREATE TABLE table (
-            id Int64 NOT NULL,
-            val Int64,
-            PRIMARY KEY(id)
+        for name in ["table", "table1", "table2"]:
+            await session_pool.execute_with_retries(
+                f"""
+                DROP TABLE IF EXISTS {name};
+                CREATE TABLE {name} (
+                id Int64 NOT NULL,
+                val Int64,
+                PRIMARY KEY(id)
+                )
+                """
             )
-            """
-        )
 
         yield session_pool
 
@@ -173,14 +172,16 @@ def session_pool_sync(
 ) -> Generator[ydb.QuerySessionPool]:
     session_pool = ydb.QuerySessionPool(driver_sync)
     with session_pool:
-        session_pool.execute_with_retries("""DROP TABLE IF EXISTS table""")
-        session_pool.execute_with_retries(
-            """
-            CREATE TABLE table (
-            id Int64 NOT NULL,
-            val Int64,
-            PRIMARY KEY(id)
+        for name in ["table", "table1", "table2"]:
+            session_pool.execute_with_retries(
+                f"""
+                DROP TABLE IF EXISTS {name};
+                CREATE TABLE {name} (
+                id Int64 NOT NULL,
+                val Int64,
+                PRIMARY KEY(id)
+                )
+                """
             )
-            """
-        )
+
         yield session_pool

tests/test_cursor.py

@@ -5,21 +5,22 @@
 import ydb
 import ydb_dbapi
 
-INSERT_YQL = """
-DELETE FROM table;
-INSERT INTO table (id, val) VALUES
-(1, 1),
-(2, 2),
-(3, 3),
-(4, 4)
-"""
-
 
 @pytest.fixture
 async def session(
     session_pool: ydb.aio.QuerySessionPool,
 ) -> AsyncGenerator[ydb.aio.QuerySession]:
-    await session_pool.execute_with_retries(INSERT_YQL)
+    for name in ["table", "table1", "table2"]:
+        await session_pool.execute_with_retries(
+            f"""
+            DELETE FROM {name};
+            INSERT INTO {name} (id, val) VALUES
+            (0, 0),
+            (1, 1),
+            (2, 2),
+            (3, 3)
+            """
+        )
 
     session = await session_pool.acquire()
     yield session
@@ -30,13 +31,27 @@ async def session(
 def session_sync(
     session_pool_sync: ydb.QuerySessionPool,
 ) -> Generator[ydb.QuerySession]:
-    session_pool_sync.execute_with_retries(INSERT_YQL)
+    for name in ["table", "table1", "table2"]:
+        session_pool_sync.execute_with_retries(
+            f"""
+            DELETE FROM {name};
+            INSERT INTO {name} (id, val) VALUES
+            (0, 0),
+            (1, 1),
+            (2, 2),
+            (3, 3)
+            """
+        )
 
     session = session_pool_sync.acquire()
     yield session
     session_pool_sync.release(session)
 
 
+RESULT_SET_LENGTH = 4
+RESULT_SET_COUNT = 3
+
+
 class TestAsyncCursor:
     @pytest.mark.asyncio
     async def test_cursor_fetch_one(
@@ -48,10 +63,10 @@ async def test_cursor_fetch_one(
             """
             await cursor.execute(query=yql_text)
 
-            for i in range(4):
+            for i in range(RESULT_SET_LENGTH):
                 res = await cursor.fetchone()
                 assert res is not None
-                assert res[0] == i + 1
+                assert res[0] == i
 
             assert await cursor.fetchone() is None
 
@@ -68,20 +83,20 @@ async def test_cursor_fetch_many(
             res = await cursor.fetchmany()
             assert res is not None
             assert len(res) == 1
-            assert res[0][0] == 1
+            assert res[0][0] == 0
 
             res = await cursor.fetchmany(size=2)
             assert res is not None
             assert len(res) == 2
-            assert res[0][0] == 2
-            assert res[1][0] == 3
+            assert res[0][0] == 1
+            assert res[1][0] == 2
 
             res = await cursor.fetchmany(size=2)
             assert res is not None
             assert len(res) == 1
-            assert res[0][0] == 4
+            assert res[0][0] == 3
 
-            assert await cursor.fetchmany(size=2) is None
+            assert await cursor.fetchmany(size=2) == []
 
     @pytest.mark.asyncio
     async def test_cursor_fetch_all(
@@ -93,15 +108,15 @@ async def test_cursor_fetch_all(
             """
             await cursor.execute(query=yql_text)
 
-            assert cursor.rowcount == 4
+            assert cursor.rowcount == RESULT_SET_LENGTH
 
             res = await cursor.fetchall()
             assert res is not None
-            assert len(res) == 4
-            for i in range(4):
-                assert res[i][0] == i + 1
+            assert len(res) == RESULT_SET_LENGTH
+            for i in range(RESULT_SET_LENGTH):
+                assert res[i][0] == i
 
-            assert await cursor.fetchall() is None
+            assert await cursor.fetchall() == []
 
     @pytest.mark.asyncio
     async def test_cursor_next_set(
@@ -127,11 +142,77 @@ async def test_cursor_next_set(
             nextset = await cursor.nextset()
             assert nextset
 
-            assert await cursor.fetchall() is None
+            assert await cursor.fetchall() == []
 
             nextset = await cursor.nextset()
             assert not nextset
 
+    @pytest.mark.asyncio
+    async def test_cursor_fetch_one_autoscroll(
+        self, session: ydb.aio.QuerySession
+    ) -> None:
+        async with ydb_dbapi.AsyncCursor(
+            session=session, auto_scroll_result_sets=True
+        ) as cursor:
+            yql_text = """
+            SELECT id, val FROM table;
+            SELECT id, val FROM table1;
+            SELECT id, val FROM table2;
+            """
+            await cursor.execute(query=yql_text)
+
+            for i in range(RESULT_SET_LENGTH * RESULT_SET_COUNT):
+                res = await cursor.fetchone()
+                assert res is not None
+                assert res[0] == i % RESULT_SET_LENGTH
+
+            assert await cursor.fetchone() is None
+            assert not await cursor.nextset()
+
+    @pytest.mark.asyncio
+    async def test_cursor_fetch_many_autoscroll(
+        self, session: ydb.aio.QuerySession
+    ) -> None:
+        async with ydb_dbapi.AsyncCursor(
+            session=session, auto_scroll_result_sets=True
+        ) as cursor:
+            yql_text = """
+            SELECT id, val FROM table;
+            SELECT id, val FROM table1;
+            SELECT id, val FROM table2;
+            """
+            await cursor.execute(query=yql_text)
+
+            halfsize = (RESULT_SET_LENGTH * RESULT_SET_COUNT) // 2
+            for _ in range(2):
+                res = await cursor.fetchmany(size=halfsize)
+                assert res is not None
+                assert len(res) == halfsize
+
+            assert await cursor.fetchmany(2) == []
+            assert not await cursor.nextset()
+
+    @pytest.mark.asyncio
+    async def test_cursor_fetch_all_autoscroll(
+        self, session: ydb.aio.QuerySession
+    ) -> None:
+        async with ydb_dbapi.AsyncCursor(
+            session=session, auto_scroll_result_sets=True
+        ) as cursor:
+            yql_text = """
+            SELECT id, val FROM table;
+            SELECT id, val FROM table1;
+            SELECT id, val FROM table2;
+            """
+            await cursor.execute(query=yql_text)
+
+            res = await cursor.fetchall()
+
+            assert len(res) == RESULT_SET_COUNT * RESULT_SET_LENGTH
+
+            assert await cursor.fetchall() == []
+            assert not await cursor.nextset()
+
 
 # The same test class as above but for Cursor
 
@@ -147,7 +228,7 @@ def test_cursor_fetch_one(self, session_sync: ydb.QuerySession) -> None:
             for i in range(4):
                 res = cursor.fetchone()
                 assert res is not None
-                assert res[0] == i + 1
+                assert res[0] == i
 
             assert cursor.fetchone() is None
 
@@ -161,20 +242,20 @@ def test_cursor_fetch_many(self, session_sync: ydb.QuerySession) -> None:
             res = cursor.fetchmany()
             assert res is not None
             assert len(res) == 1
-            assert res[0][0] == 1
+            assert res[0][0] == 0
 
             res = cursor.fetchmany(size=2)
             assert res is not None
             assert len(res) == 2
-            assert res[0][0] == 2
-            assert res[1][0] == 3
+            assert res[0][0] == 1
+            assert res[1][0] == 2
 
             res = cursor.fetchmany(size=2)
             assert res is not None
             assert len(res) == 1
-            assert res[0][0] == 4
+            assert res[0][0] == 3
 
-            assert cursor.fetchmany(size=2) is None
+            assert cursor.fetchmany(size=2) == []
 
     def test_cursor_fetch_all(self, session_sync: ydb.QuerySession) -> None:
         with ydb_dbapi.Cursor(session=session_sync) as cursor:
@@ -189,9 +270,9 @@ def test_cursor_fetch_all(self, session_sync: ydb.QuerySession) -> None:
             assert res is not None
             assert len(res) == 4
             for i in range(4):
-                assert res[i][0] == i + 1
+                assert res[i][0] == i
 
-            assert cursor.fetchall() is None
+            assert cursor.fetchall() == []
 
     def test_cursor_next_set(self, session_sync: ydb.QuerySession) -> None:
         with ydb_dbapi.Cursor(session=session_sync) as cursor:
@@ -214,22 +295,70 @@ def test_cursor_next_set(self, session_sync: ydb.QuerySession) -> None:
             nextset = cursor.nextset()
             assert nextset
 
-            assert cursor.fetchall() is None
+            assert cursor.fetchall() == []
 
             nextset = cursor.nextset()
             assert not nextset
 
-    def test_cursor_autoscroll(self, session_sync: ydb.QuerySession) -> None:
+    def test_cursor_fetch_one_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;"
+            yql_text = """
+            SELECT id, val FROM table;
+            SELECT id, val FROM table1;
+            SELECT id, val FROM table2;
+            """
             cursor.execute(query=yql_text)
 
-            for i in range(3):
+            for i in range(RESULT_SET_LENGTH * RESULT_SET_COUNT):
                 res = cursor.fetchone()
                 assert res is not None
-                assert res[0] == i + 1
+                assert res[0] == i % RESULT_SET_LENGTH
 
             assert cursor.fetchone() is None
             assert not cursor.nextset()
+
+    def test_cursor_fetch_many_autoscroll(
+        self, session_sync: ydb.QuerySession
+    ) -> None:
+        with ydb_dbapi.Cursor(
+            session=session_sync, auto_scroll_result_sets=True
+        ) as cursor:
+            yql_text = """
+            SELECT id, val FROM table;
+            SELECT id, val FROM table1;
+            SELECT id, val FROM table2;
+            """
+            cursor.execute(query=yql_text)
+
+            halfsize = (RESULT_SET_LENGTH * RESULT_SET_COUNT) // 2
+            for _ in range(2):
+                res = cursor.fetchmany(size=halfsize)
+                assert res is not None
+                assert len(res) == halfsize
+
+            assert cursor.fetchmany(2) == []
+            assert not cursor.nextset()
+
+    def test_cursor_fetch_all_autoscroll(
+        self, session_sync: ydb.QuerySession
+    ) -> None:
+        with ydb_dbapi.Cursor(
+            session=session_sync, auto_scroll_result_sets=True
+        ) as cursor:
+            yql_text = """
+            SELECT id, val FROM table;
+            SELECT id, val FROM table1;
+            SELECT id, val FROM table2;
+            """
+            cursor.execute(query=yql_text)
+
+            res = cursor.fetchall()
+
+            assert len(res) == RESULT_SET_COUNT * RESULT_SET_LENGTH
+
+            assert cursor.fetchall() == []
+            assert not cursor.nextset()