Merge branch 'main' into dependabot/pip/mkdocs-material-9.5.7

Author: ya7010

turu-core/src/turu/core/mock/store.py

@@ -9,7 +9,7 @@
 
 class TuruMockStore:
     def __init__(self):
-        self._data: List[Tuple[Optional[Type], Sequence | None | Exception]] = []
+        self._data: List[Tuple[Optional[Type], Union[Sequence, None, Exception]]] = []
         self._counter = 0
 
     @overload
@@ -37,15 +37,18 @@ def inject_response(self, row_type, response):
     def provide_response(
         self,
         row_type: Optional[Type],
-    ) -> "Sequence[Any] | None":
+    ) -> "Optional[Sequence[Any]]":
         self._counter += 1
 
         if len(self._data) == 0:
             raise TuruMockStoreDataNotFoundError(self._counter)
 
         _row_type, _response = self._data.pop(0)
 
-        if _row_type is not row_type:
+        if _row_type is not row_type and not (
+            row_type.__module__ == _row_type.__module__
+            and row_type.__name__ == _row_type.__name__
+        ):
             raise TuruMockResponseTypeMismatchError(row_type, _row_type, self._counter)
 
         if isinstance(_response, Exception):

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

@@ -5,7 +5,12 @@
 class AsyncRecordCursor(turu.core.record.AsyncRecordCursor):
     async def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
         df: PandasDataFrame = self._raw_cursor.fetch_pandas_all(**kwargs)  # type: ignore
+
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
-            df.to_csv(self._recorder.file, index=False)
+            if limit := self._recorder._options.get("limit"):
+                df.head(limit).to_csv(self._recorder.file, index=False)
+
+            else:
+                df.to_csv(self._recorder.file, index=False)
 
         return df

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

@@ -5,7 +5,12 @@
 class RecordCursor(turu.core.record.RecordCursor):
     def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
         df: PandasDataFrame = self._raw_cursor.fetch_pandas_all(**kwargs)  # type: ignore
+
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
-            df.to_csv(self._recorder.file, index=False)
+            if limit := self._recorder._options.get("limit"):
+                df.head(limit).to_csv(self._recorder.file, index=False)
+
+            else:
+                df.to_csv(self._recorder.file, index=False)
 
         return df

turu-snowflake/tests/turu/test_snowflake.py

@@ -244,22 +244,47 @@ def test_record_pandas_dataframe(self, connection: turu.snowflake.Connection):
         import pandas as pd  # type: ignore[import]
         from pandas.testing import assert_frame_equal  # type: ignore[import]
 
+        with tempfile.NamedTemporaryFile() as file:
+            with record_to_csv(
+                file.name,
+                connection.execute_map(
+                    pd.DataFrame, "select 1 as ID union all select 2 AS ID"
+                ),
+            ) as cursor:
+                expected = pd.DataFrame({"ID": [1, 2]}, dtype="int8")
+
+                assert_frame_equal(cursor.fetch_pandas_all(), expected)
+
+            assert (
+                Path(file.name).read_text()
+                == dedent(
+                    """
+                    ID
+                    1
+                    2
+                    """
+                ).lstrip()
+            )
+
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    def test_record_pandas_dataframe_with_limit_option(
+        self, connection: turu.snowflake.Connection
+    ):
+        import pandas as pd  # type: ignore[import]
+        from pandas.testing import assert_frame_equal  # type: ignore[import]
+
         with tempfile.NamedTemporaryFile() as file:
             with record_to_csv(
                 file.name,
                 connection.execute_map(
                     pd.DataFrame,
-                    "select 1 as ID union all select 2 AS ID",
+                    "select value::integer as ID from table(flatten(ARRAY_GENERATE_RANGE(1, 10)))",
                 ),
+                limit=2,
             ) as cursor:
-                expected = pd.DataFrame(
-                    {"ID": [1, 2]},
-                    dtype="int8",
-                )
+                expected = pd.DataFrame({"ID": list(range(1, 10))}, dtype="object")
 
                 assert_frame_equal(cursor.fetch_pandas_all(), expected)
-                for row in expected.values:
-                    print(row)
 
             assert (
                 Path(file.name).read_text()