Merge pull request #88 from yassun7010/add_use_method

feat: add use_* methods.

Author: ya7010

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

@@ -18,7 +18,7 @@ class AsyncRecordCursor(  # type: ignore[override]
     ],
 ):
     async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
-        df = cast(GenericPandasDataFrame, await self._cursor.fetch_pandas_all(**kwargs))  # type: ignore[assignment]
+        df = await self._sf_cursor.fetch_pandas_all(**kwargs)
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
@@ -35,10 +35,8 @@ async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
     async def fetch_pandas_batches(
         self, **kwargs
     ) -> AsyncIterator[GenericPandasDataFrame]:
-        batches = cast(
-            AsyncIterator[GenericPandasDataFrame],
-            self._cursor.fetch_pandas_batches(**kwargs),  # type: ignore[assignment]
-        )
+        batches = self._sf_cursor.fetch_pandas_batches(**kwargs)
+
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
                 async for batch in batches:
@@ -52,7 +50,7 @@ async def fetch_pandas_batches(
             yield batch
 
     async def fetch_arrow_all(self) -> GenericPyArrowTable:
-        table = cast(GenericPyArrowTable, await self._cursor.fetch_arrow_all())  # type: ignore[assignment]
+        table = await self._sf_cursor.fetch_arrow_all()
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
@@ -67,10 +65,8 @@ async def fetch_arrow_all(self) -> GenericPyArrowTable:
         return table
 
     async def fetch_arrow_batches(self) -> AsyncIterator[GenericPyArrowTable]:
-        batches = cast(
-            AsyncIterator[GenericPyArrowTable],
-            self._cursor.fetch_arrow_batches(),  # type: ignore[assignment]
-        )
+        batches = self._sf_cursor.fetch_arrow_batches()
+
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
                 async for batch in batches:
@@ -82,3 +78,39 @@ async def fetch_arrow_batches(self) -> AsyncIterator[GenericPyArrowTable]:
 
         async for batch in batches:
             yield batch
+
+    def use_warehouse(self, warehouse: str, /) -> "AsyncRecordCursor":
+        """Use a warehouse in cursor."""
+
+        self._sf_cursor.use_warehouse(warehouse)
+
+        return self
+
+    def use_database(self, database: str, /) -> "AsyncRecordCursor":
+        """Use a database in cursor."""
+
+        self._sf_cursor.use_database(database)
+
+        return self
+
+    def use_schema(self, schema: str, /) -> "AsyncRecordCursor":
+        """Use a schema in cursor."""
+
+        self._sf_cursor.use_schema(schema)
+
+        return self
+
+    def use_role(self, role: str, /) -> "AsyncRecordCursor":
+        """Use a role in cursor."""
+
+        self._sf_cursor.use_role(role)
+
+        return self
+
+    @property
+    def _sf_cursor(
+        self,
+    ) -> turu.snowflake.async_cursor.AsyncCursor[
+        GenericRowType, GenericPandasDataFrame, GenericPyArrowTable
+    ]:
+        return cast(turu.snowflake.async_cursor.AsyncCursor, self._cursor)

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

@@ -7,14 +7,15 @@
     GenericPandasDataFrame,
     GenericPyArrowTable,
 )
+from typing_extensions import Self
 
 
 class RecordCursor(  # type: ignore[override]
     turu.core.record.RecordCursor,
     Generic[GenericRowType, GenericPandasDataFrame, GenericPyArrowTable],
 ):
     def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
-        df = cast(GenericPandasDataFrame, self._cursor.fetch_pandas_all(**kwargs))  # type: ignore[assignment]
+        df = self._sf_cursor.fetch_pandas_all(**kwargs)
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
@@ -29,10 +30,8 @@ def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
         return df
 
     def fetch_pandas_batches(self, **kwargs) -> Iterator[GenericPandasDataFrame]:
-        batches = cast(
-            Iterator[GenericPandasDataFrame],
-            self._cursor.fetch_pandas_batches(**kwargs),  # type: ignore[assignment]
-        )
+        batches = self._sf_cursor.fetch_pandas_batches(**kwargs)
+
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
                 for batch in batches:
@@ -45,7 +44,7 @@ def fetch_pandas_batches(self, **kwargs) -> Iterator[GenericPandasDataFrame]:
         return batches
 
     def fetch_arrow_all(self) -> GenericPyArrowTable:
-        table = cast(GenericPyArrowTable, self._cursor.fetch_arrow_all())  # type: ignore[assignment]
+        table = self._sf_cursor.fetch_arrow_all()
 
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
@@ -60,10 +59,8 @@ def fetch_arrow_all(self) -> GenericPyArrowTable:
         return table
 
     def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]:
-        batches = cast(
-            Iterator[GenericPyArrowTable],
-            self._cursor.fetch_arrow_batches(),  # type: ignore[assignment]
-        )
+        batches = self._sf_cursor.fetch_arrow_batches()
+
         if isinstance(self._recorder, turu.core.record.CsvRecorder):
             if limit := self._recorder._options.get("limit"):
                 for batch in batches:
@@ -74,3 +71,39 @@ def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]:
                         return
 
         return batches
+
+    def use_warehouse(self, warehouse: str, /) -> Self:
+        """Use a warehouse in cursor."""
+
+        self._sf_cursor.use_warehouse(warehouse)
+
+        return self
+
+    def use_database(self, database: str, /) -> Self:
+        """Use a database in cursor."""
+
+        self._sf_cursor.use_database(database)
+
+        return self
+
+    def use_schema(self, schema: str, /) -> Self:
+        """Use a schema in cursor."""
+
+        self._sf_cursor.use_schema(schema)
+
+        return self
+
+    def use_role(self, role: str, /) -> Self:
+        """Use a role in cursor."""
+
+        self._sf_cursor.use_role(role)
+
+        return self
+
+    @property
+    def _sf_cursor(
+        self,
+    ) -> turu.snowflake.cursor.Cursor[
+        GenericRowType, GenericPandasDataFrame, GenericPyArrowTable
+    ]:
+        return cast(turu.snowflake.cursor.Cursor, self._cursor)