feat: return async iterator (#44)

Author: barakalon

examples/streaming_results.py

@@ -0,0 +1,26 @@
+import logging
+import asyncio
+
+from mysql_mimic import MysqlServer, Session
+
+
+class MySession(Session):
+    async def generate_rows(self, n):
+        for i in range(n):
+            if i % 100 == 0:
+                logging.info("Pretending to fetch another batch of results...")
+                await asyncio.sleep(1)
+            yield i,
+
+    async def query(self, expression, sql, attrs):
+        return self.generate_rows(1000), ["a"]
+
+
+async def main():
+    logging.basicConfig(level=logging.DEBUG)
+    server = MysqlServer(session_factory=MySession)
+    await server.serve_forever()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

mysql_mimic/__init__.py

@@ -9,3 +9,4 @@
 from mysql_mimic.results import AllowedResult, ResultColumn, ResultSet
 from mysql_mimic.session import Session
 from mysql_mimic.server import MysqlServer
+from mysql_mimic.types import ColumnType

mysql_mimic/connection.py

@@ -1,6 +1,6 @@
 import logging
 from ssl import SSLContext
-from typing import Optional, Dict, Any, Iterator
+from typing import Optional, Dict, Any, Iterator, AsyncIterator
 
 from mysql_mimic.auth import (
     AuthInfo,
@@ -26,7 +26,7 @@
 from mysql_mimic.session import BaseSession
 from mysql_mimic.stream import MysqlStream, ConnectionClosed
 from mysql_mimic.types import Capabilities
-from mysql_mimic.utils import seq
+from mysql_mimic.utils import seq, aiterate
 
 logger = logging.getLogger(__name__)
 
@@ -352,14 +352,16 @@ async def handle_field_list(self, data: bytes) -> None:
         sql = com_field_list_to_show_statement(com_field_list)
         result = await self.query(sql=sql, query_attrs={})
         columns = b"".join(
-            make_column_definition_41(
-                server_charset=self.server_charset,
-                table=com_field_list.table,
-                name=row[0],
-                is_com_field_list=True,
-                default=row[4],
-            )
-            for row in result.rows
+            [
+                make_column_definition_41(
+                    server_charset=self.server_charset,
+                    table=com_field_list.table,
+                    name=row[0],
+                    is_com_field_list=True,
+                    default=row[4],
+                )
+                async for row in aiterate(result.rows)
+            ]
         )
         await self.stream.write(columns)
         await self.stream.write(self.ok_or_eof())
@@ -383,7 +385,7 @@ async def handle_query(self, data: bytes) -> None:
             await self.stream.write(self.ok())
             return
 
-        for packet in self.text_resultset(result_set):
+        async for packet in self.text_resultset(result_set):
             await self.stream.write(packet)
 
     async def handle_stmt_prepare(self, data: bytes) -> None:
@@ -457,10 +459,11 @@ async def handle_stmt_execute(self, data: bytes) -> None:
                 )
             )
 
-        rows = (
-            packets.make_binary_resultrow(r, result_set.columns)
-            for r in result_set.rows
-        )
+        async def gen_rows() -> AsyncIterator[bytes]:
+            async for r in aiterate(result_set.rows):
+                yield packets.make_binary_resultrow(r, result_set.columns)
+
+        rows = gen_rows()
 
         if com_stmt_execute.use_cursor:
             com_stmt_execute.stmt.cursor = rows
@@ -470,7 +473,7 @@ async def handle_stmt_execute(self, data: bytes) -> None:
         else:
             if not self.deprecate_eof():
                 await self.stream.write(self.eof())
-            for row in rows:
+            async for row in rows:
                 await self.stream.write(row)
             await self.stream.write(self.ok_or_eof())
 
@@ -485,7 +488,10 @@ async def handle_stmt_fetch(self, data: bytes) -> None:
         stmt = self.get_stmt(com_stmt_fetch.stmt_id)
         assert stmt.cursor is not None
         count = 0
-        for _, packet in zip(range(com_stmt_fetch.num_rows), stmt.cursor):
+
+        async for packet in stmt.cursor:
+            if count >= com_stmt_fetch.num_rows:
+                break
             await self.stream.write(packet)
             count += 1
 
@@ -530,7 +536,7 @@ def get_stmt(self, stmt_id: int) -> PreparedStatement:
     async def query(self, sql: str, query_attrs: Dict[str, str]) -> ResultSet:
         logger.debug("Received query: %s", sql)
 
-        result_set = ensure_result_set(
+        result_set = await ensure_result_set(
             await self.session.handle_query(sql, query_attrs)
         )
         return result_set
@@ -574,7 +580,7 @@ def error(self, **kwargs: Any) -> bytes:
     def deprecate_eof(self) -> bool:
         return Capabilities.CLIENT_DEPRECATE_EOF in self.capabilities
 
-    def text_resultset(self, result_set: ResultSet) -> Iterator[bytes]:
+    async def text_resultset(self, result_set: ResultSet) -> AsyncIterator[bytes]:
         yield packets.make_column_count(
             capabilities=self.capabilities, column_count=len(result_set.columns)
         )
@@ -592,7 +598,7 @@ def text_resultset(self, result_set: ResultSet) -> Iterator[bytes]:
 
         affected_rows = 0
 
-        for row in result_set.rows:
+        async for row in aiterate(result_set.rows):
             affected_rows += 1
             yield packets.make_text_resultset_row(row, result_set.columns)
 

mysql_mimic/prepared.py

@@ -1,7 +1,6 @@
 import re
 from dataclasses import dataclass
-from typing import Optional, Dict, Iterable
-
+from typing import Optional, Dict, AsyncIterable
 
 # Borrowed from mysql-connector-python
 REGEX_PARAM = re.compile(r"""\?(?=(?:[^"'`]*["'`][^"'`]*["'`])*[^"'`]*$)""")
@@ -13,4 +12,4 @@ class PreparedStatement:
     sql: str
     num_params: int
     param_buffers: Optional[Dict[int, bytearray]] = None
-    cursor: Optional[Iterable] = None
+    cursor: Optional[AsyncIterable[bytes]] = None

mysql_mimic/results.py

@@ -4,12 +4,23 @@
 import struct
 from dataclasses import dataclass
 from datetime import datetime, date, timedelta
-from typing import Iterable, Sequence, Optional, Callable, Any, Union, Tuple, Dict
+from typing import (
+    Iterable,
+    Sequence,
+    Optional,
+    Callable,
+    Any,
+    Union,
+    Tuple,
+    Dict,
+    AsyncIterable,
+    cast,
+)
 
 from mysql_mimic.errors import MysqlError
 from mysql_mimic.types import ColumnType, str_len, uint_1, uint_2, uint_4
 from mysql_mimic.charset import CharacterSet
-
+from mysql_mimic.utils import aiterate
 
 Encoder = Callable[[Any, "ResultColumn"], bytes]
 
@@ -54,7 +65,7 @@ def __repr__(self) -> str:
 
 @dataclass
 class ResultSet:
-    rows: Iterable[Sequence]
+    rows: Iterable[Sequence] | AsyncIterable[Sequence]
     columns: Sequence[ResultColumn]
 
     def __bool__(self) -> bool:
@@ -63,11 +74,14 @@ def __bool__(self) -> bool:
 
 AllowedColumn = Union[ResultColumn, str]
 AllowedResult = Union[
-    ResultSet, Tuple[Sequence[Sequence[Any]], Sequence[AllowedColumn]], None
+    ResultSet,
+    Tuple[Sequence[Sequence[Any]], Sequence[AllowedColumn]],
+    Tuple[AsyncIterable[Sequence[Any]], Sequence[AllowedColumn]],
+    None,
 ]
 
 
-def ensure_result_set(result: AllowedResult) -> ResultSet:
+async def ensure_result_set(result: AllowedResult) -> ResultSet:
     if result is None:
         return ResultSet([], [])
     if isinstance(result, ResultSet):
@@ -80,39 +94,74 @@ def ensure_result_set(result: AllowedResult) -> ResultSet:
         rows = result[0]
         columns = result[1]
 
-        return ResultSet(
-            rows=rows,
-            columns=[_ensure_result_col(col, i, rows) for i, col in enumerate(columns)],
-        )
+        return await _ensure_result_cols(rows, columns)
 
     raise MysqlError(f"Unexpected result set type: {type(result)}")
 
 
-def _ensure_result_col(
-    column: AllowedColumn, idx: int, rows: Sequence[Sequence[Any]]
-) -> ResultColumn:
-    if isinstance(column, ResultColumn):
-        return column
+async def _ensure_result_cols(
+    rows: Sequence[Sequence[Any]] | AsyncIterable[Sequence[Any]],
+    columns: Sequence[AllowedColumn],
+) -> ResultSet:
+    # Which columns need to be inferred?
+    remaining = {
+        col: i for i, col in enumerate(columns) if not isinstance(col, ResultColumn)
+    }
+
+    if not remaining:
+        return ResultSet(
+            rows=rows,
+            columns=cast(Sequence[ResultColumn], columns),
+        )
 
-    if isinstance(column, str):
-        value = _find_first_non_null_value(idx, rows)
-        type_ = infer_type(value)
-        return ResultColumn(
-            name=column,
-            type=type_,
+    # Copy the columns
+    columns = list(columns)
+
+    arows = aiterate(rows)
+
+    # Keep track of rows we've consumed from the iterator so we can add them back
+    peeks = []
+
+    # Find the first non-null value for each column
+    while remaining:
+        try:
+            peek = await arows.__anext__()
+        except StopAsyncIteration:
+            break
+
+        peeks.append(peek)
+
+        inferred = []
+        for name, i in remaining.items():
+            value = peek[i]
+            if value is not None:
+                type_ = infer_type(value)
+                columns[i] = ResultColumn(
+                    name=str(name),
+                    type=type_,
+                )
+                inferred.append(name)
+
+        for name in inferred:
+            remaining.pop(name)
+
+    # If we failed to find a non-null value, set the type to NULL
+    for name, i in remaining.items():
+        columns[i] = ResultColumn(
+            name=str(name),
+            type=ColumnType.NULL,
         )
 
-    raise MysqlError(f"Unexpected result column value: {column}")
+    # Add the consumed rows back in to the iterator
+    async def gen_rows() -> AsyncIterable[Sequence[Any]]:
+        for row in peeks:
+            yield row
 
+        async for row in arows:
+            yield row
 
-def _find_first_non_null_value(
-    idx: int, rows: Sequence[Sequence[Any]]
-) -> Optional[Any]:
-    for row in rows:
-        value = row[idx]
-        if value is not None:
-            return value
-    return None
+    assert all(isinstance(col, ResultColumn) for col in columns)
+    return ResultSet(rows=gen_rows(), columns=cast(Sequence[ResultColumn], columns))
 
 
 def _binary_encode_tiny(col: ResultColumn, val: Any) -> bytes:

mysql_mimic/utils.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
+
+import inspect
 import sys
 from collections.abc import Iterator
 import random
-from typing import List
+from typing import List, TypeVar, AsyncIterable, Iterable, AsyncIterator, cast
 import string
 
 from sqlglot import expressions as exp
 from sqlglot.optimizer.scope import traverse_scope
 
 
+T = TypeVar("T")
+
 # MySQL Connector/J uses ASCII to decode nonce
 SAFE_NONCE_CHARS = (string.ascii_letters + string.digits).encode()
 
@@ -92,3 +96,13 @@ def dict_depth(d: dict) -> int:
     except StopIteration:
         # d.values() returns an empty sequence
         return 1
+
+
+async def aiterate(iterable: AsyncIterable[T] | Iterable[T]) -> AsyncIterator[T]:
+    """Iterate either an async iterable or a regular iterable"""
+    if inspect.isasyncgen(iterable):
+        async for item in iterable:
+            yield item
+    else:
+        for item in cast(Iterable, iterable):
+            yield item

tests/test_query.py

@@ -1,3 +1,4 @@
+import asyncio
 import io
 from contextlib import closing
 from datetime import date, datetime, timedelta
@@ -835,3 +836,26 @@ async def test_unsupported_commands(
     with pytest.raises(Exception) as ctx:
         await query_fixture(sql)
     assert msg in str(ctx.value)
+
+
+@pytest.mark.asyncio
+async def test_async_iterator(
+    session: MockSession,
+    server: MysqlServer,
+    query_fixture: QueryFixture,
+) -> None:
+    async def generate_rows() -> Any:
+        yield 1, None, None
+        await asyncio.sleep(0)
+        yield None, "2", None
+        await asyncio.sleep(0)
+        yield None, None, None
+
+    session.return_value = (generate_rows(), ["a", "b", "c"])
+
+    result = await query_fixture("SELECT * FROM x")
+    assert [
+        {"a": 1, "b": None, "c": None},
+        {"a": None, "b": "2", "c": None},
+        {"a": None, "b": None, "c": None},
+    ] == result