From 71c55dc427321b25100f380ac2bf337b5fb78e86 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 May 2025 16:56:23 +0900 Subject: [PATCH] Add the import_dbapi method to classes that inherit AthenaDialect to suppress SQLAlchemy WARNING messages. (fix #582) --- pyathena/sqlalchemy/arrow.py | 9 +++++++++ pyathena/sqlalchemy/pandas.py | 9 +++++++++ pyathena/sqlalchemy/rest.py | 9 +++++++++ 3 files changed, 27 insertions(+) diff --git a/pyathena/sqlalchemy/arrow.py b/pyathena/sqlalchemy/arrow.py index 67528954..c7d8dea8 100644 --- a/pyathena/sqlalchemy/arrow.py +++ b/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() diff --git a/pyathena/sqlalchemy/pandas.py b/pyathena/sqlalchemy/pandas.py index a7bc88fa..8d32a03b 100644 --- a/pyathena/sqlalchemy/pandas.py +++ b/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() diff --git a/pyathena/sqlalchemy/rest.py b/pyathena/sqlalchemy/rest.py index 3aee9f81..dc8c6c9e 100644 --- a/pyathena/sqlalchemy/rest.py +++ b/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()