From 68b25aa4988550c39f7242b61f4f2a1e2a45cc60 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Wed, 30 Sep 2026 14:28:52 +0900 Subject: [PATCH 1/2] Backport #839: Resolve the bare awsathena scheme to the REST dialect (cherry picked from commit 0e72d33ba911e0e5a9d229f95f6f9c3f105ff57d) Conflict resolution: tests/pyathena/sqlalchemy/test_base.py conflicted because master's file carries helpers, imports, and the TestAthenaDialect class from PRs that are not backported to 3.x (the class first appeared with the metadata fallback work of #777). Only #839's own changes are applied: the AthenaRestDialect import and test_bare_scheme_uses_rest_driver. On 3.x the test goes into a new TestAthenaDialect class placed where master has it, directly before TestSQLAlchemyAthena. The added lines otherwise match #839's diff exactly. pyproject.toml and docs/sqlalchemy.md apply unchanged. Co-Authored-By: Claude Opus 5.5 --- docs/sqlalchemy.md | 2 +- pyproject.toml | 2 +- tests/pyathena/sqlalchemy/test_base.py | 14 ++++++++++++++ 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/docs/sqlalchemy.md b/docs/sqlalchemy.md index 7b75c0ea6..60fd19169 100644 --- a/docs/sqlalchemy.md +++ b/docs/sqlalchemy.md @@ -99,7 +99,7 @@ awsathena+aiorest://:@athena.{region_name}.amazonaws.com:443/{schema_name}?s3_st | Dialect | Driver | Schema | Cursor | |-----------|--------|------------------|------------------------| -| awsathena | | awsathena | DefaultCursor | +| awsathena | rest | awsathena | DefaultCursor | | awsathena | rest | awsathena+rest | DefaultCursor | | awsathena | pandas | awsathena+pandas | {ref}`pandas-cursor` | | awsathena | arrow | awsathena+arrow | {ref}`arrow-cursor` | diff --git a/pyproject.toml b/pyproject.toml index de1888064..1425c00c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ documentation = "https://pyathena.dev/" issues = "https://github.com/pyathena-dev/PyAthena/issues" [project.entry-points."sqlalchemy.dialects"] -awsathena = "pyathena.sqlalchemy.base:AthenaDialect" +awsathena = "pyathena.sqlalchemy.rest:AthenaRestDialect" "awsathena.rest" = "pyathena.sqlalchemy.rest:AthenaRestDialect" "awsathena.pandas" = "pyathena.sqlalchemy.pandas:AthenaPandasDialect" "awsathena.arrow" = "pyathena.sqlalchemy.arrow:AthenaArrowDialect" diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 5dd7b1d52..ebe3fbb0a 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -16,6 +16,7 @@ from sqlalchemy.sql.schema import Column, MetaData, Table from sqlalchemy.sql.selectable import TextualSelect +from pyathena.sqlalchemy.rest import AthenaRestDialect from pyathena.sqlalchemy.types import ( TINYINT, AthenaArray, @@ -50,6 +51,19 @@ def unique_s3tables_table_name(base: str) -> str: return f"{base}_{uuid.uuid4().hex[:8]}" +class TestAthenaDialect: + def test_bare_scheme_uses_rest_driver(self): + # The bare awsathena entry point resolves to the REST dialect, like + # awsathena+rest. Requires the package to be reinstalled (uv sync) so the + # installed entry point metadata matches pyproject.toml. + url = "awsathena://athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://bucket/path/" + bare = create_engine(url) + assert type(bare.dialect) is AthenaRestDialect + assert bare.driver == "rest" + assert bare.url.get_driver_name() == "rest" + assert bare.dialect.dialect_description == "awsathena+rest" + + class TestSQLAlchemyAthena: @pytest.mark.parametrize( "engine", From bc5fc9378c2a0bf68aed9113c134472f3ce8d690 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 30 Sep 2026 15:02:59 +0900 Subject: [PATCH 2/2] Backport #905: Align the compliance-suite dialect registry with the entry points tests/sqlalchemy/__init__.py registers the dialects for the SQLAlchemy compliance suite, and a registration overrides the installed entry point in any process that imports it. After #839 pointed the bare awsathena entry point at AthenaRestDialect, the file still registered the base AthenaDialect, and it never registered awsathena.polars. Register the same classes as pyproject.toml, and add a test that compares the file's registrations with the installed entry points. (cherry picked from commit d6b930356837a7b39c8ed99bf4548e2dc9925e39) Conflict resolution: tests/pyathena/sqlalchemy/test_base.py conflicted because master's import block and the tests around TestAthenaDialect come from PRs that are not on 3.x. Only #905's own lines are applied: the importlib.metadata, runpy, Path, and PluginLoader imports and test_compliance_suite_registry_matches_entry_points, placed after test_bare_scheme_uses_rest_driver as on master. The added lines match #905's commit exactly. tests/sqlalchemy/__init__.py applies unchanged apart from context (3.x's copy has no license header). Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/sqlalchemy/test_base.py | 18 ++++++++++++++++++ tests/sqlalchemy/__init__.py | 3 ++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index ebe3fbb0a..2863fb77e 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -1,8 +1,11 @@ +import importlib.metadata import re +import runpy import textwrap import uuid from datetime import date, datetime from decimal import Decimal +from pathlib import Path from urllib.parse import quote_plus import numpy as np @@ -15,6 +18,7 @@ from sqlalchemy.sql.ddl import CreateTable from sqlalchemy.sql.schema import Column, MetaData, Table from sqlalchemy.sql.selectable import TextualSelect +from sqlalchemy.util import PluginLoader from pyathena.sqlalchemy.rest import AthenaRestDialect from pyathena.sqlalchemy.types import ( @@ -63,6 +67,20 @@ def test_bare_scheme_uses_rest_driver(self): assert bare.url.get_driver_name() == "rest" assert bare.dialect.dialect_description == "awsathena+rest" + def test_compliance_suite_registry_matches_entry_points(self, monkeypatch): + # tests/sqlalchemy registers the dialects for the compliance suite, and a + # registration overrides the installed entry point in that process. + loader = PluginLoader("sqlalchemy.dialects") + monkeypatch.setattr("sqlalchemy.dialects.registry", loader) + runpy.run_path(str(Path(__file__).parents[2] / "sqlalchemy" / "__init__.py")) + entry_points = { + entry_point.name: entry_point.load() + for entry_point in importlib.metadata.entry_points(group="sqlalchemy.dialects") + if entry_point.value.startswith("pyathena.") + } + assert entry_points + assert {name: load() for name, load in loader.impls.items()} == entry_points + class TestSQLAlchemyAthena: @pytest.mark.parametrize( diff --git a/tests/sqlalchemy/__init__.py b/tests/sqlalchemy/__init__.py index 9a9dea779..b827aa6fe 100644 --- a/tests/sqlalchemy/__init__.py +++ b/tests/sqlalchemy/__init__.py @@ -1,9 +1,10 @@ from sqlalchemy.dialects import registry -registry.register("awsathena", "pyathena.sqlalchemy.base", "AthenaDialect") +registry.register("awsathena", "pyathena.sqlalchemy.rest", "AthenaRestDialect") registry.register("awsathena.rest", "pyathena.sqlalchemy.rest", "AthenaRestDialect") registry.register("awsathena.pandas", "pyathena.sqlalchemy.pandas", "AthenaPandasDialect") registry.register("awsathena.arrow", "pyathena.sqlalchemy.arrow", "AthenaArrowDialect") +registry.register("awsathena.polars", "pyathena.sqlalchemy.polars", "AthenaPolarsDialect") 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")