bugfix - PR #179 - *Delta > Writer*: Fixed checking merge builder type in DeltaTableWriter

There was a small error in validating merge builder type when using the `merge_builder` param with `DeltaTableWriter`. It required that the value passed was either a `list` or a `delta.tables.DeltaMergeBuilder`
instance so validation failed when an instance of `delta.connect.tables.DeltaMergeBuilder` was passed.

## Motivation and Context
One cannot pass a correct DeltaMergeBuilder instance as a merge builder
for `DeltaTableWriter` if it's not a `delta.tables.DeltaMergeBuilder`
instance or a `list`. Connect merge builders should also be supported.

## How Has This Been Tested?
4 dedicated test cases have been added to test the merge builder
validation.

Author: zarembat

.github/workflows/test.yml

@@ -39,19 +39,18 @@ jobs:
       uses: actions/checkout@v4
       with:
         fetch-depth: 0
-        ref: ${{ github.event.pull_request.head.ref }}
+        ref: ${{ github.event.pull_request.head.sha }}
         repository: ${{ github.event.pull_request.head.repo.full_name }}
-    - name: Fetch target branch
-      run: git fetch origin ${{ github.event.pull_request.head.ref || 'main'}}:${{ github.event.pull_request.base.ref || 'main'}}
+
     - name: Check changes
       id: check
       run: |
-          # Set the base reference for the git diff
-          BASE_REF=${{ github.event.pull_request.base.ref || 'main' }}
-        
-          # Check for changes in this PR / commit
-          git_diff_output=$(git diff --name-only $BASE_REF ${{ github.event.after }})
-        
+          # Fetch the base branch
+          git fetch origin ${{ github.event.pull_request.base.sha }}
+
+          # Get the diff between PR base and head
+          git_diff_output=$(git diff --name-only ${{ github.event.pull_request.base.sha }} ${{ github.event.pull_request.head.sha }})
+
           # Count the number of changes to Python and TOML files
           python_changed=$(echo "$git_diff_output" | grep '\.py$' | wc -l)
           toml_changed=$(echo "$git_diff_output" | grep '\.toml$' | wc -l)

src/koheesio/spark/writers/delta/batch.py

@@ -328,11 +328,11 @@ def _validate_params(cls, params: dict) -> dict:
                     clause = merge_conf.get("clause")
                     if clause not in valid_clauses:
                         raise ValueError(f"Invalid merge clause '{clause}' provided")
-            elif (
-                not isinstance(merge_builder, DeltaMergeBuilder)
-                or not type(merge_builder).__name__ == "DeltaMergeBuilder"
+            elif not (
+                isinstance(merge_builder, DeltaMergeBuilder)
+                or type(merge_builder).__name__ == "DeltaMergeBuilder"
             ):
-                raise ValueError("merge_builder must be a list or merge clauses or a DeltaMergeBuilder instance")
+                raise ValueError("merge_builder must be a list of merge clauses or a DeltaMergeBuilder instance")
 
         return params
 

tests/spark/writers/delta/test_delta_writer.py

@@ -2,6 +2,7 @@
 from unittest.mock import MagicMock, patch
 
 from conftest import await_job_completion
+from delta.tables import DeltaMergeBuilder
 import pytest
 
 from pydantic import ValidationError
@@ -24,7 +25,9 @@
 
 def test_delta_table_writer(dummy_df, spark):
     table_name = "test_table"
-    writer = DeltaTableWriter(table=table_name, output_mode=BatchOutputMode.APPEND, df=dummy_df)
+    writer = DeltaTableWriter(
+        table=table_name, output_mode=BatchOutputMode.APPEND, df=dummy_df
+    )
     writer.execute()
     actual_count = spark.read.table(table_name).count()
     assert actual_count == 1
@@ -44,7 +47,10 @@ def test_delta_table_writer(dummy_df, spark):
 def test_delta_partitioning(spark, sample_df_to_partition):
     table_name = "partition_table"
     DeltaTableWriter(
-        table=table_name, output_mode=BatchOutputMode.OVERWRITE, df=sample_df_to_partition, partition_by=["partition"]
+        table=table_name,
+        output_mode=BatchOutputMode.OVERWRITE,
+        df=sample_df_to_partition,
+        partition_by=["partition"],
     ).execute()
     output_df = spark.read.table(table_name)
     assert output_df.count() == 2
@@ -55,7 +61,11 @@ def test_delta_table_merge_all(spark):
 
     table_name = "test_merge_all_table"
     target_df = spark.createDataFrame(
-        [{"id": 1, "value": "no_merge"}, {"id": 2, "value": "expected_merge"}, {"id": 5, "value": "xxxx"}]
+        [
+            {"id": 1, "value": "no_merge"},
+            {"id": 2, "value": "expected_merge"},
+            {"id": 5, "value": "xxxx"},
+        ]
     )
     source_df = spark.createDataFrame(
         [
@@ -73,7 +83,9 @@ def test_delta_table_merge_all(spark):
         # No merge as old value is greater
         5: "xxxx",
     }
-    DeltaTableWriter(table=table_name, output_mode=BatchOutputMode.APPEND, df=target_df).execute()
+    DeltaTableWriter(
+        table=table_name, output_mode=BatchOutputMode.APPEND, df=target_df
+    ).execute()
     merge_writer = DeltaTableWriter(
         table=table_name,
         output_mode=BatchOutputMode.MERGEALL,
@@ -89,7 +101,9 @@ def test_delta_table_merge_all(spark):
         with pytest.raises(SparkConnectDeltaTableException) as exc_info:
             merge_writer.execute()
 
-        assert str(exc_info.value).startswith("`DeltaTable.forName` is not supported due to delta calling _sc")
+        assert str(exc_info.value).startswith(
+            "`DeltaTable.forName` is not supported due to delta calling _sc"
+        )
     else:
         merge_writer.execute()
         result = {
@@ -107,12 +121,18 @@ def test_deltatablewriter_with_invalid_conditions(spark, dummy_df):
 
     if 3.4 < SPARK_MINOR_VERSION < 4.0 and is_remote_session():
         with pytest.raises(SparkConnectDeltaTableException) as exc_info:
-            builder = get_delta_table_for_name(spark_session=spark, table_name=table_name)
+            builder = get_delta_table_for_name(
+                spark_session=spark, table_name=table_name
+            )
 
-        assert str(exc_info.value).startswith("`DeltaTable.forName` is not supported due to delta calling _sc")
+        assert str(exc_info.value).startswith(
+            "`DeltaTable.forName` is not supported due to delta calling _sc"
+        )
     else:
         with pytest.raises(AnalysisException):
-            builder = get_delta_table_for_name(spark_session=spark, table_name=table_name)
+            builder = get_delta_table_for_name(
+                spark_session=spark, table_name=table_name
+            )
             merge_builder = builder.alias("target").merge(
                 condition="invalid_condition", source=dummy_df.alias("source")
             )
@@ -137,7 +157,10 @@ def test_delta_new_table_merge(spark):
         [
             {"id": 1, "value": "new_value"},
             {"id": 2, "value": "new_value"},
-            {"id": 3, "value": None},  # Will be not saved, because source.value IS NOT NULL,
+            {
+                "id": 3,
+                "value": None,
+            },  # Will be not saved, because source.value IS NOT NULL,
             {"id": 4, "value": "new_value"},
             {"id": 5, "value": "new_value"},
         ]
@@ -154,7 +177,8 @@ def test_delta_new_table_merge(spark):
         df=source_df,
     ).execute()
     result = {
-        list(row.asDict().values())[0]: list(row.asDict().values())[1] for row in spark.read.table(table_name).collect()
+        list(row.asDict().values())[0]: list(row.asDict().values())[1]
+        for row in spark.read.table(table_name).collect()
     }
     assert result == expected
 
@@ -270,7 +294,9 @@ def test_delta_with_options(spark):
     """
     sample_df = spark.createDataFrame([{"id": 1, "value": "test_value"}])
 
-    with patch("koheesio.spark.writers.delta.DeltaTableWriter.writer", new_callable=MagicMock) as mock_writer:
+    with patch(
+        "koheesio.spark.writers.delta.DeltaTableWriter.writer", new_callable=MagicMock
+    ) as mock_writer:
         delta_writer = DeltaTableWriter(
             table="test_table",
             output_mode=BatchOutputMode.APPEND,
@@ -279,9 +305,13 @@ def test_delta_with_options(spark):
             df=sample_df,
         )
         delta_writer.execute()
-        mock_writer.options.assert_called_once_with(testParam1="testValue1", testParam2="testValue2")
+        mock_writer.options.assert_called_once_with(
+            testParam1="testValue1", testParam2="testValue2"
+        )
 
-    with patch("koheesio.spark.writers.delta.DeltaTableWriter.writer", new_callable=MagicMock) as mock_writer:
+    with patch(
+        "koheesio.spark.writers.delta.DeltaTableWriter.writer", new_callable=MagicMock
+    ) as mock_writer:
         delta_writer = DeltaTableWriter(
             table="test_table",
             output_mode=BatchOutputMode.OVERWRITE,
@@ -290,7 +320,9 @@ def test_delta_with_options(spark):
             df=sample_df,
         )
         delta_writer.execute()
-        mock_writer.options.assert_called_once_with(testParam1="testValue1", testParam2="testValue2")
+        mock_writer.options.assert_called_once_with(
+            testParam1="testValue1", testParam2="testValue2"
+        )
 
 
 def test_merge_from_args(spark, dummy_df):
@@ -313,7 +345,11 @@ def test_merge_from_args(spark, dummy_df):
             output_mode=BatchOutputMode.MERGE,
             output_mode_params={
                 "merge_builder": [
-                    {"clause": "whenMatchedUpdate", "set": {"id": "source.id"}, "condition": "source.id=target.id"},
+                    {
+                        "clause": "whenMatchedUpdate",
+                        "set": {"id": "source.id"},
+                        "condition": "source.id=target.id",
+                    },
                     {
                         "clause": "whenNotMatchedInsert",
                         "values": {"id": "source.id"},
@@ -328,7 +364,9 @@ def test_merge_from_args(spark, dummy_df):
             with pytest.raises(SparkConnectDeltaTableException) as exc_info:
                 writer._merge_builder_from_args()
 
-            assert str(exc_info.value).startswith("`DeltaTable.forName` is not supported due to delta calling _sc")
+            assert str(exc_info.value).startswith(
+                "`DeltaTable.forName` is not supported due to delta calling _sc"
+            )
         else:
             writer._merge_builder_from_args()
             mock_delta_builder.whenMatchedUpdate.assert_called_once_with(
@@ -337,17 +375,21 @@ def test_merge_from_args(spark, dummy_df):
             mock_delta_builder.whenNotMatchedInsert.assert_called_once_with(
                 values={"id": "source.id"}, condition="source.id IS NOT NULL"
             )
-            assert ["clause" in c for c in writer.params["merge_builder"]] == [True] * len(
-                writer.params["merge_builder"]
-            )
+            assert ["clause" in c for c in writer.params["merge_builder"]] == [
+                True
+            ] * len(writer.params["merge_builder"])
 
 
 @pytest.mark.parametrize(
     "output_mode_params",
     [
         {
             "merge_builder": [
-                {"clause": "NOT-SUPPORTED-MERGE-CLAUSE", "set": {"id": "source.id"}, "condition": "source.id=target.id"}
+                {
+                    "clause": "NOT-SUPPORTED-MERGE-CLAUSE",
+                    "set": {"id": "source.id"},
+                    "condition": "source.id=target.id",
+                }
             ],
             "merge_cond": "source.id=target.id",
         },
@@ -368,7 +410,11 @@ def test_merge_no_table(spark):
 
     table_name = "test_merge_no_table"
     target_df = spark.createDataFrame(
-        [{"id": 1, "value": "no_merge"}, {"id": 2, "value": "expected_merge"}, {"id": 5, "value": "expected_merge"}]
+        [
+            {"id": 1, "value": "no_merge"},
+            {"id": 2, "value": "expected_merge"},
+            {"id": 5, "value": "expected_merge"},
+        ]
     )
     source_df = spark.createDataFrame(
         [
@@ -402,18 +448,26 @@ def test_merge_no_table(spark):
         "merge_cond": "source.id=target.id",
     }
     writer1 = DeltaTableWriter(
-        df=target_df, table=table_name, output_mode=BatchOutputMode.MERGE, output_mode_params=params
+        df=target_df,
+        table=table_name,
+        output_mode=BatchOutputMode.MERGE,
+        output_mode_params=params,
     )
     writer2 = DeltaTableWriter(
-        df=source_df, table=table_name, output_mode=BatchOutputMode.MERGE, output_mode_params=params
+        df=source_df,
+        table=table_name,
+        output_mode=BatchOutputMode.MERGE,
+        output_mode_params=params,
     )
     if 3.4 < SPARK_MINOR_VERSION < 4.0 and is_remote_session():
         writer1.execute()
 
         with pytest.raises(SparkConnectDeltaTableException) as exc_info:
             writer2.execute()
 
-        assert str(exc_info.value).startswith("`DeltaTable.forName` is not supported due to delta calling _sc")
+        assert str(exc_info.value).startswith(
+            "`DeltaTable.forName` is not supported due to delta calling _sc"
+        )
     else:
         writer1.execute()
         writer2.execute()
@@ -435,10 +489,14 @@ def test_log_clauses(mocker):
 
     mock_clause = mocker.MagicMock()
     mock_clause.clauseType.return_value = "test"
-    mock_clause.actions.return_value.toList.return_value.apply.return_value.toString.return_value = "test_column"
+    mock_clause.actions.return_value.toList.return_value.apply.return_value.toString.return_value = (
+        "test_column"
+    )
 
     mock_condition = mocker.MagicMock()
-    mock_condition.value.return_value.toString.return_value = "source_alias == target_alias"
+    mock_condition.value.return_value.toString.return_value = (
+        "source_alias == target_alias"
+    )
     mock_condition.toString.return_value = "None"
     mock_clause.condition.return_value = mock_condition
 
@@ -448,4 +506,62 @@ def test_log_clauses(mocker):
     result = log_clauses(mock_clauses, "source_alias", "target_alias")
 
     # Assert the result
-    assert result == "Test will perform action:Test columns (test_column) if `source_alias == target_alias`"
+    assert (
+        result
+        == "Test will perform action:Test columns (test_column) if `source_alias == target_alias`"
+    )
+
+
+def test_merge_builder_type__list_of_merge_builders(mocker, spark):
+    table_name = "test_merge_builder_type"
+    df = spark.createDataFrame([{"id": 1, "value": "test"}])
+    merge_builder = mocker.MagicMock(spec=list)  # mocks a list of merge builders
+    # No ValueError should be raised
+    DeltaTableWriter(
+        df=df,
+        table=table_name,
+        output_mode=BatchOutputMode.MERGE,
+        output_mode_params={"merge_builder": merge_builder},
+    )
+
+
+def test_merge_builder_type___delta_merge_builder(mocker, spark):
+    table_name = "test_merge_builder_type"
+    df = spark.createDataFrame([{"id": 1, "value": "test"}])
+    merge_builder = mocker.MagicMock(spec=DeltaMergeBuilder)
+    # No ValueError should be raised
+    DeltaTableWriter(
+        df=df,
+        table=table_name,
+        output_mode=BatchOutputMode.MERGE,
+        output_mode_params={"merge_builder": merge_builder},
+    )
+
+
+def test_merge_builder_type__connect_delta_merge_builder(mocker, spark):
+    table_name = "test_merge_builder_type"
+    df = spark.createDataFrame([{"id": 1, "value": "test"}])
+    # Not a delta.tables.DeltaMergeBuilder instance but the name is DeltaMergeBuilder (as in delta.connect.tables.DeltaMergeBuilder)
+    merge_builder = mocker.MagicMock()
+    merge_builder.__class__.__name__ = "DeltaMergeBuilder"
+    DeltaTableWriter(
+        df=df,
+        table=table_name,
+        output_mode=BatchOutputMode.MERGE,
+        output_mode_params={"merge_builder": merge_builder},
+    )
+
+
+def test_merge_builder_type__invalid_merge_builder(mocker, spark):
+    table_name = "test_merge_builder_type"
+    df = spark.createDataFrame([{"id": 1, "value": "test"}])
+    merge_builder = mocker.MagicMock(
+        spec=str
+    )  # Not a DeltaMergeBuilder instance nor a list
+    with pytest.raises(ValueError):
+        DeltaTableWriter(
+            df=df,
+            table=table_name,
+            output_mode=BatchOutputMode.MERGE,
+            output_mode_params={"merge_builder": merge_builder},
+        )