squash: add get_stats method

Author: dlawin

.gitignore

@@ -141,3 +141,6 @@ benchmark_*.png
 
 # IntelliJ
 .idea
+
+# VSCode
+.vscode

data_diff/__main__.py

@@ -403,58 +403,19 @@ def _main(
     ]
 
     diff_iter = differ.diff_tables(*segments)
-    info = diff_iter.info_tree.info
 
     if limit:
         diff_iter = islice(diff_iter, int(limit))
 
     if stats:
-        diff = list(diff_iter)
-        key_columns_len = len(key_columns)
-
-        diff_by_key = {}
-        for sign, values in diff:
-            k = values[:key_columns_len]
-            if k in diff_by_key:
-                assert sign != diff_by_key[k]
-                diff_by_key[k] = "!"
-            else:
-                diff_by_key[k] = sign
-
-        diff_by_sign = {k: 0 for k in "+-!"}
-        for sign in diff_by_key.values():
-            diff_by_sign[sign] += 1
-
-        table1_count = info.rowcounts[1]
-        table2_count = info.rowcounts[2]
-        unchanged = table1_count - diff_by_sign["-"] - diff_by_sign["!"]
-        diff_percent = 1 - unchanged / max(table1_count, table2_count)
-
+        # required to create this variable before get_stats
+        diff_list = list(diff_iter)
+        stats_output = diff_iter.get_stats()
         if json_output:
-            json_output = {
-                "rows_A": table1_count,
-                "rows_B": table2_count,
-                "exclusive_A": diff_by_sign["-"],
-                "exclusive_B": diff_by_sign["+"],
-                "updated": diff_by_sign["!"],
-                "unchanged": unchanged,
-                "total": sum(diff_by_sign.values()),
-                "stats": differ.stats,
-            }
-            rich.print_json(json.dumps(json_output))
+            rich.print(json.dumps(stats_output[0]))
         else:
-            rich.print(f"{table1_count} rows in table A")
-            rich.print(f"{table2_count} rows in table B")
-            rich.print(f"{diff_by_sign['-']} rows exclusive to table A (not present in B)")
-            rich.print(f"{diff_by_sign['+']} rows exclusive to table B (not present in A)")
-            rich.print(f"{diff_by_sign['!']} rows updated")
-            rich.print(f"{unchanged} rows unchanged")
-            rich.print(f"{100*diff_percent:.2f}% difference score")
-
-            if differ.stats:
-                print("\nExtra-Info:")
-                for k, v in sorted(differ.stats.items()):
-                    rich.print(f"  {k} = {v}")
+            rich.print(stats_output[1])
+
     else:
         for op, values in diff_iter:
             color = COLOR_SCHEME[op]

data_diff/diff_tables.py

@@ -11,6 +11,7 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from runtype import dataclass
+from dataclasses import field
 
 from data_diff.info_tree import InfoTree, SegmentInfo
 
@@ -82,10 +83,66 @@ def _run_in_background(self, *funcs):
 class DiffResultWrapper:
     diff: iter  # DiffResult
     info_tree: InfoTree
+    stats: dict
+    result_list: list = field(default_factory=list)
 
     def __iter__(self):
-        return iter(self.diff)
+        for i in self.diff:
+            self.result_list.append(i)
+            yield i
+
+    def get_stats(self):
+
+        diff_by_key = {}
+        if len(self.result_list) > 0:
+            for sign, values in self.result_list:
+                k = values[: len(self.info_tree.info.tables[0].key_columns)]
+                if k in diff_by_key:
+                    assert sign != diff_by_key[k]
+                    diff_by_key[k] = "!"
+                else:
+                    diff_by_key[k] = sign
+
+            diff_by_sign = {k: 0 for k in "+-!"}
+            for sign in diff_by_key.values():
+                diff_by_sign[sign] += 1
+
+            table1_count = self.info_tree.info.rowcounts[1]
+            table2_count = self.info_tree.info.rowcounts[2]
+            unchanged = table1_count - diff_by_sign["-"] - diff_by_sign["!"]
+            diff_percent = 1 - unchanged / max(table1_count, table2_count)
+
+
+            json_output = {
+                    "rows_A": table1_count,
+                    "rows_B": table2_count,
+                    "exclusive_A": diff_by_sign["-"],
+                    "exclusive_B": diff_by_sign["+"],
+                    "updated": diff_by_sign["!"],
+                    "unchanged": unchanged,
+                    "total": sum(diff_by_sign.values()),
+                    "stats": self.stats,
+                }
+
+            string_output = ""
+            string_output += f"{table1_count} rows in table A\n"
+            string_output += f"{table2_count} rows in table B\n"
+            string_output += f"{diff_by_sign['-']} rows exclusive to table A (not present in B)\n"
+            string_output += f"{diff_by_sign['+']} rows exclusive to table B (not present in A)\n"
+            string_output += f"{diff_by_sign['!']} rows updated\n"
+            string_output += f"{unchanged} rows unchanged\n"
+            string_output += f"{100*diff_percent:.2f}% difference score\n"
+
+            if self.stats:
+                string_output += "\nExtra-Info:\n"
+                for k, v in sorted(self.stats.items()):
+                    string_output += f"  {k} = {v}\n"
+        else:
+            raise RuntimeError(
+                "result_list is empty, consume the diff iterator to populate values: e.g. \ndiff_iter = diff_tables(...) \ndiff_list = list(diff_iter) \ndiff_iter.print_stats(json_output)"
+            )
 
+        return json_output, string_output
 
 class TableDiffer(ThreadBase, ABC):
     bisection_factor = 32
@@ -106,7 +163,7 @@ def diff_tables(self, table1: TableSegment, table2: TableSegment, info_tree: Inf
         """
         if info_tree is None:
             info_tree = InfoTree(SegmentInfo([table1, table2]))
-        return DiffResultWrapper(self._diff_tables_wrapper(table1, table2, info_tree), info_tree)
+        return DiffResultWrapper(self._diff_tables_wrapper(table1, table2, info_tree), info_tree, self.stats, [])
 
     def _diff_tables_wrapper(self, table1: TableSegment, table2: TableSegment, info_tree: InfoTree) -> DiffResult:
         if is_tracking_enabled():
@@ -177,6 +234,8 @@ def _bisect_and_diff_tables(self, table1, table2, info_tree):
             raise NotImplementedError("Composite key not supported yet!")
         if len(table2.key_columns) > 1:
             raise NotImplementedError("Composite key not supported yet!")
+        if len(table1.key_columns) != len(table2.key_columns):
+            raise ValueError("Tables should have an equivalent number of key columns!")
         (key1,) = table1.key_columns
         (key2,) = table2.key_columns
 

tests/test_api.py

@@ -1,12 +1,14 @@
 import unittest
+import io
+import unittest.mock
 import arrow
 from datetime import datetime
 
 from data_diff import diff_tables, connect_to_table
 from data_diff.databases import MySQL
 from data_diff.sqeleton.queries import table, commit
 
-from .common import TEST_MYSQL_CONN_STRING, get_conn
+from .common import TEST_MYSQL_CONN_STRING, get_conn, random_table_suffix
 
 
 def _commit(conn):
@@ -16,16 +18,17 @@ def _commit(conn):
 class TestApi(unittest.TestCase):
     def setUp(self) -> None:
         self.conn = get_conn(MySQL)
-        table_src_name = "test_api"
-        table_dst_name = "test_api_2"
+        suffix = random_table_suffix()
+        self.table_src_name = f"test_api{suffix}"
+        self.table_dst_name = f"test_api_2{suffix}"
 
-        self.table_src = table(table_src_name)
-        self.table_dst = table(table_dst_name)
+        self.table_src = table(self.table_src_name)
+        self.table_dst = table(self.table_dst_name)
 
         self.conn.query(self.table_src.drop(True))
         self.conn.query(self.table_dst.drop(True))
 
-        src_table = table(table_src_name, schema={"id": int, "datetime": datetime, "text_comment": str})
+        src_table = table(self.table_src_name, schema={"id": int, "datetime": datetime, "text_comment": str})
         self.conn.query(src_table.create())
         self.now = now = arrow.get()
 
@@ -53,8 +56,8 @@ def tearDown(self) -> None:
         return super().tearDown()
 
     def test_api(self):
-        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, "test_api")
-        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, ("test_api_2",))
+        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_src_name)
+        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, (self.table_dst_name,))
         diff = list(diff_tables(t1, t2))
         assert len(diff) == 1
 
@@ -65,10 +68,38 @@ def test_api(self):
         diff_id = diff[0][1][0]
         where = f"id != {diff_id}"
 
-        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, "test_api", where=where)
-        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, "test_api_2", where=where)
+        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_src_name, where=where)
+        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_dst_name, where=where)
         diff = list(diff_tables(t1, t2))
         assert len(diff) == 0
 
         t1.database.close()
         t2.database.close()
+
+    def test_api_get_stats(self):
+        expected_string = "5 rows in table A\n4 rows in table B\n1 rows exclusive to table A (not present in B)\n0 rows exclusive to table B (not present in A)\n0 rows updated\n4 rows unchanged\n20.00% difference score\n\nExtra-Info:\n  rows_downloaded = 5\n"
+        expected_dict = {'rows_A': 5, 'rows_B': 4, 'exclusive_A': 1, 'exclusive_B': 0, 'updated': 0, 'unchanged': 4, 'total': 1, 'stats': {'rows_downloaded': 5}}
+        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_src_name)
+        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_dst_name)
+        diff = diff_tables(t1, t2)
+        diff_list = list(diff)
+        output = diff.get_stats()
+
+        self.assertEqual(expected_dict, output[0])
+        self.assertEqual(expected_string, output[1])
+        self.assertIsNotNone(diff)
+        assert len(diff_list) == 1
+
+        t1.database.close()
+        t2.database.close()
+
+    def test_api_print_error(self):
+        t1 = connect_to_table(TEST_MYSQL_CONN_STRING, self.table_src_name)
+        t2 = connect_to_table(TEST_MYSQL_CONN_STRING, (self.table_dst_name,))
+        diff = diff_tables(t1, t2)
+
+        with self.assertRaises(RuntimeError):
+            diff.get_stats()
+
+        t1.database.close()
+        t2.database.close()

tests/test_cli.py

@@ -8,7 +8,7 @@
 from data_diff.databases import MySQL
 from data_diff.sqeleton.queries import table, commit
 
-from .common import TEST_MYSQL_CONN_STRING, get_conn
+from .common import TEST_MYSQL_CONN_STRING, get_conn, random_table_suffix
 
 
 def _commit(conn):
@@ -30,15 +30,16 @@ class TestCLI(unittest.TestCase):
     def setUp(self) -> None:
         self.conn = get_conn(MySQL)
 
-        table_src_name = "test_cli"
-        table_dst_name = "test_cli_2"
+        suffix = random_table_suffix()
+        self.table_src_name = f"test_api{suffix}"
+        self.table_dst_name = f"test_api_2{suffix}"
 
-        self.table_src = table(table_src_name)
-        self.table_dst = table(table_dst_name)
+        self.table_src = table(self.table_src_name)
+        self.table_dst = table(self.table_dst_name)
         self.conn.query(self.table_src.drop(True))
         self.conn.query(self.table_dst.drop(True))
 
-        src_table = table(table_src_name, schema={"id": int, "datetime": datetime, "text_comment": str})
+        src_table = table(self.table_src_name, schema={"id": int, "datetime": datetime, "text_comment": str})
         self.conn.query(src_table.create())
         self.conn.query("SET @@session.time_zone='+00:00'")
         now = self.conn.query("select now()", datetime)
@@ -67,15 +68,15 @@ def tearDown(self) -> None:
         return super().tearDown()
 
     def test_basic(self):
-        diff = run_datadiff_cli(TEST_MYSQL_CONN_STRING, "test_cli", TEST_MYSQL_CONN_STRING, "test_cli_2")
+        diff = run_datadiff_cli(TEST_MYSQL_CONN_STRING, self.table_src_name, TEST_MYSQL_CONN_STRING, self.table_dst_name)
         assert len(diff) == 1
 
     def test_options(self):
         diff = run_datadiff_cli(
             TEST_MYSQL_CONN_STRING,
-            "test_cli",
+            self.table_src_name,
             TEST_MYSQL_CONN_STRING,
-            "test_cli_2",
+            self.table_dst_name,
             "--bisection-factor",
             "16",
             "--bisection-threshold",
@@ -88,3 +89,15 @@ def test_options(self):
             "1h",
         )
         assert len(diff) == 1
+
+    def test_stats(self):
+        diff_output = run_datadiff_cli(
+            TEST_MYSQL_CONN_STRING, self.table_src_name, TEST_MYSQL_CONN_STRING, self.table_dst_name, "-s"
+        )
+        assert len(diff_output) == 11
+    
+    def test_stats_json(self):
+        diff_output = run_datadiff_cli(
+            TEST_MYSQL_CONN_STRING, self.table_src_name, TEST_MYSQL_CONN_STRING, self.table_dst_name, "-s", "--json"
+        )
+        assert len(diff_output) == 2