Merge pull request #784 from pyathena-dev/fix/metadata-string-reflection

Preserve Hive string types in metadata fallback reflection

Author: laughingman7743

docs/sqlalchemy.md

@@ -103,7 +103,8 @@ For failed metadata requests, only recognized `EntityNotFoundException` response
 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 the connection's `RetryConfig.exceptions` are still retried on that path, except `MetadataException` itself, which carries wrapped throttling; list the wrapped Glue codes instead.
 A `retry_config` in `cursor_kwargs` replaces that policy entirely, including its throttling retries, which then run before the fallback.
-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.
+The fallback maps unbounded `varchar` to SQLAlchemy `String`, matching Hive `STRING` reflection from the metadata API, and preserves explicit `VARCHAR(n)` and `CHAR(n)` lengths.
+Partition columns are marked from the `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's API retries use exponential backoff with uniform jitter; `RetryConfig` documents the default attempt count and waits.

pyathena/sqlalchemy/base.py

@@ -429,7 +429,8 @@ def _columns_from_information_schema(
         return [
             self._column(
                 column_name,
-                data_type,
+                # Athena exposes Hive STRING as unbounded VARCHAR in information_schema.
+                "string" if data_type == "varchar" else data_type,
                 comment if isinstance(comment, str) else None,
                 extra_info == "partition key" or None,
             )

tests/pyathena/sqlalchemy/test_base.py

@@ -78,7 +78,7 @@ def execute(operation, **kwargs):
         assert isinstance(columns[0]["type"], types.INTEGER)
         assert columns[0]["comment"] == "identifier"
         assert isinstance(columns[1]["type"], AthenaStruct)
-        assert isinstance(columns[2]["type"], types.VARCHAR)
+        assert type(columns[2]["type"]) is types.String
         assert columns[2]["comment"] is None
         assert [column["dialect_options"]["awsathena_partition"] for column in columns] == [
             None,

tests/sqlalchemy/test_suite.py

@@ -270,7 +270,10 @@ def test_preserves_table_metadata_until_clear_cache(self, connection, metadata):
         assert [column["name"] for column in inspector.get_columns(table.name)] == ["id", "added"]
 
     @sa_testing.combinations((String, None), (VARCHAR, 52), (CHAR, 52), argnames="type_,length")
-    def test_hive_string_length_reflection(self, connection, metadata, type_, length):
+    @sa_testing.combinations(False, True, argnames="information_schema")
+    def test_hive_string_length_reflection(
+        self, connection, metadata, type_, length, information_schema
+    ):
         table = Table(
             "string_length",
             metadata,
@@ -279,8 +282,17 @@ def test_hive_string_length_reflection(self, connection, metadata, type_, length
             awsathena_file_format="PARQUET",
         )
         table.create(connection)
-        reflected_type = inspect(connection).get_columns(table.name)[0]["type"]
-        assert isinstance(reflected_type, type_)
+        if information_schema:
+            # Exercise the fallback with a real query without forcing an API failure.
+            dialect = connection.dialect
+            raw_connection = dialect._raw_connection(connection)
+            columns = dialect._columns_from_information_schema(
+                raw_connection, raw_connection.schema_name, table.name
+            )
+        else:
+            columns = inspect(connection).get_columns(table.name)
+        reflected_type = columns[0]["type"]
+        assert type(reflected_type) is type_
         # Generic String compiles to Hive STRING without a length constraint.
         assert reflected_type.length == length