cursor methods

Author: vgvoleg

.github/workflows/tests.yml

@@ -44,4 +44,4 @@ jobs:
 
       - name: Run tests
         run: |
-          poetry run -- pytest tests
+          poetry run pytest tests

poetry.lock

@@ -36,6 +36,7 @@ ruff = "ruff check"
 lint = ["mypy", "ruff"]
 format-check = "ruff format --check"
 format = "ruff format"
+tests = "pytest -v  --docker-compose-remove-volumes --docker-compose=docker-compose.yml"
 
 [tool.ruff]
 exclude = [".venv", ".git", "__pycache__", "build", "dist", "venv"]
@@ -48,6 +49,7 @@ select = ["F", "E"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
+addopts = "-p no:warnings"
 
 [[tool.mypy.overrides]]
 module = "ydb.*"

pyproject.toml

@@ -55,3 +55,23 @@ async def driver(endpoint, database, event_loop):
     yield driver
 
     await driver.stop(timeout=10)
+
+
+@pytest.fixture
+async def session_pool(driver: ydb.aio.Driver, event_loop):
+    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)
+            )
+            """
+        )
+
+        yield session_pool

tests/conftest.py

@@ -1,33 +1,125 @@
 import pytest
-
-import ydb
-
 import ydb_dbapi
 import ydb_dbapi.cursors
 
 
-@pytest.fixture
-async def cursor(driver: ydb.aio.Driver):
-    session_pool = ydb.aio.QuerySessionPool(driver)
-    session = await session_pool.acquire()
-    cursor = ydb_dbapi.cursors.Cursor(session_pool, session)
-
-    yield cursor
-
-    await session_pool.release(session)
-
-
 @pytest.mark.asyncio
-async def test_cursor_ddl(cursor):
-    op = ydb_dbapi.cursors.YdbQuery(
-        yql_text="""
+async def test_cursor_ddl(session_pool):
+    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+
+    yql = """
         CREATE TABLE table (
         id Int64 NOT NULL,
-        i64Val Int64,
+        val Int64,
         PRIMARY KEY(id)
         )
-        """,
-        is_ddl=True,
-    )
+    """
+
+    with pytest.raises(ydb_dbapi.Error):
+        await cursor.execute(query=yql)
+
+    yql = """
+    DROP TABLE table
+    """
+
+    await cursor.execute(query=yql)
+
+    assert await cursor.fetchone() is None
+
+
+@pytest.mark.asyncio
+async def test_cursor_dml(session_pool):
+    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    yql_text = """
+    INSERT INTO table (id, val) VALUES
+    (1, 1),
+    (2, 2),
+    (3, 3)
+    """
+
+    await cursor.execute(query=yql_text)
+    assert await cursor.fetchone() is None
+
+    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+
+    yql_text = """
+    SELECT COUNT(*) FROM table as sum
+    """
+
+    await cursor.execute(query=yql_text)
+
+    res = await cursor.fetchone()
+    assert len(res) == 1
+    assert res[0] == 3
+
+
+@pytest.mark.asyncio
+async def test_cursor_fetch_one(session_pool):
+    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    yql_text = """
+    INSERT INTO table (id, val) VALUES
+    (1, 1),
+    (2, 2)
+    """
+
+    await cursor.execute(query=yql_text)
+    assert await cursor.fetchone() is None
+
+    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+
+    yql_text = """
+    SELECT id, val FROM table
+    """
+
+    await cursor.execute(query=yql_text)
+
+    res = await cursor.fetchone()
+    assert res[0] == 1
+
+    res = await cursor.fetchone()
+    assert res[0] == 2
+
+    assert await cursor.fetchone() is None
+
+
+@pytest.mark.asyncio
+async def test_cursor_fetch_many(session_pool):
+    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+    yql_text = """
+    INSERT INTO table (id, val) VALUES
+    (1, 1),
+    (2, 2),
+    (3, 3),
+    (4, 4)
+    """
 
-    await cursor.execute(op)
+    await cursor.execute(query=yql_text)
+    assert await cursor.fetchone() is None
+
+    cursor = ydb_dbapi.cursors.Cursor(session_pool=session_pool)
+
+    yql_text = """
+    SELECT id, val FROM table
+    """
+
+    await cursor.execute(query=yql_text)
+
+    res = await cursor.fetchmany()
+    assert len(res) == 1
+    assert res[0][0] == 1
+
+    res = await cursor.fetchmany(size=2)
+    assert len(res) == 2
+    assert res[0][0] == 2
+    assert res[1][0] == 3
+
+    res = await cursor.fetchmany(size=2)
+    assert len(res) == 1
+    assert res[0][0] == 4
+
+    assert await cursor.fetchmany(size=2) is None
+
+
+@pytest.mark.asyncio
+async def test_cursor_fetch_all(session_pool):
+    pass

tests/test_cursor.py

@@ -0,0 +1 @@
+from .errors import *  # noqa

ydb_dbapi/__init__.py

@@ -12,7 +12,7 @@
 )
 
 import ydb
-from .errors import DatabaseError
+from .errors import Error, DatabaseError
 from .utils import handle_ydb_errors, AsyncFromSyncIterator
 
 
@@ -40,37 +40,37 @@ class Cursor:
     def __init__(
         self,
         session_pool: ydb.aio.QuerySessionPool,
-        session: ydb.aio.QuerySession,
-        tx_mode: Optional[ydb.aio.QueryTxContext] = None,
         tx_context: Optional[ydb.aio.QueryTxContext] = None,
         table_path_prefix: str = "",
         autocommit: bool = True,
     ):
         self.arraysize: int = 1
         self._description: Optional[List[Tuple]] = None
 
-        self._pool = session_pool
-        self._session = session
-        self._tx_mode = tx_mode
-        self._tx_context: ydb.aio.QueryTxContext = tx_context
+        self._session_pool = session_pool
+        self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
         self._autocommit = autocommit
 
         self._stream: Optional[AsyncIterator] = None
         self._rows: Optional[Iterator[Dict]] = None
 
     @handle_ydb_errors
-    async def _execute_ddl_query(
+    async def _execute_generic_query(
         self, query: str, parameters: Optional[ParametersType] = None
     ) -> List[ydb.convert.ResultSet]:
-        return await self._pool.execute_with_retries(
+        return await self._session_pool.execute_with_retries(
             query=query, parameters=parameters
         )
 
     @handle_ydb_errors
-    async def _execute_dml_query(
+    async def _execute_transactional_query(
         self, query: str, parameters: Optional[ParametersType] = None
     ) -> AsyncIterator:
+        if self._tx_context is None:
+            raise Error(
+                "Unable to execute tx based queries without transaction."
+            )
         return await self._tx_context.execute(
             query=query,
             parameters=parameters,
@@ -79,17 +79,17 @@ async def _execute_dml_query(
 
     @handle_ydb_errors
     async def execute(
-        self, operation: YdbQuery, parameters: Optional[ParametersType] = None
+        self, query: str, parameters: Optional[ParametersType] = None
     ):
-        if operation.is_ddl:
-            result_sets = await self._execute_ddl_query(
-                query=operation.yql_text, parameters=parameters
+        if self._tx_context is not None:
+            self._stream = await self._execute_transactional_query(
+                query=query, parameters=parameters
             )
-            self._stream = AsyncFromSyncIterator(iter(result_sets))
         else:
-            self._stream = await self._execute_dml_query(
-                query=operation.yql_text, parameters=parameters
+            result_sets = await self._execute_generic_query(
+                query=query, parameters=parameters
             )
+            self._stream = AsyncFromSyncIterator(iter(result_sets))
 
         if self._stream is None:
             return
@@ -98,7 +98,6 @@ async def execute(
         self._update_result_set(result_set)
 
     def _update_result_set(self, result_set: ydb.convert.ResultSet):
-        # self._result_set = result_set
         self._update_description(result_set)
         self._rows = self._rows_iterable(result_set)
 
@@ -133,8 +132,13 @@ async def fetchone(self):
         return next(self._rows or iter([]), None)
 
     async def fetchmany(self, size: Optional[int] = None):
-        return list(
-            itertools.islice(self._rows or iter([]), size or self.arraysize)
+        return (
+            list(
+                itertools.islice(
+                    self._rows or iter([]), size or self.arraysize
+                )
+            )
+            or None
         )
 
     async def fetchall(self):

ydb_dbapi/cursors.py

@@ -15,9 +15,9 @@
 
 def handle_ydb_errors(func):
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
+    async def wrapper(*args, **kwargs):
         try:
-            return func(*args, **kwargs)
+            return await func(*args, **kwargs)
         except (ydb.issues.AlreadyExists, ydb.issues.PreconditionFailed) as e:
             raise IntegrityError(e.message, original_error=e) from e
         except (ydb.issues.Unsupported, ydb.issues.Unimplemented) as e: