Now automatically fixing the column case using the schema.

Added the --keep-column-case switch to disable it.

Author: erezsh

data_diff/__main__.py

@@ -53,6 +53,7 @@
 @click.option("-d", "--debug", is_flag=True, help="Print debug info")
 @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.")
 @click.option(
     "-j",
     "--threads",
@@ -79,6 +80,7 @@ def main(
     verbose,
     interactive,
     threads,
+    keep_column_case,
 ):
     if limit and stats:
         print("Error: cannot specify a limit when using the -s/--stats switch")
@@ -119,6 +121,7 @@ def main(
         options = dict(
             min_update=max_age and parse_time_before_now(max_age),
             max_update=min_age and parse_time_before_now(min_age),
+            case_sensitive=keep_column_case,
         )
     except ParseError as e:
         logging.error("Error while parsing age expression: %s" % e)

data_diff/diff_tables.py

@@ -1,10 +1,11 @@
 """Provides classes for performing a table diff
 """
 
+from abc import ABC, abstractmethod
 import time
 from operator import attrgetter, methodcaller
 from collections import defaultdict
-from typing import List, Tuple, Iterator, Optional, Mapping
+from typing import List, Tuple, Iterator, Optional
 import logging
 from concurrent.futures import ThreadPoolExecutor
 
@@ -36,24 +37,44 @@ def parse_table_name(t):
     return tuple(t.split("."))
 
 
-class CaseInsensitiveDict(Mapping):
-    def __init__(self, initial=()):
-        self._dict = {k.lower(): v for k, v in dict(initial).items()}
+class Schema(ABC):
+    @abstractmethod
+    def get_key(self, key: str) -> str:
+        ...
 
-    def __setitem__(self, key, value):
-        self._dict[key.lower()] = value
+    @abstractmethod
+    def __getitem__(self, key: str) -> str:
+        ...
 
-    def __getitem__(self, key):
-        try:
-            return self._dict[key.lower()]
-        except KeyError:
-            raise
+    @abstractmethod
+    def __setitem__(self, key: str, value):
+        ...
 
-    def __iter__(self):
-        return iter(self._dict)
+    @abstractmethod
+    def __contains__(self, key: str) -> bool:
+        ...
 
-    def __len__(self):
-        return len(self._dict)
+
+class Schema_CaseSensitive(dict, Schema):
+    def get_key(self, key):
+        return key
+
+
+class Schema_CaseInsensitive(Schema):
+    def __init__(self, initial):
+        self._dict = {k.lower(): (k, v) for k, v in dict(initial).items()}
+
+    def get_key(self, key: str) -> str:
+        return self._dict[key.lower()][0]
+
+    def __getitem__(self, key: str) -> str:
+        return self._dict[key.lower()][1]
+
+    def __setitem__(self, key: str, value):
+        self._dict[key.lower()] = key, value
+
+    def __contains__(self, key):
+        return key.lower() in self._dict
 
 
 @dataclass(frozen=False)
@@ -88,8 +109,8 @@ class TableSegment:
     min_update: DbTime = None
     max_update: DbTime = None
 
-    quote_columns: bool = True
-    _schema: Mapping[str, ColType] = None
+    case_sensitive: bool = True
+    _schema: Schema = None
 
     def __post_init__(self):
         if not self.update_column and (self.min_update or self.max_update):
@@ -110,17 +131,24 @@ def _update_column(self):
         return self._quote_column(self.update_column)
 
     def _quote_column(self, c):
-        if self.quote_columns:
-            return self.database.quote(c)
-        return c
+        if self._schema:
+            c = self._schema.get_key(c)
+        return self.database.quote(c)
 
     def with_schema(self) -> "TableSegment":
         "Queries the table schema from the database, and returns a new instance of TableSegmentWithSchema."
         if self._schema:
             return self
         schema = self.database.query_table_schema(self.table_path)
-        if not self.quote_columns:
-            schema = CaseInsensitiveDict(schema)
+        if self.case_sensitive:
+            schema = Schema_CaseSensitive(schema)
+        else:
+            if len({k.lower() for k in schema}) < len(schema):
+                logger.warn(
+                    f'Ambiguous schema for {self.database}:{".".join(self.table_path)} | Columns = {", ".join(list(schema))}'
+                )
+                logger.warn("We recommend to disable case-insensitivity (remove --any-case).")
+            schema = Schema_CaseInsensitive(schema)
         return self.new(_schema=schema)
 
     def _make_key_range(self):

tests/test_database.py

@@ -18,6 +18,7 @@ def test_md5_to_int(self):
 
         self.assertEqual(str_to_checksum(str), self.mysql.query(query, int))
 
+
 class TestConnect(unittest.TestCase):
     def test_bad_uris(self):
         self.assertRaises(ValueError, connect_to_uri, "p")

tests/test_database_types.py

@@ -214,8 +214,8 @@ def test_types(self, source_db, target_db, source_type, target_type, type_catego
         dst_conn.query(f"CREATE TABLE {dst_table}(id int, col {target_type});", None)
         _insert_to_table(dst_conn, dst_table, values_in_source)
 
-        self.table = TableSegment(self.src_conn, src_table_path, "id", None, ("col",), quote_columns=False)
-        self.table2 = TableSegment(self.dst_conn, dst_table_path, "id", None, ("col",), quote_columns=False)
+        self.table = TableSegment(self.src_conn, src_table_path, "id", None, ("col",), case_sensitive=False)
+        self.table2 = TableSegment(self.dst_conn, dst_table_path, "id", None, ("col",), case_sensitive=False)
 
         self.assertEqual(len(sample_values), self.table.count())
         self.assertEqual(len(sample_values), self.table2.count())