bug: Incorrect session closing handling

[LuckySting]

Bug Report

3.21.7

Environment

Everywhere

Current behavior:

When YDB QuerySession is attached, a poling task is created. This task reads the stream waiting for message and closes the session if get non-ok message or exception:

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:
            if not self._state._already_in(QuerySessionStateEnum.CLOSED):
                self._state.reset()
                self._state._change_state(QuerySessionStateEnum.CLOSED)

The problem is that in case of graceful shutdown, for example, due to node restart, the server sends OK message and then close the stream. The code above doesn't expect that session may be closed with OK message, also it doesn't handle stream closing.

Expected behavior:

_check_session_status_loop must close the session in case of stream closing in any case (with any response message).

Steps to reproduce:

To reproduce the issue you need to run SLO test environment and execute the following code

async def get_session_closing_status_message(driver: ydb.aio.Driver) -> None:
    session = ydb.aio.QuerySession(driver)
    await session._create_call()
    attach_stream = await session._attach_call()
    async for status in attach_stream:
        print(status, status.status) # prints status: SUCCESS 400000

Related code:

Need to modify this task