Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions tests/pyathena/sqlalchemy/test_base.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
import contextlib
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 types import SimpleNamespace
from urllib.parse import quote_plus

Expand All @@ -18,6 +21,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.aio.sqlalchemy.base import AthenaAioDialect
from pyathena.converter import DefaultTypeConverter
Expand Down Expand Up @@ -122,6 +126,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):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round two (claims, callers, operations). Base 8679f04, head d6b9303.

Claims checked:

  • "A registration overrides the installed entry point": PluginLoader.load returns self.impls[name]() before it consults entry points. Reproduced: importing tests.sqlalchemy on master makes a bare URL create AthenaDialect; with this change it creates AthenaRestDialect.
  • "An unregistered name falls back to the entry point": the same method's last loop, so the missing polars registration had no runtime effect.
  • "CI and the just test recipes run the suites in separate processes, and the compliance suite connects only with +rest / +aiorest": justfile runs tests/pyathena/ and tests/sqlalchemy/ in separate pytest commands; tests/__init__.py builds only +rest / +aiorest URLs.
  • "3.x has the same file content": false as written. 3.x's copy has the same registrations but no license header (Add MIT headers to existing files written solely by the maintainer聽#805 is not on 3.x). The PR body now says so; the cherry-pick applies to the registration lines only.
  • CI note: the changed paths match the sqla filter (^tests/sqlalchemy/, ^tests/pyathena/(aio/)?sqlalchemy/) and not the spark one, so Ready runs the PyAthena suite without Spark plus both compliance suites.

Callers: only test harness code changes; no library behavior or AWS request path changes.

Result: CLEAN after the PR body correction above.

# 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"))

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent review, relayed. Reviewer: Codex CLI 0.157.1, model gpt-6-sol, reasoning effort high, session 01a0f0e9-ceb7-7911-8021-a2fa694d6043. Static review in a read-only sandbox of a detached snapshot at head d6b9303 (base 8679f04); it received the literal diff without the PR description, commit message, or prior findings, and ran no tests. The snapshot and the PR worktree were unchanged afterwards.

Covered: the two changed files, the pyproject.toml entry points, the compliance suite and its configuration, test commands and CI selection, and AGENTS.md. The snapshot had no .venv, so the reviewer did not read SQLAlchemy's PluginLoader source; round two quotes it.

Result: CLEAN. All 11 registrations match the entry points. The test locates the file from __file__ and restores the registry through monkeypatch, with no working-directory, ordering, or pytest-xdist issue found. The compliance suites use explicit +rest / +aiorest URLs, so the changed bare registration and the added polars registration do not change their dialects.

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

@pytest.mark.parametrize("dialect_class", [AthenaDialect, AthenaAioDialect])
def test_type_compiler(self, dialect_class):
# SQLAlchemy 2.0 builds the type compiler from type_compiler_cls. A legacy
Expand Down
3 changes: 2 additions & 1 deletion tests/sqlalchemy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@

from sqlalchemy.dialects import registry

registry.register("awsathena", "pyathena.sqlalchemy.base", "AthenaDialect")
registry.register("awsathena", "pyathena.sqlalchemy.rest", "AthenaRestDialect")

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round one (implementation behavior). Base 8679f04, head d6b9303.

Covered: both changed files; how PluginLoader.load resolves a name (registered impls first, then auto_fn, then entry points); every pyathena entry point against the file's registrations; test isolation under xdist; the compliance-suite collection.

Result: CLEAN.

  • The file now registers exactly the eleven pyathena entry points; before, awsathena pointed at AthenaDialect and awsathena.polars was absent.
  • The new test swaps sqlalchemy.dialects.registry for an empty PluginLoader through monkeypatch before runpy.run_path, so it sees only this file's registrations and does not depend on which dialects other tests loaded into the global registry. monkeypatch restores the global registry afterwards.
  • The test fails on master's file with both differences and passes after the change. Loading the entry points imports only pyathena dialect modules; the polars dialect modules import polars only under TYPE_CHECKING.
  • pytest tests/sqlalchemy/ --collect-only collects 1630 tests with and without --dburi async, the same as master.

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")
Expand Down
Loading