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())