fix all tests

Author: vgvoleg

poetry.lock

@@ -49,7 +49,7 @@ select = ["F", "E"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
-addopts = "-p no:warnings"
+addopts = "--tb native -v -r fxX -p no:warnings"
 
 [[tool.mypy.overrides]]
 module = "ydb.*"

pyproject.toml

@@ -55,10 +55,11 @@ async def driver(endpoint, database, event_loop):
     yield driver
 
     await driver.stop(timeout=10)
+    del driver
 
 
 @pytest.fixture
-async def session_pool(driver: ydb.aio.Driver, event_loop):
+async def session_pool(driver: ydb.aio.Driver):
     session_pool = ydb.aio.QuerySessionPool(driver)
     async with session_pool:
         await session_pool.execute_with_retries(

tests/conftest.py

@@ -14,26 +14,32 @@ async def _test_isolation_level_read_only(
         isolation_level: str,
         read_only: bool,
     ):
-        await connection.cursor().execute(
-            "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
-        )
-        connection.set_isolation_level(isolation_level)
+        async with connection.cursor() as cursor:
+            with suppress(dbapi.DatabaseError):
+                await cursor.execute("DROP TABLE foo")
 
-        cursor = connection.cursor()
+        async with connection.cursor() as cursor:
+            await cursor.execute(
+                "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
+            )
+
+        connection.set_isolation_level(isolation_level)
 
         await connection.begin()
 
-        query = "UPSERT INTO foo(id) VALUES (1)"
-        if read_only:
-            with pytest.raises(dbapi.DatabaseError):
+        async with connection.cursor() as cursor:
+            query = "UPSERT INTO foo(id) VALUES (1)"
+            if read_only:
+                with pytest.raises(dbapi.DatabaseError):
+                    await cursor.execute(query)
+                    await cursor.finish_query()
+            else:
                 await cursor.execute(query)
-        else:
-            await cursor.execute(query)
 
         await connection.rollback()
 
-        await connection.cursor().execute("DROP TABLE foo")
-        await connection.cursor().close()
+        async with connection.cursor() as cursor:
+            cursor.execute("DROP TABLE foo")
 
     async def _test_connection(self, connection: dbapi.Connection):
         await connection.commit()
@@ -42,6 +48,7 @@ async def _test_connection(self, connection: dbapi.Connection):
         cur = connection.cursor()
         with suppress(dbapi.DatabaseError):
             await cur.execute("DROP TABLE foo")
+            await cur.finish_query()
 
         assert not await connection.check_exists("/local/foo")
         with pytest.raises(dbapi.ProgrammingError):
@@ -50,6 +57,7 @@ async def _test_connection(self, connection: dbapi.Connection):
         await cur.execute(
             "CREATE TABLE foo(id Int64 NOT NULL, PRIMARY KEY (id))"
         )
+        await cur.finish_query()
 
         assert await connection.check_exists("/local/foo")
 
@@ -66,10 +74,12 @@ async def _test_cursor_raw_query(self, connection: dbapi.Connection):
 
         with suppress(dbapi.DatabaseError):
             await cur.execute("DROP TABLE test")
+            await cur.finish_query()
 
         await cur.execute(
             "CREATE TABLE test(id Int64 NOT NULL, text Utf8, PRIMARY KEY (id))"
         )
+        await cur.finish_query()
 
         await cur.execute(
             """
@@ -91,6 +101,7 @@ async def _test_cursor_raw_query(self, connection: dbapi.Connection):
                 )
             },
         )
+        await cur.finish_query()
 
         await cur.execute("DROP TABLE test")
 
@@ -104,6 +115,7 @@ async def _test_errors(self, connection: dbapi.Connection):
 
         with suppress(dbapi.DatabaseError):
             await cur.execute("DROP TABLE test")
+            await cur.finish_query()
 
         with pytest.raises(dbapi.DataError):
             await cur.execute("SELECT 18446744073709551616")
@@ -118,8 +130,11 @@ async def _test_errors(self, connection: dbapi.Connection):
             await cur.execute("SELECT * FROM test")
 
         await cur.execute("CREATE TABLE test(id Int64, PRIMARY KEY (id))")
+        await cur.finish_query()
 
         await cur.execute("INSERT INTO test(id) VALUES(1)")
+        await cur.finish_query()
+
         with pytest.raises(dbapi.IntegrityError):
             await cur.execute("INSERT INTO test(id) VALUES(1)")
 
@@ -143,10 +158,10 @@ async def connection(self, endpoint, database):
         [
             (dbapi.IsolationLevel.SERIALIZABLE, False),
             (dbapi.IsolationLevel.AUTOCOMMIT, False),
-            # (dbapi.IsolationLevel.ONLINE_READONLY, True),
-            # (dbapi.IsolationLevel.ONLINE_READONLY_INCONSISTENT, True),
-            # (dbapi.IsolationLevel.STALE_READONLY, True),
-            # (dbapi.IsolationLevel.SNAPSHOT_READONLY, True),
+            (dbapi.IsolationLevel.ONLINE_READONLY, True),
+            (dbapi.IsolationLevel.ONLINE_READONLY_INCONSISTENT, True),
+            (dbapi.IsolationLevel.STALE_READONLY, True),
+            (dbapi.IsolationLevel.SNAPSHOT_READONLY, True),
         ],
     )
     async def test_isolation_level_read_only(

tests/test_connection.py

@@ -68,6 +68,9 @@ def __init__(
         self._tx_context: Optional[ydb.QueryTxContext] = None
         self._tx_mode: ydb.BaseQueryTxMode = ydb.QuerySerializableReadWrite()
 
+        self._current_cursor: Optional[Cursor] = None
+        self.interactive_transaction: bool = False
+
     async def _wait(self, timeout: int = 5):
         try:
             await self._driver.wait(timeout, fail_fast=True)
@@ -81,15 +84,22 @@ async def _wait(self, timeout: int = 5):
             ) from e
 
     def cursor(self):
-        return Cursor(
-            session_pool=self._session_pool, tx_context=self._tx_context
+        if self._current_cursor and not self._current_cursor._closed:
+            raise RuntimeError(
+                "Unable to create new Cursor before closing existing one."
+            )
+        self._current_cursor = Cursor(
+            session_pool=self._session_pool,
+            tx_context=self._tx_context,
+            autocommit=(not self.interactive_transaction),
         )
+        return self._current_cursor
 
     async def begin(self):
         self._tx_context = None
         self._session = await self._session_pool.acquire()
         self._tx_context = self._session.transaction(self._tx_mode)
-        await self._tx_context.begin()
+        # await self._tx_context.begin()
 
     async def commit(self):
         if self._tx_context and self._tx_context.tx_id:
@@ -107,6 +117,10 @@ async def rollback(self):
 
     async def close(self):
         await self.rollback()
+
+        if self._current_cursor:
+            await self._current_cursor.close()
+
         if not self._shared_session_pool:
             await self._session_pool.stop()
             await self._driver.stop()

ydb_dbapi/connection.py

@@ -1,4 +1,3 @@
-import dataclasses
 import itertools
 from typing import (
     Any,
@@ -12,8 +11,8 @@
 )
 
 import ydb
-from .errors import Error, DatabaseError
-from .utils import handle_ydb_errors, AsyncFromSyncIterator
+from .errors import Error, DatabaseError, InterfaceError, ProgrammingError
+from .utils import handle_ydb_errors, AsyncFromSyncIterator, CursorStatus
 
 
 ParametersType = Dict[
@@ -26,12 +25,6 @@
 ]
 
 
-@dataclasses.dataclass
-class YdbQuery:
-    yql_text: str
-    is_ddl: bool = False
-
-
 def _get_column_type(type_obj: Any) -> str:
     return str(ydb.convert.type_to_native(type_obj))
 
@@ -56,13 +49,26 @@ def __init__(
         self._rows: Optional[Iterator[Dict]] = None
         self._rows_count: int = -1
 
+        self._closed: bool = False
+        self._state = CursorStatus.ready
+
+    @property
+    def description(self):
+        return self._description
+
+    @property
+    def rowcount(self):
+        return self._rows_count
+
+    @handle_ydb_errors
     async def _execute_generic_query(
         self, query: str, parameters: Optional[ParametersType] = None
     ) -> List[ydb.convert.ResultSet]:
         return await self._session_pool.execute_with_retries(
             query=query, parameters=parameters
         )
 
+    @handle_ydb_errors
     async def _execute_transactional_query(
         self, query: str, parameters: Optional[ParametersType] = None
     ) -> AsyncIterator:
@@ -76,10 +82,14 @@ async def _execute_transactional_query(
             commit_tx=self._autocommit,
         )
 
-    @handle_ydb_errors
     async def execute(
-        self, query: str, parameters: Optional[ParametersType] = None
+        self,
+        query: str,
+        parameters: Optional[ParametersType] = None,
+        prefetch_first_set: bool = True,
     ):
+        self._check_cursor_closed()
+        self._check_pending_query()
         if self._tx_context is not None:
             self._stream = await self._execute_transactional_query(
                 query=query, parameters=parameters
@@ -93,8 +103,10 @@ async def execute(
         if self._stream is None:
             return
 
-        result_set = await self._stream.__anext__()
-        self._update_result_set(result_set)
+        self._begin_query()
+
+        if prefetch_first_set:
+            await self.nextset()
 
     def _update_result_set(self, result_set: ydb.convert.ResultSet):
         self._update_description(result_set)
@@ -144,31 +156,65 @@ async def fetchmany(self, size: Optional[int] = None):
     async def fetchall(self):
         return list(self._rows or iter([])) or None
 
+    @handle_ydb_errors
     async def nextset(self):
         if self._stream is None:
             return False
         try:
             result_set = await self._stream.__anext__()
             self._update_result_set(result_set)
-        except (StopIteration, RuntimeError):
+        except (StopIteration, StopAsyncIteration, RuntimeError):
+            self._state = CursorStatus.finished
             return False
+        except ydb.Error as e:
+            self._state = CursorStatus.finished
+            raise e
         return True
 
+    async def finish_query(self):
+        self._check_cursor_closed()
+
+        if not self._state == CursorStatus.running:
+            return
+
+        next_set_available = True
+        while next_set_available:
+            next_set_available = await self.nextset()
+
+        self._state = CursorStatus.finished
+
     def setinputsizes(self):
         pass
 
     def setoutputsize(self):
         pass
 
     async def close(self):
-        next_set_available = True
-        while next_set_available:
-            next_set_available = await self.nextset()
+        if self._closed:
+            return
 
-    @property
-    def description(self):
-        return self._description
+        await self.finish_query()
+        self._state = CursorStatus.closed
+        self._closed = True
 
-    @property
-    def rowcount(self):
-        return self._rows_count
+    def _begin_query(self):
+        self._state = CursorStatus.running
+
+    def _check_pending_query(self):
+        if self._state == CursorStatus.running:
+            raise ProgrammingError(
+                "Some records have not been fetched. "
+                "Fetch the remaining records before executing the next query."
+            )
+
+    def _check_cursor_closed(self):
+        if self._state == CursorStatus.closed:
+            raise InterfaceError(
+                "Could not perform operation: Cursor is closed."
+            )
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.close()

ydb_dbapi/cursors.py

@@ -1,3 +1,4 @@
+from enum import Enum
 import functools
 from typing import Iterator
 import ydb
@@ -50,6 +51,13 @@ async def wrapper(*args, **kwargs):
     return wrapper
 
 
+class CursorStatus(str, Enum):
+    ready = "ready"
+    running = "running"
+    finished = "finished"
+    closed = "closed"
+
+
 class AsyncFromSyncIterator:
     def __init__(self, sync_iter: Iterator):
         self._sync_iter = sync_iter