From 749efcd6bee6fdfce583ab239ec9940de9b687a5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 3 Aug 2025 12:28:59 +0900 Subject: [PATCH 1/2] Add comprehensive DDL compilation tests for complex data types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add DDL compilation tests for ARRAY, MAP, and STRUCT types in SQLAlchemy: - test_create_table_with_array_types: Tests ARRAY, nested arrays, struct arrays - test_create_table_with_map_types: Tests MAP types with complex values - test_create_table_with_struct_types: Tests ROW types with nested structures - test_create_table_with_complex_nested_types: Tests deeply nested combinations Key features tested: - Basic ARRAY, ARRAY types - Nested arrays: ARRAY> - Arrays of structs: ARRAY - Maps with struct values: MAP - Structs with array/map fields: ROW(tags ARRAY, ...) - Complex nested: ARRAY> All assertions use correct Athena SQL syntax with ROW(...) format for structs. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- tests/pyathena/sqlalchemy/test_base.py | 144 ++++++++++++++++++++++++- 1 file changed, 143 insertions(+), 1 deletion(-) diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 7013591b..f5f11791 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -17,7 +17,7 @@ from sqlalchemy.sql.schema import Column, MetaData, Table from sqlalchemy.sql.selectable import TextualSelect -from pyathena.sqlalchemy.types import TINYINT, AthenaStruct, Tinyint +from pyathena.sqlalchemy.types import TINYINT, AthenaArray, AthenaMap, AthenaStruct, Tinyint from tests.pyathena.conftest import ENV @@ -1997,3 +1997,145 @@ def test_compile_temporal_query_by_timestamp_with_hint(self, engine): f"SELECT count({ENV.schema}.{table_name}.col_1) AS count_1 \n" f"FROM {ENV.schema}.{table_name} FOR VERSION AS OF '{timestamp}'" ) + + def test_create_table_with_array_types(self, engine): + """Test DDL compilation for ARRAY types.""" + engine, conn = engine + table_name = "test_create_table_with_array_types" + table = Table( + table_name, + MetaData(schema=ENV.schema), + Column("id", types.Integer), + Column("tags", AthenaArray(types.String)), + Column("scores", AthenaArray(types.Integer)), + Column("nested_arrays", AthenaArray(AthenaArray(types.String))), + Column( + "struct_array", + AthenaArray(AthenaStruct(("name", types.String), ("age", types.Integer))), + ), + awsathena_location=f"{ENV.s3_staging_dir}{ENV.schema}/{table_name}/", + awsathena_file_format="PARQUET", + ) + + # Test DDL compilation + create_ddl = CreateTable(table).compile(dialect=engine.dialect) + ddl_string = str(create_ddl) + + # Verify ARRAY types are correctly compiled + assert "tags ARRAY" in ddl_string + assert "scores ARRAY" in ddl_string + assert "nested_arrays ARRAY>" in ddl_string + assert "struct_array ARRAY" in ddl_string + + def test_create_table_with_map_types(self, engine): + """Test DDL compilation for MAP types.""" + engine, conn = engine + table_name = "test_create_table_with_map_types" + table = Table( + table_name, + MetaData(schema=ENV.schema), + Column("id", types.Integer), + Column("attributes", AthenaMap(types.String, types.String)), + Column("metrics", AthenaMap(types.String, types.Integer)), + Column( + "complex_map", + AthenaMap( + types.String, AthenaStruct(("value", types.String), ("count", types.Integer)) + ), + ), + Column("nested_map", AthenaMap(types.String, AthenaArray(types.String))), + awsathena_location=f"{ENV.s3_staging_dir}{ENV.schema}/{table_name}/", + awsathena_file_format="PARQUET", + ) + + # Test DDL compilation + create_ddl = CreateTable(table).compile(dialect=engine.dialect) + ddl_string = str(create_ddl) + + # Verify MAP types are correctly compiled + assert "attributes MAP" in ddl_string + assert "metrics MAP" in ddl_string + assert "complex_map MAP" in ddl_string + assert "nested_map MAP>" in ddl_string + + def test_create_table_with_struct_types(self, engine): + """Test DDL compilation for STRUCT types.""" + engine, conn = engine + table_name = "test_create_table_with_struct_types" + table = Table( + table_name, + MetaData(schema=ENV.schema), + Column("id", types.Integer), + Column( + "user_info", + AthenaStruct( + ("name", types.String), ("age", types.Integer), ("email", types.String) + ), + ), + Column( + "nested_struct", + AthenaStruct( + ( + "personal", + AthenaStruct(("first_name", types.String), ("last_name", types.String)), + ), + ("preferences", AthenaMap(types.String, types.String)), + ), + ), + Column( + "struct_with_array", + AthenaStruct( + ("tags", AthenaArray(types.String)), ("scores", AthenaArray(types.Integer)) + ), + ), + awsathena_location=f"{ENV.s3_staging_dir}{ENV.schema}/{table_name}/", + awsathena_file_format="PARQUET", + ) + + # Test DDL compilation + create_ddl = CreateTable(table).compile(dialect=engine.dialect) + ddl_string = str(create_ddl) + + # Verify STRUCT types are correctly compiled + assert "user_info ROW(name STRING, age INTEGER, email STRING)" in ddl_string + assert ( + "nested_struct ROW(personal ROW(first_name STRING, last_name STRING), " + "preferences MAP)" in ddl_string + ) + assert "struct_with_array ROW(tags ARRAY, scores ARRAY)" in ddl_string + + def test_create_table_with_complex_nested_types(self, engine): + """Test DDL compilation for complex nested combinations of ARRAY, MAP, and STRUCT.""" + engine, conn = engine + table_name = "test_create_table_with_complex_nested_types" + table = Table( + table_name, + MetaData(schema=ENV.schema), + Column("id", types.Integer), + Column( + "data", + AthenaArray( + AthenaMap( + types.String, + AthenaStruct( + ("value", types.String), + ("metadata", AthenaMap(types.String, types.String)), + ("tags", AthenaArray(types.String)), + ), + ) + ), + ), + awsathena_location=f"{ENV.s3_staging_dir}{ENV.schema}/{table_name}/", + awsathena_file_format="PARQUET", + ) + + # Test DDL compilation + create_ddl = CreateTable(table).compile(dialect=engine.dialect) + ddl_string = str(create_ddl) + + # Verify complex nested type is correctly compiled + expected_type = ( + "data ARRAY, " + "tags ARRAY)>>" + ) + assert expected_type in ddl_string From 7f00044d06f429a6c23732312b36223dfc660fb7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 3 Aug 2025 12:46:23 +0900 Subject: [PATCH 2/2] Fix DDL test assertions to match actual compiler output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix MAP type assertions to include spaces in generated DDL: - MAP -> MAP - MAP -> MAP This matches the actual SQLAlchemy compiler output which includes spaces after commas in type parameter lists. Fixes CI test failures: - test_create_table_with_map_types - test_create_table_with_struct_types (nested MAP parts) - test_create_table_with_complex_nested_types 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- tests/pyathena/sqlalchemy/test_base.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index f5f11791..7419aeda 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -2053,10 +2053,10 @@ def test_create_table_with_map_types(self, engine): ddl_string = str(create_ddl) # Verify MAP types are correctly compiled - assert "attributes MAP" in ddl_string - assert "metrics MAP" in ddl_string - assert "complex_map MAP" in ddl_string - assert "nested_map MAP>" in ddl_string + assert "attributes MAP" in ddl_string + assert "metrics MAP" in ddl_string + assert "complex_map MAP" in ddl_string + assert "nested_map MAP>" in ddl_string def test_create_table_with_struct_types(self, engine): """Test DDL compilation for STRUCT types.""" @@ -2100,7 +2100,7 @@ def test_create_table_with_struct_types(self, engine): assert "user_info ROW(name STRING, age INTEGER, email STRING)" in ddl_string assert ( "nested_struct ROW(personal ROW(first_name STRING, last_name STRING), " - "preferences MAP)" in ddl_string + "preferences MAP)" in ddl_string ) assert "struct_with_array ROW(tags ARRAY, scores ARRAY)" in ddl_string @@ -2135,7 +2135,7 @@ def test_create_table_with_complex_nested_types(self, engine): # Verify complex nested type is correctly compiled expected_type = ( - "data ARRAY, " + "data ARRAY, " "tags ARRAY)>>" ) assert expected_type in ddl_string