|
| 1 | +import contextlib |
1 | 2 | import re |
2 | 3 | import textwrap |
3 | 4 | import uuid |
|
26 | 27 | ) |
27 | 28 | from tests.pyathena.conftest import ENV |
28 | 29 |
|
| 30 | +# Amazon S3 Tables tests need a pre-provisioned table-bucket catalog and namespace. |
| 31 | +# Skip them unless AWS_ATHENA_S3_TABLES_CATALOG / AWS_ATHENA_S3_TABLES_NAMESPACE are set. |
| 32 | +requires_s3_tables = pytest.mark.skipif( |
| 33 | + not ENV.s3tables_catalog or not ENV.s3tables_namespace, |
| 34 | + reason="AWS_ATHENA_S3_TABLES_CATALOG / AWS_ATHENA_S3_TABLES_NAMESPACE are not configured", |
| 35 | +) |
| 36 | + |
| 37 | + |
| 38 | +def unique_s3tables_table_name(base: str) -> str: |
| 39 | + """Return a per-run-unique S3 Tables table name. |
| 40 | +
|
| 41 | + Other integration tests isolate themselves with a random per-process |
| 42 | + ``ENV.schema``, but the S3 Tables tests share one fixed namespace |
| 43 | + (``ENV.s3tables_namespace``). The CI matrix runs ``tests/pyathena`` once per |
| 44 | + Python version in parallel against the same account, so a fixed table name |
| 45 | + would collide across those concurrent jobs; a random suffix keeps them apart. |
| 46 | + """ |
| 47 | + return f"{base}_{uuid.uuid4().hex[:8]}" |
| 48 | + |
29 | 49 |
|
30 | 50 | class TestSQLAlchemyAthena: |
31 | 51 | @pytest.mark.parametrize( |
@@ -2019,6 +2039,104 @@ def test_create_iceberg_table_with_multiple_partition_transform(self, engine): |
2019 | 2039 | tblproperties = actual.dialect_options["awsathena"]["tblproperties"] |
2020 | 2040 | assert tblproperties["table_type"] == "ICEBERG" |
2021 | 2041 |
|
| 2042 | + @requires_s3_tables |
| 2043 | + @pytest.mark.parametrize("engine", [{"catalog_name": ENV.s3tables_catalog}], indirect=True) |
| 2044 | + def test_create_s3tables_iceberg_table(self, engine): |
| 2045 | + engine, conn = engine |
| 2046 | + # S3 Tables select the catalog via the connection ``catalog_name`` and use |
| 2047 | + # the namespace as the schema, so the DDL is a two-part ``namespace.table`` |
| 2048 | + # identifier with no LOCATION (managed storage). |
| 2049 | + schema = ENV.s3tables_namespace |
| 2050 | + table_name = unique_s3tables_table_name("test_create_s3tables_iceberg_table") |
| 2051 | + table = Table( |
| 2052 | + table_name, |
| 2053 | + MetaData(schema=schema), |
| 2054 | + Column("col_1", types.String), |
| 2055 | + Column("col_2", types.Integer), |
| 2056 | + awsathena_tblproperties={"table_type": "ICEBERG"}, |
| 2057 | + ) |
| 2058 | + ddl = CreateTable(table).compile(bind=conn) |
| 2059 | + |
| 2060 | + assert str(ddl) == textwrap.dedent( |
| 2061 | + f""" |
| 2062 | + CREATE TABLE {ENV.s3tables_namespace}.{table_name} ( |
| 2063 | + \tcol_1 STRING, |
| 2064 | + \tcol_2 INT |
| 2065 | + ) |
| 2066 | + TBLPROPERTIES ( |
| 2067 | + \t'table_type' = 'ICEBERG' |
| 2068 | + ) |
| 2069 | + """ |
| 2070 | + ) |
| 2071 | + |
| 2072 | + table.create(bind=conn) |
| 2073 | + try: |
| 2074 | + fqtn = f'"{ENV.s3tables_namespace}"."{table_name}"' |
| 2075 | + conn.execute(text(f"SELECT col_1, col_2 FROM {fqtn} LIMIT 0")) |
| 2076 | + finally: |
| 2077 | + with contextlib.suppress(Exception): |
| 2078 | + table.drop(bind=conn) |
| 2079 | + |
| 2080 | + @requires_s3_tables |
| 2081 | + @pytest.mark.parametrize("engine", [{"catalog_name": ENV.s3tables_catalog}], indirect=True) |
| 2082 | + def test_create_s3tables_iceberg_table_with_partition_transform(self, engine): |
| 2083 | + engine, conn = engine |
| 2084 | + schema = ENV.s3tables_namespace |
| 2085 | + table_name = unique_s3tables_table_name("test_create_s3tables_partition_transform") |
| 2086 | + table = Table( |
| 2087 | + table_name, |
| 2088 | + MetaData(schema=schema), |
| 2089 | + Column("col_1", types.String), |
| 2090 | + Column("dt", types.Date, awsathena_partition=True, awsathena_partition_transform="day"), |
| 2091 | + awsathena_tblproperties={"table_type": "ICEBERG"}, |
| 2092 | + ) |
| 2093 | + ddl = CreateTable(table).compile(bind=conn) |
| 2094 | + |
| 2095 | + assert str(ddl) == textwrap.dedent( |
| 2096 | + f""" |
| 2097 | + CREATE TABLE {ENV.s3tables_namespace}.{table_name} ( |
| 2098 | + \tcol_1 STRING, |
| 2099 | + \tdt DATE |
| 2100 | + ) |
| 2101 | + PARTITIONED BY ( |
| 2102 | + \tday(dt) |
| 2103 | + ) |
| 2104 | + TBLPROPERTIES ( |
| 2105 | + \t'table_type' = 'ICEBERG' |
| 2106 | + ) |
| 2107 | + """ |
| 2108 | + ) |
| 2109 | + |
| 2110 | + table.create(bind=conn) |
| 2111 | + with contextlib.suppress(Exception): |
| 2112 | + table.drop(bind=conn) |
| 2113 | + |
| 2114 | + @requires_s3_tables |
| 2115 | + @pytest.mark.parametrize("engine", [{"catalog_name": ENV.s3tables_catalog}], indirect=True) |
| 2116 | + def test_create_s3tables_table_as_select(self, engine): |
| 2117 | + engine, conn = engine |
| 2118 | + # CTAS is not modeled as a SQLAlchemy construct; exercise it as raw SQL. |
| 2119 | + # With the catalog selected on the connection the identifier is two-part |
| 2120 | + # (namespace.table); managed Iceberg CTAS requires is_external=false. |
| 2121 | + # Identifiers are left unquoted: DROP TABLE is Hive DDL and rejects the |
| 2122 | + # double quotes that the Trino CTAS/SELECT statements accept. |
| 2123 | + table_name = unique_s3tables_table_name("test_create_s3tables_table_as_select") |
| 2124 | + fqtn = f"{ENV.s3tables_namespace}.{table_name}" |
| 2125 | + conn.execute(text(f"DROP TABLE IF EXISTS {fqtn}")) |
| 2126 | + try: |
| 2127 | + conn.execute( |
| 2128 | + text( |
| 2129 | + f"CREATE TABLE {fqtn} " |
| 2130 | + "WITH (table_type = 'ICEBERG', is_external = false) AS " |
| 2131 | + "SELECT 1 AS id, 'a' AS name" |
| 2132 | + ) |
| 2133 | + ) |
| 2134 | + rows = conn.execute(text(f"SELECT id, name FROM {fqtn}")).fetchall() |
| 2135 | + assert rows == [(1, "a")] |
| 2136 | + finally: |
| 2137 | + with contextlib.suppress(Exception): |
| 2138 | + conn.execute(text(f"DROP TABLE IF EXISTS {fqtn}")) |
| 2139 | + |
2022 | 2140 | def test_insert_from_select_cte_follows_insert_one(self, engine): |
2023 | 2141 | engine, conn = engine |
2024 | 2142 | metadata = MetaData(schema=ENV.schema) |
|
0 commit comments