Keep the connection's non-throttling retries on the column fallback path

Derive the single-attempt policy from the connection's RetryConfig by removing only the throttling codes, so other configured codes are still retried before a metadata request fails; the async adapter exposes retry_config for that. Sort information_schema rows by ordinal_position on the client and treat non-string comments as missing, covering UNLOAD-backed and pandas cursors.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: laughingman7743

docs/sqlalchemy.md

@@ -100,7 +100,8 @@ A later listing preserves metadata already fetched for a table.
 Table-metadata lookups propagate throttling and permission errors rather than reporting missing tables.
 `has_table()` propagates permission failures, including access denied by Lake Formation, instead of returning or caching `False`.
 For failed metadata requests, only recognized `EntityNotFoundException` responses establish absence; unrecognized errors are propagated rather than guessed to mean a missing table.
-Column reflection and `has_table()` make one table-metadata request; when Athena throttles it, they read `information_schema.columns` instead, executed without query result reuse, and log a warning.
+Column reflection and `has_table()` do not retry a throttled table-metadata request; when Athena reports throttling, they read `information_schema.columns` instead, executed without query result reuse, and log a warning.
+Other error codes listed in `RetryConfig.exceptions` are still retried with the connection's policy on that path; a `retry_config` in `cursor_kwargs` takes precedence.
 Columns reflected this way carry the type names Athena reports there, such as `varchar` for `string`, and partition columns are marked from its `extra_info` column.
 This fallback does not populate the table-metadata cache, and table comments and table options still come from the metadata API with the configured retries, so they propagate the throttling error.
 Athena applies its metadata API rate limits per account, and they are not listed in Service Quotas.

pyathena/aio/sqlalchemy/base.py

@@ -22,6 +22,7 @@
     ProgrammingError,
 )
 from pyathena.sqlalchemy.base import AthenaDialect
+from pyathena.util import RetryConfig
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -142,6 +143,10 @@ def schema_name(self) -> str | None:
     def cursor_kwargs(self) -> dict[str, Any]:
         return self._connection.cursor_kwargs  # type: ignore[no-any-return]
 
+    @property
+    def retry_config(self) -> RetryConfig:
+        return self._connection.retry_config  # type: ignore[no-any-return]
+
     def cursor(self, **kwargs: Any) -> AsyncAdapt_pyathena_cursor:
         raw_cursor = self._connection.cursor(**kwargs)
         return AsyncAdapt_pyathena_cursor(raw_cursor)

pyathena/sqlalchemy/base.py

@@ -57,10 +57,6 @@
 
 _logger = logging.getLogger(__name__)
 
-# Column reflection makes one metadata attempt and then answers from
-# information_schema, so a throttled request does not wait out the retry policy.
-_SINGLE_ATTEMPT = RetryConfig(attempt=1)
-
 
 ischema_names: dict[str, type[Any]] = {
     "boolean": types.BOOLEAN,
@@ -345,8 +341,12 @@ def _get_columns(self, connection, table_name: str, schema: str | None = None, *
         # A throttled metadata request switches to information_schema at once
         # instead of waiting out the retry policy; the query answers existence
         # and columns, while table comments and options still need the API.
+        # Other retryable codes keep the connection's policy.
+        retry_config = self._without_throttling_retries(
+            self._cursor_option(raw_connection, "retry_config")
+        )
         with raw_connection.driver_connection.cursor(  # type: ignore[union-attr]
-            retry_config=_SINGLE_ATTEMPT
+            retry_config=retry_config
         ) as cursor:
             try:
                 metadata = self._lookup_table(cursor, schema, name, table_name)
@@ -368,6 +368,16 @@ def _get_columns(self, connection, table_name: str, schema: str | None = None, *
         info_cache[metadata_key] = metadata
         return self._columns_from_metadata(metadata)
 
+    @staticmethod
+    def _without_throttling_retries(retry_config: RetryConfig) -> RetryConfig:
+        return RetryConfig(
+            exceptions=[c for c in retry_config.exceptions if c not in THROTTLING_ERROR_CODES],
+            attempt=retry_config.attempt,
+            multiplier=retry_config.multiplier,
+            max_delay=retry_config.max_delay,
+            exponential_base=retry_config.exponential_base,
+        )
+
     def _column(self, name: str | None, type_: str, comment: str | None, partition: bool | None):
         return {
             "name": name,
@@ -394,16 +404,24 @@ def _columns_from_information_schema(
         table_name = table_name.lower().replace("'", "''")
         with raw_connection.driver_connection.cursor() as cursor:  # type: ignore[union-attr]
             cursor.execute(
-                "SELECT column_name, data_type, comment, extra_info "
+                "SELECT ordinal_position, column_name, data_type, comment, extra_info "
                 "FROM information_schema.columns "
-                f"WHERE table_schema = '{schema}' AND table_name = '{table_name}' "
-                "ORDER BY ordinal_position",
+                f"WHERE table_schema = '{schema}' AND table_name = '{table_name}'",
                 result_reuse_enable=False,
             )
             rows = cursor.fetchall()
+        # Sort here: UNLOAD-backed cursors do not preserve ORDER BY. Cursors that
+        # read NULL as NaN must not turn a missing comment into a value.
         return [
-            self._column(column_name, data_type, comment, extra_info == "partition key" or None)
-            for column_name, data_type, comment, extra_info in rows
+            self._column(
+                column_name,
+                data_type,
+                comment if isinstance(comment, str) else None,
+                extra_info == "partition key" or None,
+            )
+            for _, column_name, data_type, comment, extra_info in sorted(
+                rows, key=lambda row: int(row[0])
+            )
         ]
 
     def _get_tables(self, connection, schema: str | None = None, **kw):

tests/pyathena/sqlalchemy/test_base.py

@@ -27,6 +27,7 @@
     Tinyint,
     get_double_type,
 )
+from pyathena.util import RetryConfig
 from tests.pyathena.conftest import ENV
 
 # Amazon S3 Tables tests need a pre-provisioned table-bucket catalog and namespace.
@@ -51,10 +52,11 @@ def unique_s3tables_table_name(base: str) -> str:
 
 class TestAthenaDialect:
     def test_columns_from_information_schema(self):
+        # Rows arrive unordered, and a cursor may read a NULL comment as NaN.
         rows = [
-            ("id", "integer", "identifier", None),
-            ("payload", "row(a integer, b array(varchar))", None, None),
-            ("dt", "varchar", None, "partition key"),
+            ("3", "dt", "varchar", float("nan"), "partition key"),
+            ("1", "id", "integer", "identifier", None),
+            ("2", "payload", "row(a integer, b array(varchar))", None, None),
         ]
         executed = []
 
@@ -75,16 +77,29 @@ def execute(operation, **kwargs):
         assert columns[0]["comment"] == "identifier"
         assert isinstance(columns[1]["type"], AthenaStruct)
         assert isinstance(columns[2]["type"], types.VARCHAR)
+        assert columns[2]["comment"] is None
         assert [column["dialect_options"]["awsathena_partition"] for column in columns] == [
             None,
             None,
             True,
         ]
         ((operation, kwargs),) = executed
         assert "WHERE table_schema = 'my_schema' AND table_name = 'o''neil'" in operation
-        assert "ORDER BY ordinal_position" in operation
         assert kwargs == {"result_reuse_enable": False}
 
+    def test_without_throttling_retries_keeps_other_codes(self):
+        policy = RetryConfig(
+            exceptions=("ThrottlingException", "InternalServerException"),
+            attempt=10,
+            multiplier=2,
+            max_delay=30,
+            exponential_base=3,
+        )
+        derived = AthenaDialect._without_throttling_retries(policy)
+        assert derived.exceptions == ("InternalServerException",)
+        assert (derived.attempt, derived.multiplier, derived.max_delay) == (10, 2, 30)
+        assert derived.exponential_base == 3
+
     def test_get_table_matches_long_names_case_insensitively(self):
         # GetTableMetadata rejects names over 128 characters, so the lookup lists
         # with a lowercase filter and matches the catalog's own casing.

tests/sqlalchemy/test_suite.py

@@ -356,25 +356,30 @@ def record_query(params, **kwargs):
     )
     def test_metadata_errors_do_not_establish_absence(self, connection, monkeypatch, code):
         raw_connection = _raw_connection(connection)
-        if code == "InternalServerException":
-            # A code the retry policy covers still does not reach the existence fallback.
+        retried = code == "InternalServerException"
+        if retried:
+            # A code the retry policy covers is retried with that policy, then
+            # propagates; it does not reach the information_schema fallback.
             monkeypatch.setattr(
                 raw_connection.retry_config, "exceptions", ("ThrottlingException", code)
             )
+            monkeypatch.setattr(raw_connection.retry_config, "multiplier", 0)
         message = (
             "Catalog error (Service: AmazonDataCatalog; Status Code: 400; "
             f"Error Code: {code}; Request ID: example; Proxy: null)"
             if code
             else "Table not found"
         )
         error = _metadata_error("MetadataException", message)
-        calls = _fail_get_table_metadata(monkeypatch, raw_connection, error)
+        calls = _fail_get_table_metadata(
+            monkeypatch, raw_connection, error, attempt=2 if retried else 1
+        )
         inspector = inspect(connection)
         for _ in range(2):
             with pytest.raises(OperationalError) as caught:
                 inspector.has_table("unavailable_metadata")
             assert caught.value.__cause__ is error
-        assert len(calls) == 2
+        assert len(calls) == (4 if retried else 2)
 
     @sa_testing.combinations((True, sa_testing.requires.schemas), False, argnames="use_schema")
     def test_has_table_cache_drop(self, connection, metadata, use_schema):