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

Signed-off-by: Daisuke Taniwaki <daisuketaniwaki@gmail.com>

Author: dtaniwaki

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

@@ -9,6 +9,9 @@
     {% 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 table_type = optional_args.get("table_type", "") %}
+    {% 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 +28,48 @@ 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 == "bucket":
+            return F.bucket(int(args[1]), F.col(args[0]))
+        if func == "truncate":
+            return F.truncate(int(args[1]), 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 }}/")
+    {% if extra_table_properties is not none %}
+    {% for prop_name, prop_value in extra_table_properties.items() %}
+    _writer = _writer.tableProperty("{{ prop_name }}", "{{ prop_value }}")
+    {% endfor %}
+    {% endif %}
+    {% if partitioned_by is not none %}
+    _writer = _writer.partitionedBy(
+        {%- for part_expr in partitioned_by %}
+        _parse_iceberg_partition("{{ part_expr }}"){{ "," 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,9 @@
   {%- endif -%}
 
   {%- if language == 'python' -%}
+    {%- set use_iceberg_write_to = config.get('use_iceberg_write_to', false) -%}
     {%- 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 +86,10 @@
           'write_compression': write_compression,
           'bucket_count': bucket_count,
           'field_delimiter': field_delimiter,
-          'spark_ctas': spark_ctas
+          'table_type': table_type,
+          '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,24 @@
     {%- 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 -%}
+      {%- 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' -%}