aio query session

Author: vgvoleg

tests/aio/query/__init__.py

@@ -0,0 +1,13 @@
+import pytest
+from ydb.aio.query.session import QuerySessionAsync
+
+@pytest.fixture
+async def session(driver):
+    session = QuerySessionAsync(driver)
+
+    yield session
+
+    try:
+        await session.delete()
+    except BaseException:
+        pass

tests/aio/query/conftest.py

@@ -0,0 +1,110 @@
+import pytest
+from ydb.aio.query.session import QuerySessionAsync
+
+
+def _check_session_state_empty(session: QuerySessionAsync):
+    assert session._state.session_id is None
+    assert session._state.node_id is None
+    assert not session._state.attached
+
+
+def _check_session_state_full(session: QuerySessionAsync):
+    assert session._state.session_id is not None
+    assert session._state.node_id is not None
+    assert session._state.attached
+
+
+class TestAsyncQuerySession:
+    @pytest.mark.asyncio
+    async def test_session_normal_lifecycle(self, session: QuerySessionAsync):
+        _check_session_state_empty(session)
+
+        await session.create()
+        _check_session_state_full(session)
+
+        await session.delete()
+        _check_session_state_empty(session)
+
+    @pytest.mark.asyncio
+    async def test_second_create_do_nothing(self, session: QuerySessionAsync):
+        await session.create()
+        _check_session_state_full(session)
+
+        session_id_before = session._state.session_id
+        node_id_before = session._state.node_id
+
+        await session.create()
+        _check_session_state_full(session)
+
+        assert session._state.session_id == session_id_before
+        assert session._state.node_id == node_id_before
+
+    @pytest.mark.asyncio
+    async def test_second_delete_do_nothing(self, session: QuerySessionAsync):
+        await session.create()
+
+        await session.delete()
+        await session.delete()
+
+    @pytest.mark.asyncio
+    async def test_delete_before_create_not_possible(self, session: QuerySessionAsync):
+        with pytest.raises(RuntimeError):
+            await session.delete()
+
+    @pytest.mark.asyncio
+    async def test_create_after_delete_not_possible(self, session: QuerySessionAsync):
+        await session.create()
+        await session.delete()
+        with pytest.raises(RuntimeError):
+            await session.create()
+
+    # def test_transaction_before_create_raises(self, session: QuerySessionAsync):
+    #     with pytest.raises(RuntimeError):
+    #         session.transaction()
+
+    # def test_transaction_after_delete_raises(self, session: QuerySessionAsync):
+    #     session.create()
+
+    #     session.delete()
+
+    #     with pytest.raises(RuntimeError):
+    #         session.transaction()
+
+    # def test_transaction_after_create_not_raises(self, session: QuerySessionAsync):
+    #     session.create()
+    #     session.transaction()
+
+    @pytest.mark.asyncio
+    async def test_execute_before_create_raises(self, session: QuerySessionAsync):
+        with pytest.raises(RuntimeError):
+            await session.execute("select 1;")
+
+    @pytest.mark.asyncio
+    async def test_execute_after_delete_raises(self, session: QuerySessionAsync):
+        await session.create()
+        await session.delete()
+        with pytest.raises(RuntimeError):
+            await session.execute("select 1;")
+
+    @pytest.mark.asyncio
+    async def test_basic_execute(self, session: QuerySessionAsync):
+        await session.create()
+        it = await session.execute("select 1;")
+        result_sets = [result_set async for result_set in it]
+
+        assert len(result_sets) == 1
+        assert len(result_sets[0].rows) == 1
+        assert len(result_sets[0].columns) == 1
+        assert list(result_sets[0].rows[0].values()) == [1]
+
+    @pytest.mark.asyncio
+    async def test_two_results(self, session: QuerySessionAsync):
+        await session.create()
+        res = []
+
+        async with await session.execute("select 1; select 2") as results:
+            async for result_set in results:
+                if len(result_set.rows) > 0:
+                    res.append(list(result_set.rows[0].values()))
+
+        assert res == [[1], [2]]

tests/aio/query/test_query_session.py

@@ -0,0 +1,132 @@
+import asyncio
+
+from typing import (
+    Optional,
+)
+
+from .. import _utilities
+from ... import issues
+from ..._grpc.grpcwrapper import common_utils
+from ...query import base
+from ...query.session import (
+    BaseQuerySession,
+    QuerySessionStateEnum,
+)
+
+
+class AsyncResponseContextIterator(_utilities.AsyncResponseIterator):
+    async def __aenter__(self) -> "AsyncResponseContextIterator":
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        async for _ in self:
+            pass
+
+
+class QuerySessionAsync(BaseQuerySession):
+    """Session object for Query Service. It is not recommended to control
+    session's lifecycle manually - use a QuerySessionPool is always a better choise.
+    """
+
+    _loop: asyncio.AbstractEventLoop
+    _status_stream: _utilities.AsyncResponseIterator = None
+
+    def __init__(self, driver: base.SupportedDriverType, settings: Optional[base.QueryClientSettings] = None):
+        super(QuerySessionAsync, self).__init__(driver, settings)
+        self._loop = asyncio.get_running_loop()
+
+    async def _attach(self) -> None:
+        self._stream = await self._attach_call()
+        self._status_stream = _utilities.AsyncResponseIterator(
+            self._stream,
+            lambda response: common_utils.ServerStatus.from_proto(response),
+        )
+
+        first_response = await self._status_stream.next()
+        if first_response.status != issues.StatusCode.SUCCESS:
+            pass
+
+        self._state.set_attached(True)
+        self._state._change_state(QuerySessionStateEnum.CREATED)
+
+        self._loop.create_task(self._check_session_status_loop(), name="check session status task")
+
+
+    async def _check_session_status_loop(self) -> None:
+        try:
+            async for status in self._status_stream:
+                if status.status != issues.StatusCode.SUCCESS:
+                    self._state.reset()
+                    self._state._change_state(QuerySessionStateEnum.CLOSED)
+        except Exception:
+            pass
+
+    async def delete(self) -> None:
+        """WARNING: This API is experimental and could be changed.
+
+        Deletes a Session of Query Service on server side and releases resources.
+
+        :return: None
+        """
+        if self._state._already_in(QuerySessionStateEnum.CLOSED):
+            return
+
+        self._state._check_invalid_transition(QuerySessionStateEnum.CLOSED)
+        await self._delete_call()
+        self._stream.cancel()
+
+    async def create(self) -> "QuerySessionAsync":
+        """WARNING: This API is experimental and could be changed.
+
+        Creates a Session of Query Service on server side and attaches it.
+
+        :return: QuerySessionSync object.
+        """
+        if self._state._already_in(QuerySessionStateEnum.CREATED):
+            return
+
+        self._state._check_invalid_transition(QuerySessionStateEnum.CREATED)
+        await self._create_call()
+        await self._attach()
+
+        return self
+
+    async def transaction(self, tx_mode) -> base.IQueryTxContext:
+        return super().transaction(tx_mode)
+
+    async def execute(
+        self,
+        query: str,
+        parameters: dict = None,
+        commit_tx: bool = False,
+        syntax: base.QuerySyntax = None,
+        exec_mode: base.QueryExecMode = None,
+        concurrent_result_sets: bool = False,
+    ) -> AsyncResponseContextIterator:
+        """WARNING: This API is experimental and could be changed.
+
+        Sends a query to Query Service
+        :param query: (YQL or SQL text) to be executed.
+        :param syntax: Syntax of the query, which is a one from the following choises:
+         1) QuerySyntax.YQL_V1, which is default;
+         2) QuerySyntax.PG.
+        :param parameters: dict with parameters and YDB types;
+        :param concurrent_result_sets: A flag to allow YDB mix parts of different result sets. Default is False;
+
+        :return: Iterator with result sets
+        """
+        self._state._check_session_ready_to_use()
+
+        stream_it = await self._execute_call(
+            query=query,
+            commit_tx=True,
+            syntax=syntax,
+            exec_mode=exec_mode,
+            parameters=parameters,
+            concurrent_result_sets=concurrent_result_sets,
+        )
+
+        return AsyncResponseContextIterator(
+            stream_it,
+            lambda resp: base.wrap_execute_query_response(rpc_state=None, response_pb=resp),
+        )