cli: only json from standard diff with --json

Author: sirupsen

data_diff/__main__.py

@@ -51,7 +51,7 @@
 @click.option("--max-age", default=None, help="Considers only rows younger than specified. See --min-age.")
 @click.option("-s", "--stats", is_flag=True, help="Print stats instead of a detailed diff")
 @click.option("-d", "--debug", is_flag=True, help="Print debug info")
-@click.option("--json", 'json_output', is_flag=True, help="Print JSON output for --stats")
+@click.option("--json", 'json_output', is_flag=True, help="Print JSONL output for machine readability")
 @click.option("-v", "--verbose", is_flag=True, help="Print extra info")
 @click.option("-i", "--interactive", is_flag=True, help="Confirm queries, implies --debug")
 @click.option("--keep-column-case", is_flag=True, help="Don't use the schema to fix the case of given column names.")
@@ -160,16 +160,21 @@ def main(
                 "different_-": minus,
                 "total": max_table_count,
             }
-            print(json.dumps(json_output, indent=2))
+            print(json.dumps(json_output))
         else:
             print(f"Diff-Total: {len(diff)} changed rows out of {max_table_count}")
             print(f"Diff-Percent: {percent:.14f}%")
             print(f"Diff-Split: +{plus}  -{minus}")
     else:
         for op, key in diff_iter:
             color = COLOR_SCHEME[op]
-            jsonl = json.dumps([op, list(key)])
-            rich.print(f"[{color}]{jsonl}[/{color}]")
+
+            if json_output:
+                jsonl = json.dumps([op, list(key)])
+                rich.print(f"[{color}]{jsonl}[/{color}]")
+            else:
+                rich.print(f"[{color}]{op} {key!r}[/{color}]")
+
             sys.stdout.flush()
 
     end = time.time()

data_diff/diff_tables.py

@@ -382,9 +382,10 @@ def _bisect_and_diff_tables(self, table1, table2, level=0, max_rows=None):
             rows1, rows2 = self._threaded_call("get_values", [table1, table2])
             diff = list(diff_sets(rows1, rows2))
 
-            # This happens when the initial bisection threshold is larger than
-            # the table itself.
-            if level == 0 and not self.stats.get("table1_count", False):
+            # Initial bisection_threshold larger than count. Normally we always
+            # checksum and count segments, even if we get the values. At the
+            # first level, however, that won't be true.
+            if level == 0:
                 self.stats["table1_count"] = self.stats.get("table1_count", 0) + len(rows1)
                 self.stats["table2_count"] = self.stats.get("table2_count", 0) + len(rows2)
 
@@ -427,8 +428,8 @@ def _diff_tables(self, table1, table2, level=0, segment_index=None, segment_coun
             return
 
         if level == 1:
-            self.stats["table1_count"] = self.stats.get("table_count1", 0) + count1
-            self.stats["table2_count"] = self.stats.get("table_count2", 0) + count2
+            self.stats["table1_count"] = self.stats.get("table1_count", 0) + count1
+            self.stats["table2_count"] = self.stats.get("table2_count", 0) + count2
 
         if checksum1 != checksum2:
             yield from self._bisect_and_diff_tables(table1, table2, level=level, max_rows=max(count1, count2))

tests/test_diff_tables.py

@@ -155,6 +155,8 @@ def test_diff_small_tables(self):
         diff = list(self.differ.diff_tables(self.table, self.table2))
         expected = [("-", ("2", time + ".000000"))]
         self.assertEqual(expected, diff)
+        self.assertEqual(2, self.differ.stats["table1_count"])
+        self.assertEqual(1, self.differ.stats["table2_count"])
 
     def test_diff_table_above_bisection_threshold(self):
         time = "2022-01-01 00:00:00"
@@ -176,6 +178,8 @@ def test_diff_table_above_bisection_threshold(self):
         diff = list(self.differ.diff_tables(self.table, self.table2))
         expected = [("-", ("5", time + ".000000"))]
         self.assertEqual(expected, diff)
+        self.assertEqual(5, self.differ.stats["table1_count"])
+        self.assertEqual(4, self.differ.stats["table2_count"])
 
     def test_return_empty_array_when_same(self):
         time = "2022-01-01 00:00:00"