diff --git a/data_diff/format.py b/data_diff/format.py index af2c0009..5a16f274 100644 --- a/data_diff/format.py +++ b/data_diff/format.py @@ -213,16 +213,16 @@ def _jsonify_diff(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, Json continue if field.startswith("is_diff_"): - column_name = field.replace("is_diff_", "") + column_name = field[len("is_diff_") :] columns[column_name]["isDiff"] = bool(value) elif field.endswith("_a"): - column_name = field.replace("_a", "") + column_name = field[: -len("_a")] columns[column_name]["dataset1"] = value columns[column_name]["isPK"] = column_name in key_columns elif field.endswith("_b"): - column_name = field.replace("_b", "") + column_name = field[: -len("_b")] columns[column_name]["dataset2"] = value columns[column_name]["isPK"] = column_name in key_columns @@ -237,11 +237,11 @@ def _jsonify_exclusive(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, if field.startswith("is_diff_"): continue if field.endswith("_b") and row["is_exclusive_b"]: - column_name = field.replace("_b", "") + column_name = field[: -len("_b")] columns[column_name]["isPK"] = column_name in key_columns columns[column_name]["value"] = value elif field.endswith("_a") and row["is_exclusive_a"]: - column_name = field.replace("_a", "") + column_name = field[: -len("_a")] columns[column_name]["isPK"] = column_name in key_columns columns[column_name]["value"] = value return {column: JsonExclusiveRowValue(**data) for column, data in columns.items()} diff --git a/tests/test_format.py b/tests/test_format.py index 2ba0f999..393a4e95 100644 --- a/tests/test_format.py +++ b/tests/test_format.py @@ -104,3 +104,60 @@ def test_jsonify_diff_no_difeference(self): "columns": None, }, ) + + def test_jsonify_column_suffix_fix(self): + diff = DiffResultWrapper( + info_tree=InfoTree( + info=SegmentInfo( + tables=[ + TableSegment(table_path=("db", "schema", "table1"), key_columns=("id_a",), database=Database()), + TableSegment(table_path=("db", "schema", "table2"), key_columns=("id_a",), database=Database()), + ], + diff_schema=( + ("is_exclusive_a", bool), + ("is_exclusive_b", bool), + ("is_diff_id_a", int), + ("is_diff_value_b", int), + ("id_a_a", str), + ("id_a_b", str), + ("value_b_a", str), + ("value_b_b", str), + ), + diff=[ + (False, False, 0, 1, "1", "1", "3", "201"), + (True, False, 1, 1, "2", None, "4", None), + (False, True, 1, 1, None, "3", None, "202"), + ], + ) + ), + diff=[], + stats={}, + ) + json_diff = jsonify(diff, dbt_model="my_model") + self.assertEqual( + json_diff, + { + "version": "1.0.0", + "status": "success", + "result": "different", + "model": "my_model", + "dataset1": ["db", "schema", "table1"], + "dataset2": ["db", "schema", "table2"], + "rows": { + "exclusive": { + "dataset1": [{"id_a": {"isPK": True, "value": "2"}, "value_b": {"isPK": False, "value": "4"}}], + "dataset2": [ + {"id_a": {"isPK": True, "value": "3"}, "value_b": {"isPK": False, "value": "202"}} + ], + }, + "diff": [ + { + "id_a": {"isPK": True, "dataset1": "1", "dataset2": "1", "isDiff": False}, + "value_b": {"isPK": False, "dataset1": "3", "dataset2": "201", "isDiff": True}, + }, + ], + }, + "summary": None, + "columns": None, + }, + )