Add tests for async SQLAlchemy dialects

- Add representative async tests in tests/pyathena/aio/sqlalchemy/:
  basic query, reflection, schema inspection, dialect properties
- Add ASYNC_SQLALCHEMY_CONNECTION_STRING to tests/__init__.py
- Register async dialects in tests/sqlalchemy/__init__.py
- Support --dburi async in tests/sqlalchemy/conftest.py for running
  the SQLAlchemy standard test suite with async dialects
- Add make test-sqla-async target

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: laughingman7743

Makefile

@@ -21,6 +21,10 @@ test: chk
 test-sqla:
 	uv run pytest -n 8 --cov pyathena --cov-report html --cov-report term tests/sqlalchemy/
 
+.PHONY: test-sqla-async
+test-sqla-async:
+	uv run pytest -n 8 --cov pyathena --cov-report html --cov-report term tests/sqlalchemy/ --dburi async
+
 .PHONY: tox
 tox:
 	uvx tox@$(TOX_VERSION) -c pyproject.toml run

tests/__init__.py

@@ -7,6 +7,10 @@
     "awsathena+rest://athena.{region_name}.amazonaws.com:443/"
     "{schema_name}?s3_staging_dir={s3_staging_dir}&location={location}"
 )
+ASYNC_SQLALCHEMY_CONNECTION_STRING = (
+    "awsathena+aiorest://athena.{region_name}.amazonaws.com:443/"
+    "{schema_name}?s3_staging_dir={s3_staging_dir}&location={location}"
+)
 
 
 class Env:

tests/pyathena/aio/sqlalchemy/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

tests/pyathena/aio/sqlalchemy/test_base.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+import pytest
+import sqlalchemy
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.sql.schema import MetaData, Table
+
+from tests import ASYNC_SQLALCHEMY_CONNECTION_STRING, ENV
+
+
+def _async_conn_str(**kwargs):
+    conn_str = ASYNC_SQLALCHEMY_CONNECTION_STRING
+    return conn_str.format(
+        region_name=ENV.region_name,
+        schema_name=ENV.schema,
+        s3_staging_dir=ENV.s3_staging_dir,
+        location=ENV.s3_staging_dir,
+        **kwargs,
+    )
+
+
+@pytest.fixture
+async def async_engine():
+    engine = create_async_engine(_async_conn_str())
+    try:
+        async with engine.connect() as conn:
+            yield engine, conn
+    finally:
+        await engine.dispose()
+
+
+class TestAsyncSQLAlchemyAthena:
+    async def test_basic_query(self, async_engine):
+        engine, conn = async_engine
+        rows = (await conn.execute(text("SELECT * FROM one_row"))).fetchall()
+        assert len(rows) == 1
+        assert rows[0].number_of_rows == 1
+        assert len(rows[0]) == 1
+
+    async def test_unicode(self, async_engine):
+        engine, conn = async_engine
+        unicode_str = "密林"
+        returned_str = (
+            await conn.execute(
+                sqlalchemy.select(
+                    sqlalchemy.sql.expression.bindparam(
+                        "あまぞん", unicode_str, type_=sqlalchemy.types.String()
+                    )
+                )
+            )
+        ).scalar()
+        assert returned_str == unicode_str
+
+    async def test_reflect_table(self, async_engine):
+        engine, conn = async_engine
+        one_row = await conn.run_sync(
+            lambda sync_conn: Table("one_row", MetaData(schema=ENV.schema), autoload_with=sync_conn)
+        )
+        assert len(one_row.c) == 1
+        assert one_row.c.number_of_rows is not None
+        assert one_row.comment == "table comment"
+
+    async def test_reflect_schemas(self, async_engine):
+        engine, conn = async_engine
+
+        def _inspect(sync_conn):
+            insp = sqlalchemy.inspect(sync_conn)
+            return insp.get_schema_names()
+
+        schemas = await conn.run_sync(_inspect)
+        assert ENV.schema in schemas
+        assert "default" in schemas
+
+    async def test_get_table_names(self, async_engine):
+        engine, conn = async_engine
+
+        def _inspect(sync_conn):
+            insp = sqlalchemy.inspect(sync_conn)
+            return insp.get_table_names(schema=ENV.schema)
+
+        table_names = await conn.run_sync(_inspect)
+        assert "many_rows" in table_names
+
+    async def test_has_table(self, async_engine):
+        engine, conn = async_engine
+
+        def _inspect(sync_conn):
+            insp = sqlalchemy.inspect(sync_conn)
+            return (
+                insp.has_table("one_row", schema=ENV.schema),
+                insp.has_table("this_table_does_not_exist", schema=ENV.schema),
+            )
+
+        exists, not_exists = await conn.run_sync(_inspect)
+        assert exists
+        assert not not_exists
+
+    async def test_get_columns(self, async_engine):
+        engine, conn = async_engine
+
+        def _inspect(sync_conn):
+            insp = sqlalchemy.inspect(sync_conn)
+            return insp.get_columns(table_name="one_row", schema=ENV.schema)
+
+        columns = await conn.run_sync(_inspect)
+        actual = columns[0]
+        assert actual["name"] == "number_of_rows"
+        assert isinstance(actual["type"], sqlalchemy.types.INTEGER)
+        assert actual["nullable"]
+        assert actual["default"] is None
+        assert not actual["autoincrement"]
+        assert actual["comment"] == "some comment"
+
+
+class TestAsyncDialectProperties:
+    async def test_aiorest_dialect(self):
+        engine = create_async_engine(_async_conn_str())
+        try:
+            assert engine.dialect.is_async is True
+            assert engine.dialect.driver == "aiorest"
+            assert engine.dialect.supports_statement_cache is True
+        finally:
+            await engine.dispose()
+
+    async def test_aiopandas_dialect(self):
+        conn_str = _async_conn_str().replace("+aiorest", "+aiopandas")
+        engine = create_async_engine(conn_str)
+        try:
+            assert engine.dialect.is_async is True
+            assert engine.dialect.driver == "aiopandas"
+        finally:
+            await engine.dispose()
+
+    async def test_aioarrow_dialect(self):
+        conn_str = _async_conn_str().replace("+aiorest", "+aioarrow")
+        engine = create_async_engine(conn_str)
+        try:
+            assert engine.dialect.is_async is True
+            assert engine.dialect.driver == "aioarrow"
+        finally:
+            await engine.dispose()
+
+    async def test_aiopolars_dialect(self):
+        conn_str = _async_conn_str().replace("+aiorest", "+aiopolars")
+        engine = create_async_engine(conn_str)
+        try:
+            assert engine.dialect.is_async is True
+            assert engine.dialect.driver == "aiopolars"
+        finally:
+            await engine.dispose()
+
+    async def test_aios3fs_dialect(self):
+        conn_str = _async_conn_str().replace("+aiorest", "+aios3fs")
+        engine = create_async_engine(conn_str)
+        try:
+            assert engine.dialect.is_async is True
+            assert engine.dialect.driver == "aios3fs"
+        finally:
+            await engine.dispose()

tests/sqlalchemy/__init__.py

@@ -6,3 +6,8 @@
 registry.register("awsathena.pandas", "pyathena.sqlalchemy.pandas", "AthenaPandasDialect")
 registry.register("awsathena.arrow", "pyathena.sqlalchemy.arrow", "AthenaArrowDialect")
 registry.register("awsathena.s3fs", "pyathena.sqlalchemy.s3fs", "AthenaS3FSDialect")
+registry.register("awsathena.aiorest", "pyathena.aio.sqlalchemy.rest", "AthenaAioRestDialect")
+registry.register("awsathena.aiopandas", "pyathena.aio.sqlalchemy.pandas", "AthenaAioPandasDialect")
+registry.register("awsathena.aioarrow", "pyathena.aio.sqlalchemy.arrow", "AthenaAioArrowDialect")
+registry.register("awsathena.aiopolars", "pyathena.aio.sqlalchemy.polars", "AthenaAioPolarsDialect")
+registry.register("awsathena.aios3fs", "pyathena.aio.sqlalchemy.s3fs", "AthenaAioS3FSDialect")

tests/sqlalchemy/conftest.py

@@ -13,13 +13,15 @@
     temp_table_keyword_args,
 )
 
-from tests import ENV, SQLALCHEMY_CONNECTION_STRING
+from tests import ASYNC_SQLALCHEMY_CONNECTION_STRING, ENV, SQLALCHEMY_CONNECTION_STRING
 
 
 def pytest_sessionstart(session):
-    conn_str = (
-        SQLALCHEMY_CONNECTION_STRING + "&tblproperties=" + quote_plus("'table_type'='ICEBERG'")
+    use_async = session.config.getoption("--dburi", None) == ["async"]
+    base_conn_str = (
+        ASYNC_SQLALCHEMY_CONNECTION_STRING if use_async else SQLALCHEMY_CONNECTION_STRING
     )
+    conn_str = base_conn_str + "&tblproperties=" + quote_plus("'table_type'='ICEBERG'")
     session.config.option.dburi = [
         conn_str.format(
             region_name=ENV.region_name,