support combo pks in --dbt local_diff

Author: dlawin

data_diff/dbt.py

@@ -73,26 +73,16 @@ def dbt_diff(
 
         if is_cloud and len(diff_vars.primary_keys) > 0:
             _cloud_diff(diff_vars)
-        elif is_cloud:
-            rich.print(
-                "[red]"
-                + ".".join(diff_vars.prod_path)
-                + " <> "
-                + ".".join(diff_vars.dev_path)
-                + "[/] \n"
-                + "Skipped due to missing primary-key tag\n"
-            )
-
-        if not is_cloud and len(diff_vars.primary_keys) == 1:
+        elif not is_cloud and len(diff_vars.primary_keys) > 0:
             _local_diff(diff_vars)
-        elif not is_cloud:
+        else:
             rich.print(
                 "[red]"
                 + ".".join(diff_vars.prod_path)
                 + " <> "
                 + ".".join(diff_vars.dev_path)
                 + "[/] \n"
-                + "Skipped due to missing primary-key tag or multi-column primary-key (unsupported for non --cloud diffs)\n"
+                + "Skipped due to missing primary-key tag\n"
             )
 
         rich.print("Diffs Complete!")
@@ -127,10 +117,9 @@ def _local_diff(diff_vars: DiffVars) -> None:
     column_diffs_str = ""
     dev_qualified_string = ".".join(diff_vars.dev_path)
     prod_qualified_string = ".".join(diff_vars.prod_path)
-    primary_key = diff_vars.primary_keys[0]
 
-    table1 = connect_to_table(diff_vars.connection, dev_qualified_string, primary_key)
-    table2 = connect_to_table(diff_vars.connection, prod_qualified_string, primary_key)
+    table1 = connect_to_table(diff_vars.connection, dev_qualified_string, tuple(diff_vars.primary_keys))
+    table2 = connect_to_table(diff_vars.connection, prod_qualified_string, tuple(diff_vars.primary_keys))
 
     table1_columns = list(table1.get_schema())
     try:
@@ -159,7 +148,7 @@ def _local_diff(diff_vars: DiffVars) -> None:
     if table2_set_diff:
         column_diffs_str += "Column(s) removed: " + str(table2_set_diff) + "\n"
 
-    mutual_set.discard(primary_key)
+    mutual_set = mutual_set - set(diff_vars.primary_keys)
     extra_columns = tuple(mutual_set)
 
     diff = diff_tables(table1, table2, threaded=True, algorithm=Algorithm.JOINDIFF, extra_columns=extra_columns)

tests/test_dbt.py

@@ -351,17 +351,17 @@ def test_local_diff(self, mock_diff_tables):
         mock_diff.__iter__.return_value = [1, 2, 3]
         dev_qualified_list = ["dev_db", "dev_schema", "dev_table"]
         prod_qualified_list = ["prod_db", "prod_schema", "prod_table"]
-        expected_key = "key"
-        diff_vars = DiffVars(dev_qualified_list, prod_qualified_list, [expected_key], None, mock_connection)
+        expected_keys = ["key"]
+        diff_vars = DiffVars(dev_qualified_list, prod_qualified_list, expected_keys, None, mock_connection)
         with patch("data_diff.dbt.connect_to_table", side_effect=[mock_table1, mock_table2]) as mock_connect:
             _local_diff(diff_vars)
 
         mock_diff_tables.assert_called_once_with(
             mock_table1, mock_table2, threaded=True, algorithm=Algorithm.JOINDIFF, extra_columns=tuple(column_set)
         )
         self.assertEqual(mock_connect.call_count, 2)
-        mock_connect.assert_any_call(mock_connection, ".".join(dev_qualified_list), expected_key)
-        mock_connect.assert_any_call(mock_connection, ".".join(prod_qualified_list), expected_key)
+        mock_connect.assert_any_call(mock_connection, ".".join(dev_qualified_list), tuple(expected_keys))
+        mock_connect.assert_any_call(mock_connection, ".".join(prod_qualified_list), tuple(expected_keys))
         mock_diff.get_stats_string.assert_called_once()
 
     @patch("data_diff.dbt.diff_tables")
@@ -377,17 +377,17 @@ def test_local_diff_no_diffs(self, mock_diff_tables):
         mock_diff.__iter__.return_value = []
         dev_qualified_list = ["dev_db", "dev_schema", "dev_table"]
         prod_qualified_list = ["prod_db", "prod_schema", "prod_table"]
-        expected_key = "primary_key_column"
-        diff_vars = DiffVars(dev_qualified_list, prod_qualified_list, [expected_key], None, mock_connection)
+        expected_keys = ["primary_key_column"]
+        diff_vars = DiffVars(dev_qualified_list, prod_qualified_list, expected_keys, None, mock_connection)
         with patch("data_diff.dbt.connect_to_table", side_effect=[mock_table1, mock_table2]) as mock_connect:
             _local_diff(diff_vars)
 
         mock_diff_tables.assert_called_once_with(
             mock_table1, mock_table2, threaded=True, algorithm=Algorithm.JOINDIFF, extra_columns=tuple(column_set)
         )
         self.assertEqual(mock_connect.call_count, 2)
-        mock_connect.assert_any_call(mock_connection, ".".join(dev_qualified_list), expected_key)
-        mock_connect.assert_any_call(mock_connection, ".".join(prod_qualified_list), expected_key)
+        mock_connect.assert_any_call(mock_connection, ".".join(dev_qualified_list), tuple(expected_keys))
+        mock_connect.assert_any_call(mock_connection, ".".join(prod_qualified_list), tuple(expected_keys))
         mock_diff.get_stats_string.assert_not_called()
 
     @patch("data_diff.dbt.rich.print")