Merge pull request #3 from ydb-platform/execute_retries

vgvoleg · web-flow · commit c241e4940616 · 2024-11-01T22:06:06.000+03:00
Add pool retries to execute calls
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -7,7 +7,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.9"
-ydb = "^3.18.3"
+ydb = "^3.18.8"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^4.0.1"
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -163,6 +163,17 @@ async def session_pool(
                 """
             )
 
+            await session_pool.execute_with_retries(
+                f"""
+                DELETE FROM {name};
+                INSERT INTO {name} (id, val) VALUES
+                (0, 0),
+                (1, 1),
+                (2, 2),
+                (3, 3)
+                """
+            )
+
         yield session_pool
 
 
@@ -184,46 +195,15 @@ def session_pool_sync(
                 """
             )
 
-        yield session_pool
-
-
-@pytest.fixture
-async def session(
-    session_pool: ydb.aio.QuerySessionPool,
-) -> AsyncGenerator[ydb.aio.QuerySession]:
-    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
-    await session_pool.release(session)
-
-
-@pytest.fixture
-def session_sync(
-    session_pool_sync: ydb.QuerySessionPool,
-) -> Generator[ydb.QuerySession]:
-    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_pool.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)
+        yield session_pool
diff --git a/tests/test_connections.py b/tests/test_connections.py
@@ -35,6 +35,8 @@ def _test_isolation_level_read_only(
 
         connection.set_isolation_level(isolation_level)
         cursor = connection.cursor()
+        maybe_await(connection.begin())
+
         query = "UPSERT INTO foo(id) VALUES (1)"
         if read_only:
             with pytest.raises(dbapi.DatabaseError):
diff --git a/tests/test_cursors.py b/tests/test_cursors.py
@@ -139,8 +139,10 @@ def _test_cursor_fetch_all_multiple_result_sets(
 
 class TestCursor(BaseCursorTestSuit):
     @pytest.fixture
-    def sync_cursor(self, session_sync: ydb.QuerySession) -> Generator[Cursor]:
-        cursor = Cursor(session_sync, ydb.QuerySerializableReadWrite())
+    def sync_cursor(
+        self, session_pool_sync: ydb.QuerySessionPool
+    ) -> Generator[Cursor]:
+        cursor = Cursor(session_pool_sync, ydb.QuerySerializableReadWrite())
         yield cursor
         cursor.close()
 
@@ -173,9 +175,9 @@ def test_cursor_fetch_all_multiple_result_sets(
 class TestAsyncCursor(BaseCursorTestSuit):
     @pytest.fixture
     async def async_cursor(
-        self, session: ydb.aio.QuerySession
+        self, session_pool: ydb.aio.QuerySessionPool
     ) -> AsyncGenerator[Cursor]:
-        cursor = AsyncCursor(session, ydb.QuerySerializableReadWrite())
+        cursor = AsyncCursor(session_pool, ydb.QuerySerializableReadWrite())
         yield cursor
         await greenlet_spawn(cursor.close)
 
diff --git a/ydb_dbapi/connections.py b/ydb_dbapi/connections.py
@@ -87,12 +87,15 @@ def __init__(
         if ydb_session_pool is not None:
             self._shared_session_pool = True
             self._session_pool = ydb_session_pool
+            settings = self._get_client_settings()
+            self._session_pool._query_client_settings = settings
             self._driver = self._session_pool._driver
         else:
             driver_config = ydb.DriverConfig(
                 endpoint=self.endpoint,
                 database=self.database,
                 credentials=self.credentials,
+                query_client_settings=self._get_client_settings(),
             )
             self._driver = self._driver_cls(driver_config)
             self._session_pool = self._pool_cls(self._driver, size=5)
@@ -126,11 +129,15 @@ def get_isolation_level(self) -> str:
         msg = f"{self._tx_mode.name} is not supported"
         raise NotSupportedError(msg)
 
-    def _maybe_init_tx(
-        self, session: ydb.QuerySession | ydb.aio.QuerySession
-    ) -> None:
-        if self._tx_context is None and self.interactive_transaction:
-            self._tx_context = session.transaction(self._tx_mode)
+    def _get_client_settings(self) -> ydb.QueryClientSettings:
+        return (
+            ydb.QueryClientSettings()
+            .with_native_date_in_result_sets(True)
+            .with_native_datetime_in_result_sets(True)
+            .with_native_timestamp_in_result_sets(True)
+            .with_native_interval_in_result_sets(True)
+            .with_native_json_in_result_sets(False)
+        )
 
 
 class Connection(BaseConnection):
@@ -160,17 +167,11 @@ def __init__(
         self._current_cursor: Cursor | None = None
 
     def cursor(self) -> Cursor:
-        if self._session is None:
-            raise RuntimeError("Connection is not ready, use wait_ready.")
-
-        self._maybe_init_tx(self._session)
-
-        self._current_cursor = self._cursor_cls(
-            session=self._session,
+        return self._cursor_cls(
+            session_pool=self._session_pool,
             tx_mode=self._tx_mode,
             tx_context=self._tx_context,
         )
-        return self._current_cursor
 
     def wait_ready(self, timeout: int = 10) -> None:
         try:
@@ -185,27 +186,33 @@ def wait_ready(self, timeout: int = 10) -> None:
             )
             raise InterfaceError(msg) from e
 
-        self._session = self._session_pool.acquire()
+    @handle_ydb_errors
+    def begin(self) -> None:
+        self._tx_context = None
+        if self.interactive_transaction:
+            self._session = self._session_pool.acquire()
+            self._tx_context = self._session.transaction(self._tx_mode)
 
     @handle_ydb_errors
     def commit(self) -> None:
         if self._tx_context and self._tx_context.tx_id:
             self._tx_context.commit()
+            self._session_pool.release(self._session)
             self._tx_context = None
+            self._session = None
 
     @handle_ydb_errors
     def rollback(self) -> None:
         if self._tx_context and self._tx_context.tx_id:
             self._tx_context.rollback()
+            self._session_pool.release(self._session)
             self._tx_context = None
+            self._session = None
 
     @handle_ydb_errors
     def close(self) -> None:
         self.rollback()
 
-        if self._current_cursor:
-            self._current_cursor.close()
-
         if self._session:
             self._session_pool.release(self._session)
 
@@ -287,17 +294,11 @@ def __init__(
         self._current_cursor: AsyncCursor | None = None
 
     def cursor(self) -> AsyncCursor:
-        if self._session is None:
-            raise RuntimeError("Connection is not ready, use wait_ready.")
-
-        self._maybe_init_tx(self._session)
-
-        self._current_cursor = self._cursor_cls(
-            session=self._session,
+        return self._cursor_cls(
+            session_pool=self._session_pool,
             tx_mode=self._tx_mode,
             tx_context=self._tx_context,
         )
-        return self._current_cursor
 
     async def wait_ready(self, timeout: int = 10) -> None:
         try:
@@ -312,27 +313,33 @@ async def wait_ready(self, timeout: int = 10) -> None:
             )
             raise InterfaceError(msg) from e
 
-        self._session = await self._session_pool.acquire()
+    @handle_ydb_errors
+    async def begin(self) -> None:
+        self._tx_context = None
+        if self.interactive_transaction:
+            self._session = await self._session_pool.acquire()
+            self._tx_context = self._session.transaction(self._tx_mode)
 
     @handle_ydb_errors
     async def commit(self) -> None:
-        if self._tx_context and self._tx_context.tx_id:
+        if self._session and self._tx_context and self._tx_context.tx_id:
             await self._tx_context.commit()
+            await self._session_pool.release(self._session)
+            self._session = None
             self._tx_context = None
 
     @handle_ydb_errors
     async def rollback(self) -> None:
-        if self._tx_context and self._tx_context.tx_id:
+        if self._session and self._tx_context and self._tx_context.tx_id:
             await self._tx_context.rollback()
+            await self._session_pool.release(self._session)
+            self._session = None
             self._tx_context = None
 
     @handle_ydb_errors
     async def close(self) -> None:
         await self.rollback()
 
-        if self._current_cursor:
-            await self._current_cursor.close()
-
         if self._session:
             await self._session_pool.release(self._session)
 
diff --git a/ydb_dbapi/cursors.py b/ydb_dbapi/cursors.py
@@ -138,13 +138,13 @@ def _fetchall_from_buffer(self) -> list:
 class Cursor(BufferedCursor):
     def __init__(
         self,
-        session: ydb.QuerySession,
+        session_pool: ydb.QuerySessionPool,
         tx_mode: ydb.BaseQueryTxMode,
         tx_context: ydb.QueryTxContext | None = None,
         table_path_prefix: str = "",
     ) -> None:
         super().__init__()
-        self._session = session
+        self._session_pool = session_pool
         self._tx_mode = tx_mode
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
@@ -165,19 +165,32 @@ def fetchall(self) -> list:
     def _execute_generic_query(
         self, query: str, parameters: ParametersType | None = None
     ) -> Iterator[ydb.convert.ResultSet]:
-        return self._session.execute(query=query, parameters=parameters)
+        def callee(
+            session: ydb.QuerySession,
+        ) -> Iterator[ydb.convert.ResultSet]:
+            return session.execute(
+                query=query,
+                parameters=parameters,
+            )
+
+        return self._session_pool.retry_operation_sync(callee)
 
     @handle_ydb_errors
     def _execute_session_query(
         self,
         query: str,
         parameters: ParametersType | None = None,
     ) -> Iterator[ydb.convert.ResultSet]:
-        return self._session.transaction(self._tx_mode).execute(
-            query=query,
-            parameters=parameters,
-            commit_tx=True,
-        )
+        def callee(
+            session: ydb.QuerySession,
+        ) -> Iterator[ydb.convert.ResultSet]:
+            return session.transaction(self._tx_mode).execute(
+                query=query,
+                parameters=parameters,
+                commit_tx=True,
+            )
+
+        return self._session_pool.retry_operation_sync(callee)
 
     @handle_ydb_errors
     def _execute_transactional_query(
@@ -276,13 +289,13 @@ def __exit__(
 class AsyncCursor(BufferedCursor):
     def __init__(
         self,
-        session: ydb.aio.QuerySession,
+        session_pool: ydb.aio.QuerySessionPool,
         tx_mode: ydb.BaseQueryTxMode,
         tx_context: ydb.aio.QueryTxContext | None = None,
         table_path_prefix: str = "",
     ) -> None:
         super().__init__()
-        self._session = session
+        self._session_pool = session_pool
         self._tx_mode = tx_mode
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
@@ -303,19 +316,32 @@ async def fetchall(self) -> list:
     async def _execute_generic_query(
         self, query: str, parameters: ParametersType | None = None
     ) -> AsyncIterator[ydb.convert.ResultSet]:
-        return await self._session.execute(query=query, parameters=parameters)
+        async def callee(
+            session: ydb.aio.QuerySession,
+        ) -> AsyncIterator[ydb.convert.ResultSet]:
+            return await session.execute(
+                query=query,
+                parameters=parameters,
+            )
+
+        return await self._session_pool.retry_operation_async(callee)
 
     @handle_ydb_errors
     async def _execute_session_query(
         self,
         query: str,
         parameters: ParametersType | None = None,
     ) -> AsyncIterator[ydb.convert.ResultSet]:
-        return await self._session.transaction(self._tx_mode).execute(
-            query=query,
-            parameters=parameters,
-            commit_tx=True,
-        )
+        async def callee(
+            session: ydb.aio.QuerySession,
+        ) -> AsyncIterator[ydb.convert.ResultSet]:
+            return await session.transaction(self._tx_mode).execute(
+                query=query,
+                parameters=parameters,
+                commit_tx=True,
+            )
+
+        return await self._session_pool.retry_operation_async(callee)
 
     @handle_ydb_errors
     async def _execute_transactional_query(
