Add the import_dbapi method to classes that inherit AthenaDialect to suppress SQLAlchemy WARNING messages. (fix #582)

pyathena/sqlalchemy/arrow.py

@@ -1,7 +1,12 @@
 # -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
 from pyathena.sqlalchemy.base import AthenaDialect
 from pyathena.util import strtobool
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
 
 class AthenaArrowDialect(AthenaDialect):
     driver = "arrow"
@@ -18,3 +23,7 @@ def create_connect_args(self, url):
         if cursor_kwargs:
             opts.update({"cursor_kwargs": cursor_kwargs})
         return [[], opts]
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()

pyathena/sqlalchemy/pandas.py

@@ -1,7 +1,12 @@
 # -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
 from pyathena.sqlalchemy.base import AthenaDialect
 from pyathena.util import strtobool
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
 
 class AthenaPandasDialect(AthenaDialect):
     driver = "pandas"
@@ -22,3 +27,7 @@ def create_connect_args(self, url):
         if cursor_kwargs:
             opts.update({"cursor_kwargs": cursor_kwargs})
         return [[], opts]
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()

pyathena/sqlalchemy/rest.py

@@ -1,7 +1,16 @@
 # -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
 from pyathena.sqlalchemy.base import AthenaDialect
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
 
 class AthenaRestDialect(AthenaDialect):
     driver = "rest"
     supports_statement_cache = True
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()