Merge pull request #75 from yassun7010/use_header_option

fix: record_to_csv header option

Author: ya7010

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

@@ -8,9 +8,15 @@ async def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
-                df.head(limit).to_csv(self._recorder.file, index=False)
+                record_df = df.head(limit)
 
             else:
-                df.to_csv(self._recorder.file, index=False)
+                record_df = df
+
+            record_df.to_csv(
+                self._recorder.file,
+                index=False,
+                header=self._recorder._options.get("header", True),
+            )
 
         return df

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

@@ -8,9 +8,15 @@ def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame":
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
-                df.head(limit).to_csv(self._recorder.file, index=False)
+                record_df = df.head(limit)
 
             else:
-                df.to_csv(self._recorder.file, index=False)
+                record_df = df
+
+            record_df.to_csv(
+                self._recorder.file,
+                index=False,
+                header=self._recorder._options.get("header", True),
+            )
 
         return df

turu-snowflake/tests/turu/test_snowflake.py

@@ -266,6 +266,35 @@ def test_record_pandas_dataframe(self, connection: turu.snowflake.Connection):
                 ).lstrip()
             )
 
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    def test_record_pandas_dataframe_without_header_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"
+                ),
+                header=False,
+            ) 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(
+                    """
+                    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

turu-snowflake/tests/turu/test_snowflake_async.py

@@ -379,14 +379,72 @@ async def test_record_pandas_dataframe(
                     "select 1 as ID union all select 2 AS ID",
                 ),
             ) as cursor:
-                expected = pd.DataFrame(
-                    {"ID": [1, 2]},
-                    dtype="int8",
-                )
+                expected = pd.DataFrame({"ID": [1, 2]}, dtype="int8")
+
+                assert_frame_equal(await 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")
+    @pytest.mark.asyncio
+    async def test_record_pandas_dataframe_without_header_option(
+        self, async_connection: turu.snowflake.AsyncConnection
+    ):
+        import pandas as pd  # type: ignore[import]
+        from pandas.testing import assert_frame_equal  # type: ignore[import]
+
+        with tempfile.NamedTemporaryFile() as file:
+            async with record_to_csv(
+                file.name,
+                await async_connection.execute_map(
+                    pd.DataFrame,
+                    "select 1 as ID union all select 2 AS ID",
+                ),
+                header=False,
+            ) as cursor:
+                expected = pd.DataFrame({"ID": [1, 2]}, dtype="int8")
+
+                assert_frame_equal(await cursor.fetch_pandas_all(), expected)
+
+            assert (
+                Path(file.name).read_text()
+                == dedent(
+                    """
+                    1
+                    2
+                    """
+                ).lstrip()
+            )
+
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    @pytest.mark.asyncio
+    async def test_record_pandas_dataframe_with_limit_option(
+        self, async_connection: turu.snowflake.AsyncConnection
+    ):
+        import pandas as pd  # type: ignore[import]
+        from pandas.testing import assert_frame_equal  # type: ignore[import]
+
+        with tempfile.NamedTemporaryFile() as file:
+            async with record_to_csv(
+                file.name,
+                await async_connection.execute_map(
+                    pd.DataFrame,
+                    "select value::integer as ID from table(flatten(ARRAY_GENERATE_RANGE(1, 10)))",
+                ),
+                limit=2,
+            ) as cursor:
+                expected = pd.DataFrame({"ID": list(range(1, 10))}, dtype="object")
 
                 assert_frame_equal(await cursor.fetch_pandas_all(), expected)
-                for row in expected.values:
-                    print(row)
 
             assert (
                 Path(file.name).read_text()