From 3a27327d67eef89055cfb2c62968e515110ca0a3 Mon Sep 17 00:00:00 2001 From: Abdel Jaidi Date: Thu, 14 Mar 2024 13:48:05 +0000 Subject: [PATCH] fix: allow s3_output in athena.to_iceberg --- awswrangler/athena/_write_iceberg.py | 4 ++++ tests/unit/test_athena_iceberg.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/awswrangler/athena/_write_iceberg.py b/awswrangler/athena/_write_iceberg.py index 05236c264..3da02fc38 100644 --- a/awswrangler/athena/_write_iceberg.py +++ b/awswrangler/athena/_write_iceberg.py @@ -241,6 +241,7 @@ def to_iceberg( merge_cols: list[str] | None = None, keep_files: bool = True, data_source: str | None = None, + s3_output: str | None = None, workgroup: str = "primary", mode: Literal["append", "overwrite", "overwrite_partitions"] = "append", encryption: str | None = None, @@ -288,6 +289,8 @@ def to_iceberg( Whether staging files produced by Athena are retained. 'True' by default. data_source : str, optional Data Source / Catalog name. If None, 'AwsDataCatalog' will be used by default. + s3_output : str, optional + Amazon S3 path used for query execution. workgroup : str Athena workgroup. Primary by default. mode: str @@ -498,6 +501,7 @@ def to_iceberg( wg_config=wg_config, database=database, data_source=data_source, + s3_output=s3_output, encryption=encryption, kms_key=kms_key, boto3_session=boto3_session, diff --git a/tests/unit/test_athena_iceberg.py b/tests/unit/test_athena_iceberg.py index ca2e4ed3d..cf93e8930 100644 --- a/tests/unit/test_athena_iceberg.py +++ b/tests/unit/test_athena_iceberg.py @@ -68,6 +68,7 @@ def test_athena_to_iceberg( def test_athena_to_iceberg_append( path: str, path2: str, + path3: str, glue_database: str, glue_table: str, partition_cols: list[str] | None, @@ -102,6 +103,7 @@ def test_athena_to_iceberg_append( partition_cols=partition_cols, keep_files=False, mode="append", + s3_output=path3, ) df_actual = wr.athena.read_sql_query(