Merge pull request #71 from yassun7010/fix_limit

Fix limit
yassun7010 · Feb 8, 2024 · b21fcdd · b21fcdd
2 parents e265ee3 + 7dd342f
commit b21fcdd
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py b/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
diff --git a/turu-snowflake/src/turu/snowflake/record/record_cursor.py b/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