query session pool async

Author: vgvoleg

tests/aio/query/conftest.py

@@ -1,5 +1,6 @@
 import pytest
 from ydb.aio.query.session import QuerySessionAsync
+from ydb.aio.query.pool import QuerySessionPoolAsync
 
 
 @pytest.fixture
@@ -25,3 +26,9 @@ async def tx(session):
         await transaction.rollback()
     except BaseException:
         pass
+
+
+@pytest.fixture
+def pool(driver):
+    pool = QuerySessionPoolAsync(driver)
+    yield pool

tests/aio/query/test_query_session_pool.py

@@ -0,0 +1,56 @@
+import pytest
+import ydb
+from ydb.aio.query.pool import QuerySessionPoolAsync
+from ydb.aio.query.session import QuerySessionAsync, QuerySessionStateEnum
+
+
+class TestQuerySessionPoolAsync:
+    @pytest.mark.asyncio
+    async def test_checkout_provides_created_session(self, pool: QuerySessionPoolAsync):
+        async with pool.checkout() as session:
+            assert session._state._state == QuerySessionStateEnum.CREATED
+
+        assert session._state._state == QuerySessionStateEnum.CLOSED
+
+    @pytest.mark.asyncio
+    async def test_oneshot_query_normal(self, pool: QuerySessionPoolAsync):
+        res = await pool.execute_with_retries("select 1;")
+        assert len(res) == 1
+
+    @pytest.mark.asyncio
+    async def test_oneshot_ddl_query(self, pool: QuerySessionPoolAsync):
+        await pool.execute_with_retries("create table Queen(key UInt64, PRIMARY KEY (key));")
+        await pool.execute_with_retries("drop table Queen;")
+
+    @pytest.mark.asyncio
+    async def test_oneshot_query_raises(self, pool: QuerySessionPoolAsync):
+        with pytest.raises(ydb.GenericError):
+            await pool.execute_with_retries("Is this the real life? Is this just fantasy?")
+
+    @pytest.mark.asyncio
+    async def test_retry_op_uses_created_session(self, pool: QuerySessionPoolAsync):
+        async def callee(session: QuerySessionAsync):
+            assert session._state._state == QuerySessionStateEnum.CREATED
+
+        await pool.retry_operation_async(callee)
+
+    @pytest.mark.asyncio
+    async def test_retry_op_normal(self, pool: QuerySessionPoolAsync):
+        async def callee(session: QuerySessionAsync):
+            async with session.transaction() as tx:
+                iterator = await tx.execute("select 1;", commit_tx=True)
+                return [result_set async for result_set in iterator]
+
+        res = await pool.retry_operation_async(callee)
+        assert len(res) == 1
+
+    @pytest.mark.asyncio
+    async def test_retry_op_raises(self, pool: QuerySessionPoolAsync):
+        class CustomException(Exception):
+            pass
+
+        async def callee(session: QuerySessionAsync):
+            raise CustomException()
+
+        with pytest.raises(CustomException):
+            await pool.retry_operation_async(callee)

ydb/aio/query/pool.py

@@ -0,0 +1,91 @@
+import logging
+from typing import (
+    Callable,
+    Optional,
+    List,
+)
+
+from ...query import base
+from .session import (
+    QuerySessionAsync,
+)
+from ...retries import (
+    RetrySettings,
+    retry_operation_async,
+)
+from ... import convert
+
+logger = logging.getLogger(__name__)
+
+
+class QuerySessionPoolAsync:
+    """QuerySessionPoolAsync is an object to simplify operations with sessions of Query Service."""
+
+    def __init__(self, driver: base.SupportedDriverType):
+        """
+        :param driver: A driver instance
+        """
+
+        logger.warning("QuerySessionPoolAsync is an experimental API, which could be changed.")
+        self._driver = driver
+
+    def checkout(self) -> "SimpleQuerySessionCheckoutAsync":
+        """WARNING: This API is experimental and could be changed.
+        Return a Session context manager, that opens session on enter and closes session on exit.
+        """
+
+        return SimpleQuerySessionCheckoutAsync(self)
+
+    async def retry_operation_async(self, callee: Callable, retry_settings: Optional[RetrySettings] = None, *args, **kwargs):
+        """WARNING: This API is experimental and could be changed.
+        Special interface to execute a bunch of commands with session in a safe, retriable way.
+
+        :param callee: A function, that works with session.
+        :param retry_settings: RetrySettings object.
+
+        :return: Result sets or exception in case of execution errors.
+        """
+
+        retry_settings = RetrySettings() if retry_settings is None else retry_settings
+
+        async def wrapped_callee():
+            async with self.checkout() as session:
+                return await callee(session, *args, **kwargs)
+
+        return await retry_operation_async(wrapped_callee, retry_settings)
+
+    async def execute_with_retries(
+        self, query: str, retry_settings: Optional[RetrySettings] = None, *args, **kwargs
+    ) -> List[convert.ResultSet]:
+        """WARNING: This API is experimental and could be changed.
+        Special interface to execute a one-shot queries in a safe, retriable way.
+        Note: this method loads all data from stream before return, do not use this
+        method with huge read queries.
+
+        :param query: A query, yql or sql text.
+        :param retry_settings: RetrySettings object.
+
+        :return: Result sets or exception in case of execution errors.
+        """
+
+        retry_settings = RetrySettings() if retry_settings is None else retry_settings
+
+        async def wrapped_callee():
+            async with self.checkout() as session:
+                it = await session.execute(query, *args, **kwargs)
+                return [result_set async for result_set in it]
+
+        return await retry_operation_async(wrapped_callee, retry_settings)
+
+
+class SimpleQuerySessionCheckoutAsync:
+    def __init__(self, pool: QuerySessionPoolAsync):
+        self._pool = pool
+        self._session = QuerySessionAsync(pool._driver)
+
+    async def __aenter__(self) -> base.IQuerySession:
+        await self._session.create()
+        return self._session
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._session.delete()

ydb/retries.py

@@ -1,3 +1,4 @@
+import asyncio
 import random
 import time
 
@@ -134,3 +135,27 @@ def retry_operation_sync(callee, retry_settings=None, *args, **kwargs):
             time.sleep(next_opt.timeout)
         else:
             return next_opt.result
+
+
+async def retry_operation_async(callee, retry_settings=None, *args, **kwargs):  # pylint: disable=W1113
+    """
+    The retry operation helper can be used to retry a coroutine that raises YDB specific
+    exceptions.
+
+    :param callee: A coroutine to retry.
+    :param retry_settings: An instance of ydb.RetrySettings that describes how the coroutine
+    should be retried. If None, default instance of retry settings will be used.
+    :param args: A tuple with positional arguments to be passed into the coroutine.
+    :param kwargs: A dictionary with keyword arguments to be passed into the coroutine.
+
+    Returns awaitable result of coroutine. If retries are not succussful exception is raised.
+    """
+    opt_generator = retry_operation_impl(callee, retry_settings, *args, **kwargs)
+    for next_opt in opt_generator:
+        if isinstance(next_opt, YdbRetryOperationSleepOpt):
+            await asyncio.sleep(next_opt.timeout)
+        else:
+            try:
+                return await next_opt.result
+            except BaseException as e:  # pylint: disable=W0703
+                next_opt.set_exception(e)