feat(athena): add use_iceberg_write_to config for Iceberg Python models

dbt-athena/.changes/unreleased/Features-20260422-013953.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: Add use_iceberg_write_to config for Iceberg Python models
+time: 2026-04-22T01:39:53.640855+09:00
+custom:
+    Author: dtaniwaki
+    Issue: "1882"

dbt-athena/src/dbt/include/athena/macros/adapters/python_submissions.sql

@@ -9,6 +9,8 @@
     {% set merge_schema = optional_args.get("merge_schema", true) %}
     {% set bucket_count = optional_args.get("bucket_count") %}
     {% set field_delimiter = optional_args.get("field_delimiter") %}
+    {% set extra_table_properties = optional_args.get("extra_table_properties") %}
+    {% set use_iceberg_write_to = optional_args.get("use_iceberg_write_to", false) %}
     {% set spark_ctas = optional_args.get("spark_ctas", "") %}
 
 import pyspark
@@ -25,7 +27,51 @@ def materialize(spark_session, df, target_relation):
         msg = f"{type(df)} is not a supported type for dbt Python materialization"
         raise Exception(msg)
 
-{% if spark_ctas|length > 0 %}
+{% if use_iceberg_write_to %}
+    import re
+    from pyspark.sql import functions as F
+
+    def _parse_iceberg_partition(expr_str):
+        expr_str = expr_str.strip()
+        m = re.match(r"(\w+)\((.+)\)", expr_str)
+        if not m:
+            return F.col(expr_str)
+        func = m.group(1).lower()
+        args = [a.strip() for a in m.group(2).split(",")]
+        if func in ("day", "days"):
+            return F.days(F.col(args[0]))
+        if func in ("month", "months"):
+            return F.months(F.col(args[0]))
+        if func in ("year", "years"):
+            return F.years(F.col(args[0]))
+        if func in ("hour", "hours"):
+            return F.hours(F.col(args[0]))
+        if func in ("bucket", "truncate"):
+            if len(args) != 2:
+                raise ValueError(
+                    f"Iceberg partition transform '{func}' requires 2 arguments (column, n), got: {expr_str}"
+                )
+            n = int(args[1])
+            return F.bucket(n, F.col(args[0])) if func == "bucket" else F.truncate(n, F.col(args[0]))
+        raise ValueError(f"Unknown Iceberg partition transform: {func}")
+
+    _writer = df.writeTo("{{ target_relation.schema | replace('\"', '`') }}.{{ target_relation.identifier | replace('\"', '`') }}")
+    _writer = _writer.using("iceberg")
+    _writer = _writer.tableProperty("location", {{ (location ~ "/") | tojson }})
+    {% if extra_table_properties is not none %}
+    {% for prop_name, prop_value in extra_table_properties.items() %}
+    _writer = _writer.tableProperty({{ prop_name | tojson }}, {{ prop_value | string | tojson }})
+    {% endfor %}
+    {% endif %}
+    {% if partitioned_by is not none %}
+    _writer = _writer.partitionedBy(
+        {%- for part_expr in partitioned_by %}
+        _parse_iceberg_partition({{ part_expr | tojson }}){{ "," if not loop.last }}
+        {%- endfor %}
+    )
+    {% endif %}
+    _writer.createOrReplace()
+{% elif spark_ctas|length > 0 %}
     df.createOrReplaceTempView("{{ target_relation.schema}}_{{ target_relation.identifier }}_tmpvw")
     spark_session.sql("""
     {{ spark_ctas }}

dbt-athena/src/dbt/include/athena/macros/materializations/models/table/create_table_as.sql

@@ -43,8 +43,12 @@
   {%- endif -%}
 
   {%- if language == 'python' -%}
+    {%- set use_iceberg_write_to = config.get('use_iceberg_write_to', false) -%}
+    {%- if use_iceberg_write_to and table_type != 'iceberg' -%}
+      {{ exceptions.raise_compiler_error("The 'use_iceberg_write_to' config is only supported when table_type='iceberg'.") }}
+    {%- endif -%}
     {%- set spark_ctas = '' -%}
-    {%- if table_type == 'iceberg' -%}
+    {%- if table_type == 'iceberg' and not use_iceberg_write_to -%}
       {%- set spark_ctas -%}
           create table {{ relation.schema | replace('\"', '`') }}.{{ relation.identifier | replace('\"', '`') }}
           using iceberg
@@ -85,7 +89,9 @@
           'write_compression': write_compression,
           'bucket_count': bucket_count,
           'field_delimiter': field_delimiter,
-          'spark_ctas': spark_ctas
+          'extra_table_properties': extra_table_properties,
+          'use_iceberg_write_to': use_iceberg_write_to,
+          'spark_ctas': spark_ctas,
         }
       )
     }}

dbt-athena/src/dbt/include/athena/macros/materializations/models/table/table.sql

@@ -91,7 +91,27 @@
     {%- endif -%}
   {%- else -%}
 
-    {%- if old_relation is none -%}
+    {%- set use_iceberg_write_to = language == 'python' and config.get('use_iceberg_write_to', false) -%}
+
+    {%- if use_iceberg_write_to -%}
+      -- Python + use_iceberg_write_to: writeTo().createOrReplace() handles atomic replacement,
+      -- so we skip the __ha intermediate table and write directly to target.
+      -- Clean up leftover __ha / __bkp tables from previous HA-flow failures.
+      {%- if old_tmp_relation is not none -%}
+        {%- do drop_relation(old_tmp_relation) -%}
+      {%- endif -%}
+      {%- if old_bkp_relation is not none -%}
+        {%- do drop_relation(old_bkp_relation) -%}
+      {%- endif -%}
+      {%- if old_relation is not none and old_relation.is_view -%}
+        {%- do drop_relation(old_relation) -%}
+      {%- endif -%}
+      {%- set query_result = safe_create_table_as(False, target_relation, compiled_code, language, force_batch) -%}
+      {% call statement('create_table', language=language) %}
+        {{ query_result }}
+      {% endcall %}
+
+    {%- elif old_relation is none -%}
       {%- set query_result = safe_create_table_as(False, target_relation, compiled_code, language, force_batch) -%}
       -- Execute python code that is available in query result object
       {%- if language == 'python' -%}

dbt-athena/tests/functional/adapter/test_use_iceberg_write_to.py

@@ -0,0 +1,187 @@
+"""Functional tests for the use_iceberg_write_to config.
+
+The new branch in athena__py_save_table_as routes Iceberg Python models
+through DataFrameWriterV2 (writeTo().createOrReplace()) instead of the
+spark_ctas SQL path, enabling Iceberg-native partition transforms
+(day/bucket/...) and atomic replacement without the __ha intermediate
+table.
+
+Gated on DBT_TEST_ATHENA_SPARK_WORK_GROUP — must point to an Athena
+Spark workgroup with Iceberg support.
+"""
+
+import os
+
+import pytest
+
+from dbt.tests.util import run_dbt
+
+requires_spark_workgroup = pytest.mark.skipif(
+    not os.getenv("DBT_TEST_ATHENA_SPARK_WORK_GROUP"),
+    reason="DBT_TEST_ATHENA_SPARK_WORK_GROUP must point to an Athena Spark workgroup.",
+)
+
+
+_iceberg_writeto_partitioned = """
+def model(dbt, spark_session):
+    dbt.config(
+        materialized='table',
+        table_type='iceberg',
+        use_iceberg_write_to=True,
+        partitioned_by=['day(created_at)', 'bucket(user_id, 4)'],
+    )
+    from pyspark.sql import functions as F
+    rows = [
+        (1, '2026-01-01 00:00:00', 'a'),
+        (2, '2026-01-02 00:00:00', 'b'),
+        (3, '2026-01-03 00:00:00', 'c'),
+    ]
+    df = spark_session.createDataFrame(rows, ['user_id', 'created_at', 'name'])
+    return df.withColumn('created_at', F.to_timestamp('created_at'))
+"""
+
+
+@requires_spark_workgroup
+class TestUseIcebergWriteToPartitioned:
+    """writeTo() with day/bucket partition transforms must produce a working
+    Iceberg table. The spark_ctas path can't express these transforms, so
+    this case is the primary motivation for the feature."""
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {"iceberg_writeto_partitioned.py": _iceberg_writeto_partitioned}
+
+    def test_writes_rows_and_applies_iceberg_partition_transforms(self, project):
+        results = run_dbt(["run"])
+        assert len(results) == 1
+        assert results[0].status == "success"
+
+        records = sorted(
+            project.run_sql(
+                f"select user_id, name from {project.test_schema}.iceberg_writeto_partitioned",
+                fetch="all",
+            )
+        )
+        assert records == [(1, "a"), (2, "b"), (3, "c")]
+
+        # SHOW CREATE TABLE confirms day()/bucket() were applied as Iceberg
+        # partition transforms — not coerced into plain identity partitions
+        # (which is what the spark_ctas SQL path produces). Athena renders
+        # the column names with backticks in the partition spec.
+        ddl = project.run_sql(
+            f"show create table {project.test_schema}.iceberg_writeto_partitioned",
+            fetch="all",
+        )
+        ddl_text = "\n".join(row[0] for row in ddl)
+        assert "day(`created_at`)" in ddl_text
+        assert "bucket(4, `user_id`)" in ddl_text
+
+    def test_idempotent_replace(self, project):
+        # createOrReplace must succeed against an existing target without
+        # falling through the legacy __ha rename flow.
+        run_dbt(["run"])
+        results = run_dbt(["run"])
+        assert all(r.status == "success" for r in results)
+
+
+_iceberg_writeto_unpartitioned = """
+def model(dbt, spark_session):
+    dbt.config(
+        materialized='table',
+        table_type='iceberg',
+        use_iceberg_write_to=True,
+    )
+    return spark_session.createDataFrame([(1,), (2,), (3,)], ['id'])
+"""
+
+
+@requires_spark_workgroup
+class TestUseIcebergWriteToUnpartitioned:
+    """use_iceberg_write_to without partitioned_by should still succeed
+    (the partitionedBy() call is conditional on partitioned_by being set)."""
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {"iceberg_writeto_unpartitioned.py": _iceberg_writeto_unpartitioned}
+
+    def test_writes_rows_without_partitions(self, project):
+        results = run_dbt(["run"])
+        assert len(results) == 1
+        assert results[0].status == "success"
+
+        records = sorted(
+            project.run_sql(
+                f"select id from {project.test_schema}.iceberg_writeto_unpartitioned",
+                fetch="all",
+            )
+        )
+        assert records == [(1,), (2,), (3,)]
+
+
+_iceberg_writeto_with_table_properties = """
+def model(dbt, spark_session):
+    dbt.config(
+        materialized='table',
+        table_type='iceberg',
+        use_iceberg_write_to=True,
+        table_properties={'write.parquet.compression-codec': 'zstd'},
+    )
+    return spark_session.createDataFrame([(1,), (2,)], ['id'])
+"""
+
+
+@requires_spark_workgroup
+class TestUseIcebergWriteToTableProperties:
+    """table_properties must be forwarded as DataFrameWriterV2 .tableProperty()
+    calls and survive JSON escaping."""
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {"iceberg_writeto_props.py": _iceberg_writeto_with_table_properties}
+
+    def test_table_properties_are_propagated(self, project):
+        results = run_dbt(["run"])
+        assert len(results) == 1
+        assert results[0].status == "success"
+
+        records = sorted(
+            project.run_sql(
+                f"select id from {project.test_schema}.iceberg_writeto_props",
+                fetch="all",
+            )
+        )
+        assert records == [(1,), (2,)]
+
+        # Iceberg surfaces table properties through the $properties metadata
+        # table rather than Glue TBLPROPERTIES.
+        props = project.run_sql(
+            f'select key, value from "{project.test_schema}"."iceberg_writeto_props$properties"',
+            fetch="all",
+        )
+        props_dict = dict(props)
+        assert props_dict.get("write.parquet.compression-codec") == "zstd"
+
+
+_iceberg_writeto_on_hive_table = """
+def model(dbt, spark_session):
+    dbt.config(
+        materialized='table',
+        table_type='hive',
+        use_iceberg_write_to=True,
+    )
+    return spark_session.createDataFrame([(1,)], ['id'])
+"""
+
+
+@requires_spark_workgroup
+class TestUseIcebergWriteToRequiresIceberg:
+    """use_iceberg_write_to with a non-iceberg table_type must surface a
+    compiler error rather than silently writing the wrong format."""
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {"iceberg_writeto_invalid.py": _iceberg_writeto_on_hive_table}
+
+    def test_compiles_to_error(self, project):
+        results = run_dbt(["run"], expect_pass=False)
+        assert any("use_iceberg_write_to" in (r.message or "") for r in results)