feat: add conditional_merge support to to_iceberg

Author: wasim-builds

awswrangler/athena/_write_iceberg.py

@@ -240,7 +240,8 @@ def _validate_args(
     mode: Literal["append", "overwrite", "overwrite_partitions"],
     partition_cols: list[str] | None,
     merge_cols: list[str] | None,
-    merge_condition: Literal["update", "ignore"],
+    merge_condition: Literal["update", "ignore", "conditional_merge"],
+    conditional_merge_string: str | None,
 ) -> None:
     if df.empty is True:
         raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
@@ -260,9 +261,13 @@ def _validate_args(
                 "When mode is 'overwrite_partitions' merge_cols must not be specified."
             )
 
-    if merge_cols and merge_condition not in ["update", "ignore"]:
+    if merge_cols and merge_condition not in ["update", "ignore", "conditional_merge"]:
         raise exceptions.InvalidArgumentValue(
-            f"Invalid merge_condition: {merge_condition}. Valid values: ['update', 'ignore']"
+            f"Invalid merge_condition: {merge_condition}. Valid values: ['update', 'ignore', 'conditional_merge']"
+        )
+    if merge_cols and merge_condition == "conditional_merge" and not conditional_merge_string:
+        raise exceptions.InvalidArgumentCombination(
+            "When merge_condition is 'conditional_merge', conditional_merge_string must be specified."
         )
 
 
@@ -300,7 +305,8 @@ def _merge_iceberg(
     table: str,
     source_table: str,
     merge_cols: list[str] | None = None,
-    merge_condition: Literal["update", "ignore"] = "update",
+    merge_condition: Literal["update", "ignore", "conditional_merge"] = "update",
+    conditional_merge_string: str | None = None,
     merge_match_nulls: bool = False,
     partition_cols: list[str] | None = None,
     kms_key: str | None = None,
@@ -332,7 +338,9 @@ def _merge_iceberg(
 
         https://docs.aws.amazon.com/athena/latest/ug/merge-into-statement.html
     merge_condition: str, optional
-        The condition to be used in the MERGE INTO statement. Valid values: ['update', 'ignore'].
+        The condition to be used in the MERGE INTO statement. Valid values: ['update', 'ignore', 'conditional_merge'].
+    conditional_merge_string: str, optional
+        The condition string to be used if merge_condition is 'conditional_merge'.
     merge_match_nulls: bool, optional
         Instruct whether to have nulls in the merge condition match other nulls
     partition_cols: List[str], optional
@@ -364,6 +372,10 @@ def _merge_iceberg(
         if merge_condition == "update":
             match_condition = f"""WHEN MATCHED THEN
                 UPDATE SET {", ".join([f'"{x}" = source."{x}"' for x in df.columns])}"""
+        elif merge_condition == "conditional_merge":
+            match_condition = f"""WHEN MATCHED AND
+                {conditional_merge_string} THEN
+                UPDATE SET {", ".join([f'"{x}" = source."{x}"' for x in df.columns])}"""
         else:
             match_condition = ""
 
@@ -417,7 +429,8 @@ def to_iceberg(  # noqa: PLR0913
     table_location: str | None = None,
     partition_cols: list[str] | None = None,
     merge_cols: list[str] | None = None,
-    merge_condition: Literal["update", "ignore"] = "update",
+    merge_condition: Literal["update", "ignore", "conditional_merge"] = "update",
+    conditional_merge_string: str | None = None,
     merge_match_nulls: bool = False,
     keep_files: bool = True,
     data_source: str | None = None,
@@ -467,8 +480,10 @@ def to_iceberg(  # noqa: PLR0913
 
         https://docs.aws.amazon.com/athena/latest/ug/merge-into-statement.html
     merge_condition
-        The condition to be used in the MERGE INTO statement. Valid values: ['update', 'ignore'].
+        The condition to be used in the MERGE INTO statement. Valid values: ['update', 'ignore', 'conditional_merge'].
         Default is ``update``.
+    conditional_merge_string
+        The condition string to be used if merge_condition is 'conditional_merge'.
     merge_match_nulls
         Instruct whether to have nulls in the merge condition match other nulls.
     keep_files
@@ -556,6 +571,7 @@ def to_iceberg(  # noqa: PLR0913
         partition_cols=partition_cols,
         merge_cols=merge_cols,
         merge_condition=merge_condition,
+        conditional_merge_string=conditional_merge_string,
     )
 
     glue_table_settings = glue_table_settings if glue_table_settings else {}
@@ -681,6 +697,7 @@ def to_iceberg(  # noqa: PLR0913
             source_table=temp_table,
             merge_cols=merge_cols,
             merge_condition=merge_condition,
+            conditional_merge_string=conditional_merge_string,
             merge_match_nulls=merge_match_nulls,
             partition_cols=partition_cols,
             kms_key=kms_key,