make cursor client-side

Author: vgvoleg

.pre-commit-config.yaml

@@ -1,3 +1,4 @@
+files: ydb_dbapi
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
@@ -23,5 +24,3 @@ repos:
     hooks:
       - id: mypy
         name: mypy
-
-exclude: '.github/'

poetry.lock

@@ -18,6 +18,7 @@ types-protobuf = "^5.28.0.20240924"
 testcontainers = "^4.8.2"
 pytest = "^8.3.3"
 pytest-asyncio = "^0.24.0"
+sqlalchemy = "^2.0.36"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -34,7 +34,7 @@ def _test_isolation_level_read_only(
             if read_only:
                 with pytest.raises(dbapi.DatabaseError):
                     cursor.execute(query)
-                    cursor.finish_query()
+                    cursor._scroll_stream()
 
             else:
                 cursor.execute(query)
@@ -53,14 +53,14 @@ def _test_connection(self, connection: dbapi.Connection) -> None:
         cur = connection.cursor()
         with suppress(dbapi.DatabaseError):
             cur.execute("DROP TABLE foo")
-            cur.finish_query()
+            cur._scroll_stream()
 
         assert not connection.check_exists("/local/foo")
         with pytest.raises(dbapi.ProgrammingError):
             connection.describe("/local/foo")
 
         cur.execute("CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))")
-        cur.finish_query()
+        cur._scroll_stream()
 
         assert connection.check_exists("/local/foo")
 
@@ -77,12 +77,12 @@ def _test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
 
         with suppress(dbapi.DatabaseError):
             cur.execute("DROP TABLE test")
-            cur.finish_query()
+            cur._scroll_stream()
 
         cur.execute(
             "CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
         )
-        cur.finish_query()
+        cur._scroll_stream()
 
         cur.execute(
             """
@@ -104,7 +104,7 @@ def _test_cursor_raw_query(self, connection: dbapi.Connection) -> None:
                 )
             },
         )
-        cur.finish_query()
+        cur._scroll_stream()
 
         cur.execute("DROP TABLE test")
 
@@ -121,7 +121,7 @@ def _test_errors(self, connection: dbapi.Connection) -> None:
 
         with suppress(dbapi.DatabaseError):
             cur.execute("DROP TABLE test")
-            cur.finish_query()
+            cur._scroll_stream()
 
         with pytest.raises(dbapi.DataError):
             cur.execute("SELECT 18446744073709551616")
@@ -136,10 +136,10 @@ def _test_errors(self, connection: dbapi.Connection) -> None:
             cur.execute("SELECT * FROM test")
 
         cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))")
-        cur.finish_query()
+        cur._scroll_stream()
 
         cur.execute("INSERT INTO test(id) VALUES(1)")
-        cur.finish_query()
+        cur._scroll_stream()
 
         with pytest.raises(dbapi.IntegrityError):
             cur.execute("INSERT INTO test(id) VALUES(1)")
@@ -214,7 +214,7 @@ async def _test_isolation_level_read_only(
             if read_only:
                 with pytest.raises(dbapi.DatabaseError):
                     await cursor.execute(query)
-                    await cursor.finish_query()
+                    await cursor._scroll_stream()
 
             else:
                 await cursor.execute(query)
@@ -235,7 +235,7 @@ async def _test_connection(
         cur = connection.cursor()
         with suppress(dbapi.DatabaseError):
             await cur.execute("DROP TABLE foo")
-            await cur.finish_query()
+            await cur._scroll_stream()
 
         assert not await connection.check_exists("/local/foo")
         with pytest.raises(dbapi.ProgrammingError):
@@ -244,7 +244,7 @@ async def _test_connection(
         await cur.execute(
             "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
         )
-        await cur.finish_query()
+        await cur._scroll_stream()
 
         assert await connection.check_exists("/local/foo")
 
@@ -263,12 +263,12 @@ async def _test_cursor_raw_query(
 
         with suppress(dbapi.DatabaseError):
             await cur.execute("DROP TABLE test")
-            await cur.finish_query()
+            await cur._scroll_stream()
 
         await cur.execute(
             "CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
         )
-        await cur.finish_query()
+        await cur._scroll_stream()
 
         await cur.execute(
             """
@@ -290,7 +290,7 @@ async def _test_cursor_raw_query(
                 )
             },
         )
-        await cur.finish_query()
+        await cur._scroll_stream()
 
         await cur.execute("DROP TABLE test")
 
@@ -307,7 +307,7 @@ async def _test_errors(self, connection: dbapi.AsyncConnection) -> None:
 
         with suppress(dbapi.DatabaseError):
             await cur.execute("DROP TABLE test")
-            await cur.finish_query()
+            await cur._scroll_stream()
 
         with pytest.raises(dbapi.DataError):
             await cur.execute("SELECT 18446744073709551616")
@@ -322,10 +322,10 @@ async def _test_errors(self, connection: dbapi.AsyncConnection) -> None:
             await cur.execute("SELECT * FROM test")
 
         await cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))")
-        await cur.finish_query()
+        await cur._scroll_stream()
 
         await cur.execute("INSERT INTO test(id) VALUES(1)")
-        await cur.finish_query()
+        await cur._scroll_stream()
 
         with pytest.raises(dbapi.IntegrityError):
             await cur.execute("INSERT INTO test(id) VALUES(1)")

tests/test_connection.py

@@ -119,48 +119,19 @@ async def test_cursor_fetch_all(
             assert await cursor.fetchall() == []
 
     @pytest.mark.asyncio
-    async def test_cursor_next_set(
+    async def test_cursor_fetch_one_multiple_result_sets(
         self, session: ydb.aio.QuerySession
     ) -> None:
         async with ydb_dbapi.AsyncCursor(session=session) as cursor:
-            yql_text = """SELECT 1 as val; SELECT 2 as val;"""
-            await cursor.execute(query=yql_text)
-
-            res = await cursor.fetchall()
-            assert res is not None
-            assert len(res) == 1
-            assert res[0][0] == 1
-
-            nextset = await cursor.nextset()
-            assert nextset
-
-            res = await cursor.fetchall()
-            assert res is not None
-            assert len(res) == 1
-            assert res[0][0] == 2
-
-            nextset = await cursor.nextset()
-            assert nextset
-
-            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)
 
+            assert cursor.rowcount == 12
+
             for i in range(RESULT_SET_LENGTH * RESULT_SET_COUNT):
                 res = await cursor.fetchone()
                 assert res is not None
@@ -170,19 +141,19 @@ async def test_cursor_fetch_one_autoscroll(
             assert not await cursor.nextset()
 
     @pytest.mark.asyncio
-    async def test_cursor_fetch_many_autoscroll(
+    async def test_cursor_fetch_many_multiple_result_sets(
         self, session: ydb.aio.QuerySession
     ) -> None:
-        async with ydb_dbapi.AsyncCursor(
-            session=session, auto_scroll_result_sets=True
-        ) as cursor:
+        async with ydb_dbapi.AsyncCursor(session=session) 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)
 
+            assert cursor.rowcount == 12
+
             halfsize = (RESULT_SET_LENGTH * RESULT_SET_COUNT) // 2
             for _ in range(2):
                 res = await cursor.fetchmany(size=halfsize)
@@ -193,19 +164,19 @@ async def test_cursor_fetch_many_autoscroll(
             assert not await cursor.nextset()
 
     @pytest.mark.asyncio
-    async def test_cursor_fetch_all_autoscroll(
+    async def test_cursor_fetch_all_multiple_result_sets(
         self, session: ydb.aio.QuerySession
     ) -> None:
-        async with ydb_dbapi.AsyncCursor(
-            session=session, auto_scroll_result_sets=True
-        ) as cursor:
+        async with ydb_dbapi.AsyncCursor(session=session) 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)
 
+            assert cursor.rowcount == 12
+
             res = await cursor.fetchall()
 
             assert len(res) == RESULT_SET_COUNT * RESULT_SET_LENGTH
@@ -274,45 +245,19 @@ def test_cursor_fetch_all(self, session_sync: ydb.QuerySession) -> None:
 
             assert cursor.fetchall() == []
 
-    def test_cursor_next_set(self, session_sync: ydb.QuerySession) -> None:
-        with ydb_dbapi.Cursor(session=session_sync) as cursor:
-            yql_text = """SELECT 1 as val; SELECT 2 as val;"""
-            cursor.execute(query=yql_text)
-
-            res = cursor.fetchall()
-            assert res is not None
-            assert len(res) == 1
-            assert res[0][0] == 1
-
-            nextset = cursor.nextset()
-            assert nextset
-
-            res = cursor.fetchall()
-            assert res is not None
-            assert len(res) == 1
-            assert res[0][0] == 2
-
-            nextset = cursor.nextset()
-            assert nextset
-
-            assert cursor.fetchall() == []
-
-            nextset = cursor.nextset()
-            assert not nextset
-
-    def test_cursor_fetch_one_autoscroll(
+    def test_cursor_fetch_one_multiple_result_sets(
         self, session_sync: ydb.QuerySession
     ) -> None:
-        with ydb_dbapi.Cursor(
-            session=session_sync, auto_scroll_result_sets=True
-        ) as cursor:
+        with ydb_dbapi.Cursor(session=session_sync) as cursor:
             yql_text = """
             SELECT id, val FROM table;
             SELECT id, val FROM table1;
             SELECT id, val FROM table2;
             """
             cursor.execute(query=yql_text)
 
+            assert cursor.rowcount == 12
+
             for i in range(RESULT_SET_LENGTH * RESULT_SET_COUNT):
                 res = cursor.fetchone()
                 assert res is not None
@@ -321,19 +266,19 @@ def test_cursor_fetch_one_autoscroll(
             assert cursor.fetchone() is None
             assert not cursor.nextset()
 
-    def test_cursor_fetch_many_autoscroll(
+    def test_cursor_fetch_many_multiple_result_sets(
         self, session_sync: ydb.QuerySession
     ) -> None:
-        with ydb_dbapi.Cursor(
-            session=session_sync, auto_scroll_result_sets=True
-        ) as cursor:
+        with ydb_dbapi.Cursor(session=session_sync) as cursor:
             yql_text = """
             SELECT id, val FROM table;
             SELECT id, val FROM table1;
             SELECT id, val FROM table2;
             """
             cursor.execute(query=yql_text)
 
+            assert cursor.rowcount == 12
+
             halfsize = (RESULT_SET_LENGTH * RESULT_SET_COUNT) // 2
             for _ in range(2):
                 res = cursor.fetchmany(size=halfsize)
@@ -343,19 +288,19 @@ def test_cursor_fetch_many_autoscroll(
             assert cursor.fetchmany(2) == []
             assert not cursor.nextset()
 
-    def test_cursor_fetch_all_autoscroll(
+    def test_cursor_fetch_all_multiple_result_sets(
         self, session_sync: ydb.QuerySession
     ) -> None:
-        with ydb_dbapi.Cursor(
-            session=session_sync, auto_scroll_result_sets=True
-        ) as cursor:
+        with ydb_dbapi.Cursor(session=session_sync) as cursor:
             yql_text = """
             SELECT id, val FROM table;
             SELECT id, val FROM table1;
             SELECT id, val FROM table2;
             """
             cursor.execute(query=yql_text)
 
+            assert cursor.rowcount == 12
+
             res = cursor.fetchall()
 
             assert len(res) == RESULT_SET_COUNT * RESULT_SET_LENGTH