Fix empty result sets from stream

Author: vgvoleg

tests/aio/query/test_query_session.py

@@ -103,10 +103,13 @@ async def test_basic_execute(self, session: QuerySession):
     async def test_two_results(self, session: QuerySession):
         await session.create()
         res = []
+        counter = 0
 
         async with await session.execute("select 1; select 2") as results:
             async for result_set in results:
+                counter += 1
                 if len(result_set.rows) > 0:
                     res.append(list(result_set.rows[0].values()))
 
         assert res == [[1], [2]]
+        assert counter == 2

tests/aio/query/test_query_transaction.py

@@ -92,3 +92,18 @@ async def test_execute_as_context_manager(self, tx: QueryTxContext):
             res = [result_set async for result_set in results]
 
         assert len(res) == 1
+
+    @pytest.mark.asyncio
+    async def test_execute_two_results(self, tx: QueryTxContext):
+        await tx.begin()
+        counter = 0
+        res = []
+
+        async with await tx.execute("select 1; select 2") as results:
+            async for result_set in results:
+                counter += 1
+                if len(result_set.rows) > 0:
+                    res.append(list(result_set.rows[0].values()))
+
+        assert res == [[1], [2]]
+        assert counter == 2

tests/query/test_query_session.py

@@ -98,13 +98,16 @@ def test_basic_execute(self, session: QuerySession):
     def test_two_results(self, session: QuerySession):
         session.create()
         res = []
+        counter = 0
 
         with session.execute("select 1; select 2") as results:
             for result_set in results:
+                counter += 1
                 if len(result_set.rows) > 0:
                     res.append(list(result_set.rows[0].values()))
 
         assert res == [[1], [2]]
+        assert counter == 2
 
     def test_thread_leaks(self, session: QuerySession):
         session.create()

tests/query/test_query_transaction.py

@@ -79,3 +79,16 @@ def test_execute_as_context_manager(self, tx: QueryTxContext):
             res = [result_set for result_set in results]
 
         assert len(res) == 1
+
+    def test_execute_two_results(self, tx: QueryTxContext):
+        tx.begin()
+        counter = 0
+        res = []
+
+        with tx.execute("select 1; select 2") as results:
+            for result_set in results:
+                counter += 1
+                res.append(list(result_set.rows[0].values()))
+
+        assert res == [[1], [2]]
+        assert counter == 2

ydb/_utilities.py

@@ -149,9 +149,10 @@ def __next__(self):
 
 
 class SyncResponseIterator(object):
-    def __init__(self, it, wrapper):
+    def __init__(self, it, wrapper, _filter=None):
         self.it = it
         self.wrapper = wrapper
+        self.filter = _filter
 
     def cancel(self):
         self.it.cancel()
@@ -161,7 +162,11 @@ def __iter__(self):
         return self
 
     def _next(self):
-        return self.wrapper(next(self.it))
+        res = next(self.it)
+        wrapped = self.wrapper(res)
+        if self.filter is None or self.filter(res):
+            return wrapped
+        return self._next()
 
     def next(self):
         return self._next()

ydb/aio/_utilities.py

@@ -2,9 +2,10 @@
 
 
 class AsyncResponseIterator(object):
-    def __init__(self, it, wrapper):
+    def __init__(self, it, wrapper, _filter=None):
         self.it = it.__aiter__()
         self.wrapper = wrapper
+        self.filter = _filter
 
     def cancel(self):
         self.it.cancel()
@@ -17,7 +18,11 @@ def __aiter__(self):
         return self
 
     async def _next(self):
-        return self.wrapper(await self.it.__anext__())
+        res = await self.it.__anext__()
+        wrapped = self.wrapper(res)
+        if self.filter is None or self.filter(res):
+            return wrapped
+        return await self._next()
 
     async def next(self):
         return await self._next()

ydb/aio/query/session.py

@@ -149,4 +149,5 @@ async def execute(
                 session_state=self._state,
                 settings=self._settings,
             ),
+            lambda resp: resp.HasField("result_set"),
         )

ydb/aio/query/transaction.py

@@ -144,5 +144,6 @@ async def execute(
                 commit_tx=commit_tx,
                 settings=self.session._settings,
             ),
+            lambda resp: resp.HasField("result_set"),
         )
         return self._prev_stream

ydb/query/session.py

@@ -353,4 +353,5 @@ def execute(
                 session_state=self._state,
                 settings=self._settings,
             ),
+            lambda resp: resp.HasField("result_set"),
         )

ydb/query/transaction.py

@@ -425,5 +425,6 @@ def execute(
                 commit_tx=commit_tx,
                 settings=self.session._settings,
             ),
+            lambda resp: resp.HasField("result_set"),
         )
         return self._prev_stream