Merge pull request #82 from yassun7010/change_record_limit

feat: change limit behavor.
yassun7010 · Feb 15, 2024 · f96ff6c · f96ff6c
2 parents 4127080 + 6544fc9
commit f96ff6c
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 12 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
@@ -8,12 +8,9 @@ async def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
 
  if isinstance(self._recorder, turu.core.record.CsvRecorder):
  if limit := self._recorder._options.get("limit"):
- record_df = df.head(limit)
+ df = df.head(limit)
 
- else:
- record_df = df
-
- record_df.to_csv(
+ df.to_csv(
  self._recorder.file,
  index=False,
  header=self._recorder._options.get("header", True),

diff --git a/turu-snowflake/src/turu/snowflake/record/record_cursor.py b/turu-snowflake/src/turu/snowflake/record/record_cursor.py
@@ -8,12 +8,9 @@ def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
 
  if isinstance(self._recorder, turu.core.record.CsvRecorder):
  if limit := self._recorder._options.get("limit"):
- record_df = df.head(limit)
+ df = df.head(limit)
 
- else:
- record_df = df
-
- record_df.to_csv(
+ df.to_csv(
  self._recorder.file,
  index=False,
  header=self._recorder._options.get("header", True),

diff --git a/turu-snowflake/tests/turu/test_snowflake.py b/turu-snowflake/tests/turu/test_snowflake.py
@@ -311,7 +311,7 @@ def test_record_pandas_dataframe_with_limit_option(
  ),
  limit=2,
  ) as cursor:
- expected = pd.DataFrame({"ID": list(range(1, 10))}, dtype="object")
+ expected = pd.DataFrame({"ID": [1, 2]}, dtype="object")
 
  assert_frame_equal(cursor.fetch_pandas_all(), expected)
 

diff --git a/turu-snowflake/tests/turu/test_snowflake_async.py b/turu-snowflake/tests/turu/test_snowflake_async.py
@@ -442,7 +442,7 @@ async def test_record_pandas_dataframe_with_limit_option(
  ),
  limit=2,
  ) as cursor:
- expected = pd.DataFrame({"ID": list(range(1, 10))}, dtype="object")
+ expected = pd.DataFrame({"ID": [1, 2]}, dtype="object")
 
  assert_frame_equal(await cursor.fetch_pandas_all(), expected)