Revert missing cursor default arguments. (fix #506)

Author: laughingman7743

pyathena/arrow/async_cursor.py

@@ -22,14 +22,34 @@
 class AsyncArrowCursor(AsyncCursor):
     def __init__(
         self,
+        s3_staging_dir: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        catalog_name: Optional[str] = None,
+        work_group: Optional[str] = None,
+        poll_interval: float = 1,
+        encryption_option: Optional[str] = None,
+        kms_key: Optional[str] = None,
+        kill_on_interrupt: bool = True,
         max_workers: int = (cpu_count() or 1) * 5,
         arraysize: int = CursorIterator.DEFAULT_FETCH_SIZE,
         unload: bool = False,
+        result_reuse_enable: bool = False,
+        result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
         **kwargs,
     ) -> None:
         super().__init__(
+            s3_staging_dir=s3_staging_dir,
+            schema_name=schema_name,
+            catalog_name=catalog_name,
+            work_group=work_group,
+            poll_interval=poll_interval,
+            encryption_option=encryption_option,
+            kms_key=kms_key,
+            kill_on_interrupt=kill_on_interrupt,
             max_workers=max_workers,
             arraysize=arraysize,
+            result_reuse_enable=result_reuse_enable,
+            result_reuse_minutes=result_reuse_minutes,
             **kwargs,
         )
         self._unload = unload

pyathena/arrow/cursor.py

@@ -23,10 +23,32 @@
 class ArrowCursor(BaseCursor, CursorIterator, WithResultSet):
     def __init__(
         self,
+        s3_staging_dir: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        catalog_name: Optional[str] = None,
+        work_group: Optional[str] = None,
+        poll_interval: float = 1,
+        encryption_option: Optional[str] = None,
+        kms_key: Optional[str] = None,
+        kill_on_interrupt: bool = True,
         unload: bool = False,
+        result_reuse_enable: bool = False,
+        result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
         **kwargs,
     ) -> None:
-        super().__init__(**kwargs)
+        super().__init__(
+            s3_staging_dir=s3_staging_dir,
+            schema_name=schema_name,
+            catalog_name=catalog_name,
+            work_group=work_group,
+            poll_interval=poll_interval,
+            encryption_option=encryption_option,
+            kms_key=kms_key,
+            kill_on_interrupt=kill_on_interrupt,
+            result_reuse_enable=result_reuse_enable,
+            result_reuse_minutes=result_reuse_minutes,
+            **kwargs,
+        )
         self._unload = unload
         self._query_id: Optional[str] = None
         self._result_set: Optional[AthenaArrowResultSet] = None

pyathena/async_cursor.py

@@ -19,11 +19,33 @@
 class AsyncCursor(BaseCursor):
     def __init__(
         self,
+        s3_staging_dir: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        catalog_name: Optional[str] = None,
+        work_group: Optional[str] = None,
+        poll_interval: float = 1,
+        encryption_option: Optional[str] = None,
+        kms_key: Optional[str] = None,
+        kill_on_interrupt: bool = True,
         max_workers: int = (cpu_count() or 1) * 5,
         arraysize: int = CursorIterator.DEFAULT_FETCH_SIZE,
+        result_reuse_enable: bool = False,
+        result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
         **kwargs,
     ) -> None:
-        super().__init__(**kwargs)
+        super().__init__(
+            s3_staging_dir=s3_staging_dir,
+            schema_name=schema_name,
+            catalog_name=catalog_name,
+            work_group=work_group,
+            poll_interval=poll_interval,
+            encryption_option=encryption_option,
+            kms_key=kms_key,
+            kill_on_interrupt=kill_on_interrupt,
+            result_reuse_enable=result_reuse_enable,
+            result_reuse_minutes=result_reuse_minutes,
+            **kwargs,
+        )
         self._max_workers = max_workers
         self._executor = ThreadPoolExecutor(max_workers=max_workers)
         self._arraysize = arraysize

pyathena/cursor.py

@@ -13,8 +13,33 @@
 
 
 class Cursor(BaseCursor, CursorIterator, WithResultSet):
-    def __init__(self, **kwargs) -> None:
-        super().__init__(**kwargs)
+    def __init__(
+        self,
+        s3_staging_dir: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        catalog_name: Optional[str] = None,
+        work_group: Optional[str] = None,
+        poll_interval: float = 1,
+        encryption_option: Optional[str] = None,
+        kms_key: Optional[str] = None,
+        kill_on_interrupt: bool = True,
+        result_reuse_enable: bool = False,
+        result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            s3_staging_dir=s3_staging_dir,
+            schema_name=schema_name,
+            catalog_name=catalog_name,
+            work_group=work_group,
+            poll_interval=poll_interval,
+            encryption_option=encryption_option,
+            kms_key=kms_key,
+            kill_on_interrupt=kill_on_interrupt,
+            result_reuse_enable=result_reuse_enable,
+            result_reuse_minutes=result_reuse_minutes,
+            **kwargs,
+        )
         self._query_id: Optional[str] = None
         self._result_set: Optional[AthenaResultSet] = None
         self._result_set_class = AthenaResultSet

pyathena/pandas/async_cursor.py

@@ -22,14 +22,36 @@
 class AsyncPandasCursor(AsyncCursor):
     def __init__(
         self,
+        s3_staging_dir: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        catalog_name: Optional[str] = None,
+        work_group: Optional[str] = None,
+        poll_interval: float = 1,
+        encryption_option: Optional[str] = None,
+        kms_key: Optional[str] = None,
+        kill_on_interrupt: bool = True,
         max_workers: int = (cpu_count() or 1) * 5,
         arraysize: int = CursorIterator.DEFAULT_FETCH_SIZE,
         unload: bool = False,
         engine: str = "auto",
         chunksize: Optional[int] = None,
+        result_reuse_enable: bool = False,
+        result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
         **kwargs,
     ) -> None:
-        super().__init__(max_workers=max_workers, arraysize=arraysize, **kwargs)
+        super().__init__(
+            s3_staging_dir=s3_staging_dir,
+            schema_name=schema_name,
+            catalog_name=catalog_name,
+            work_group=work_group,
+            poll_interval=poll_interval,
+            encryption_option=encryption_option,
+            kms_key=kms_key,
+            kill_on_interrupt=kill_on_interrupt,
+            max_workers=max_workers,
+            arraysize=arraysize,
+            **kwargs,
+        )
         self._unload = unload
         self._engine = engine
         self._chunksize = chunksize

pyathena/pandas/cursor.py

@@ -35,15 +35,37 @@
 class PandasCursor(BaseCursor, CursorIterator, WithResultSet):
     def __init__(
         self,
+        s3_staging_dir: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        catalog_name: Optional[str] = None,
+        work_group: Optional[str] = None,
+        poll_interval: float = 1,
+        encryption_option: Optional[str] = None,
+        kms_key: Optional[str] = None,
+        kill_on_interrupt: bool = True,
         unload: bool = False,
         engine: str = "auto",
         chunksize: Optional[int] = None,
         block_size: Optional[int] = None,
         cache_type: Optional[str] = None,
         max_workers: int = (cpu_count() or 1) * 5,
+        result_reuse_enable: bool = False,
+        result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
         **kwargs,
     ) -> None:
-        super().__init__(**kwargs)
+        super().__init__(
+            s3_staging_dir=s3_staging_dir,
+            schema_name=schema_name,
+            catalog_name=catalog_name,
+            work_group=work_group,
+            poll_interval=poll_interval,
+            encryption_option=encryption_option,
+            kms_key=kms_key,
+            kill_on_interrupt=kill_on_interrupt,
+            result_reuse_enable=result_reuse_enable,
+            result_reuse_minutes=result_reuse_minutes,
+            **kwargs,
+        )
         self._unload = unload
         self._engine = engine
         self._chunksize = chunksize

pyathena/spark/async_cursor.py

@@ -14,8 +14,24 @@
 
 
 class AsyncSparkCursor(SparkBaseCursor):
-    def __init__(self, max_workers: int = (cpu_count() or 1) * 5, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(
+        self,
+        session_id: Optional[str] = None,
+        description: Optional[str] = None,
+        engine_configuration: Optional[Dict[str, Any]] = None,
+        notebook_version: Optional[str] = None,
+        session_idle_timeout_minutes: Optional[int] = None,
+        max_workers: int = (cpu_count() or 1) * 5,
+        **kwargs,
+    ):
+        super().__init__(
+            session_id=session_id,
+            description=description,
+            engine_configuration=engine_configuration,
+            notebook_version=notebook_version,
+            session_idle_timeout_minutes=session_idle_timeout_minutes,
+            **kwargs,
+        )
         self._max_workers = max_workers
         self._executor = ThreadPoolExecutor(max_workers=max_workers)
 

pyathena/spark/cursor.py

@@ -14,9 +14,21 @@
 class SparkCursor(SparkBaseCursor, WithCalculationExecution):
     def __init__(
         self,
+        session_id: Optional[str] = None,
+        description: Optional[str] = None,
+        engine_configuration: Optional[Dict[str, Any]] = None,
+        notebook_version: Optional[str] = None,
+        session_idle_timeout_minutes: Optional[int] = None,
         **kwargs,
     ) -> None:
-        super().__init__(**kwargs)
+        super().__init__(
+            session_id=session_id,
+            description=description,
+            engine_configuration=engine_configuration,
+            notebook_version=notebook_version,
+            session_idle_timeout_minutes=session_idle_timeout_minutes,
+            **kwargs,
+        )
 
     @property
     def calculation_execution(self) -> Optional[AthenaCalculationExecution]: