add: test case.

Author: ya7010

turu-snowflake/src/turu/snowflake/record/async_record_cursor.py

@@ -1,4 +1,4 @@
-from typing import AsyncIterator
+from typing import AsyncIterator, Generic, cast
 
 import turu.core.record
 import turu.snowflake.async_cursor
@@ -11,14 +11,14 @@
 
 class AsyncRecordCursor(  # type: ignore[override]
     turu.core.record.AsyncRecordCursor,
-    turu.snowflake.async_cursor.AsyncCursor[
+    Generic[
         GenericRowType,
         GenericPandasDataFrame,
         GenericPyArrowTable,
     ],
 ):
     async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
-        df = await super().fetch_pandas_all(**kwargs)
+        df = cast(GenericPandasDataFrame, await self._cursor.fetch_pandas_all(**kwargs))  # type: ignore[assignment]
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
@@ -35,7 +35,10 @@ async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
     async def fetch_pandas_batches(
         self, **kwargs
     ) -> AsyncIterator[GenericPandasDataFrame]:
-        batches = super().fetch_pandas_batches(**kwargs)
+        batches = cast(
+            AsyncIterator[GenericPandasDataFrame],
+            self._cursor.fetch_pandas_batches(**kwargs),  # type: ignore[assignment]
+        )
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
                 async for batch in batches:

turu-snowflake/src/turu/snowflake/record/record_cursor.py

@@ -1,4 +1,4 @@
-from typing import Iterator
+from typing import Generic, Iterator, cast
 
 import turu.core.record
 import turu.snowflake.cursor
@@ -11,12 +11,10 @@
 
 class RecordCursor(  # type: ignore[override]
     turu.core.record.RecordCursor,
-    turu.snowflake.cursor.Cursor[
-        GenericRowType, GenericPandasDataFrame, GenericPyArrowTable
-    ],
+    Generic[GenericRowType, GenericPandasDataFrame, GenericPyArrowTable],
 ):
     def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
-        df = super().fetch_pandas_all(**kwargs)
+        df = cast(GenericPandasDataFrame, self._cursor.fetch_pandas_all(**kwargs))  # type: ignore[assignment]
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
@@ -31,7 +29,10 @@ def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
         return df
 
     def fetch_pandas_batches(self, **kwargs) -> Iterator[GenericPandasDataFrame]:
-        batches = super().fetch_pandas_batches(**kwargs)
+        batches = cast(
+            Iterator[GenericPandasDataFrame],
+            self._cursor.fetch_pandas_batches(**kwargs),  # type: ignore[assignment]
+        )
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
                 for batch in batches:

turu-snowflake/tests/turu/test_snowflake_mock.py

@@ -5,6 +5,7 @@
 
 import pytest
 import turu.snowflake
+from turu.core.record import record_to_csv
 from turu.snowflake.features import (
     USE_PANDAS,
     USE_PANDERA,
@@ -416,3 +417,38 @@ class RowModel(pa.DataFrameModel):
         ).execute_map(RowModel, "select 1 as ID union all select 2 ID") as cursor:
             with pytest.raises(pandera.errors.SchemaInitError):
                 cursor.fetch_pandas_all()
+
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    def test_record_to_csv_and_fetch_pandas_all_with_limit_options(
+        self, mock_connection: turu.snowflake.MockConnection
+    ):
+        import pandas as pd  # type: ignore[import]
+
+        with tempfile.NamedTemporaryFile() as file:
+            with record_to_csv(
+                file.name,
+                mock_connection.inject_response(
+                    pd.DataFrame, pd.DataFrame({"ID": list(range(10))})
+                ).execute_map(pd.DataFrame, "select 1"),
+                limit=5,
+            ) as cursor:
+                assert cursor.fetch_pandas_all().equals(
+                    pd.DataFrame({"ID": list(range(5))})
+                )
+
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    def test_record_to_csv_and_fetch_pandas_batches_with_limit_options(
+        self, mock_connection: turu.snowflake.MockConnection
+    ):
+        import pandas as pd  # type: ignore[import]
+
+        with tempfile.NamedTemporaryFile() as file:
+            with record_to_csv(
+                file.name,
+                mock_connection.inject_response(
+                    pd.DataFrame, pd.DataFrame({"ID": list(range(10))})
+                ).execute_map(pd.DataFrame, "select 1"),
+                limit=5,
+            ) as cursor:
+                for batch in cursor.fetch_pandas_batches():
+                    assert batch.equals(pd.DataFrame({"ID": list(range(5))}))

turu-snowflake/tests/turu/test_snowflake_mock_async.py

@@ -5,6 +5,7 @@
 
 import pytest
 import turu.snowflake
+from turu.core.record import record_to_csv
 from turu.snowflake.features import USE_PANDAS, USE_PANDERA, USE_PYARROW, PyArrowTable
 from typing_extensions import Never
 
@@ -508,3 +509,40 @@ class RowModel(pa.DataFrameModel):
                 .execute_map(RowModel, "select 1 as ID union all select 2 as ID")
             ) as cursor:
                 assert (await cursor.fetch_pandas_all()).equals(expected)
+
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    @pytest.mark.asyncio
+    async def test_record_to_csv_and_fetch_pandas_all_with_limit_options(
+        self, mock_async_connection: turu.snowflake.MockAsyncConnection
+    ):
+        import pandas as pd  # type: ignore[import]
+
+        with tempfile.NamedTemporaryFile() as file:
+            async with record_to_csv(
+                file.name,
+                await mock_async_connection.inject_response(
+                    pd.DataFrame, pd.DataFrame({"ID": list(range(10))})
+                ).execute_map(pd.DataFrame, "select 1"),
+                limit=5,
+            ) as cursor:
+                assert (await cursor.fetch_pandas_all()).equals(
+                    pd.DataFrame({"ID": list(range(5))})
+                )
+
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    @pytest.mark.asyncio
+    async def test_record_to_csv_and_fetch_pandas_batches_with_limit_options(
+        self, mock_async_connection: turu.snowflake.MockAsyncConnection
+    ):
+        import pandas as pd  # type: ignore[import]
+
+        with tempfile.NamedTemporaryFile() as file:
+            async with record_to_csv(
+                file.name,
+                await mock_async_connection.inject_response(
+                    pd.DataFrame, pd.DataFrame({"ID": list(range(10))})
+                ).execute_map(pd.DataFrame, "select 1"),
+                limit=5,
+            ) as cursor:
+                async for batch in cursor.fetch_pandas_batches():
+                    assert batch.equals(pd.DataFrame({"ID": list(range(5))}))