fix crashing when parsing column names including _a, _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()}

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,
+            },
+        )