Skip to content

Commit d9c4a18

Browse files
committed
fix: limit option.
1 parent a574850 commit d9c4a18

File tree

2 files changed

+12
-2
lines changed

2 files changed

+12
-2
lines changed

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,12 @@
55
class AsyncRecordCursor(turu.core.record.AsyncRecordCursor):
66
async def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
77
df: PandasDataFrame = self._raw_cursor.fetch_pandas_all(**kwargs) # type: ignore
8+
89
if isinstance(self._recorder, turu.core.record.CsvRecorder):
9-
df.to_csv(self._recorder.file, index=False)
10+
if limit := self._recorder._options.get("limit"):
11+
df.head(limit).to_csv(self._recorder.file, index=False)
12+
13+
else:
14+
df.to_csv(self._recorder.file, index=False)
1015

1116
return df

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,12 @@
55
class RecordCursor(turu.core.record.RecordCursor):
66
def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
77
df: PandasDataFrame = self._raw_cursor.fetch_pandas_all(**kwargs) # type: ignore
8+
89
if isinstance(self._recorder, turu.core.record.CsvRecorder):
9-
df.to_csv(self._recorder.file, index=False)
10+
if limit := self._recorder._options.get("limit"):
11+
df.head(limit).to_csv(self._recorder.file, index=False)
12+
13+
else:
14+
df.to_csv(self._recorder.file, index=False)
1015

1116
return df

0 commit comments

Comments
 (0)