Initial support for Alphanumeric IDs + tests (Issue #59)

Author: erezsh

data_diff/databases/base.py

@@ -16,6 +16,7 @@
  Float,
  ColType_UUID,
  Native_UUID,
+ String_Alphanum,
  String_UUID,
  TemporalType,
  UnknownColType,
@@ -221,6 +222,23 @@ def _refine_coltypes(self, table_path: DbPath, col_dict: Dict[str, ColType]):
  assert col_name in col_dict
  col_dict[col_name] = String_UUID()
 
+ alphanum_samples = [s for s in samples if s and String_Alphanum.test_value(s)]
+ if alphanum_samples:
+ if len(alphanum_samples) != len(samples):
+ logger.warning(
+ f"Mixed Alphanum/Non-Alphanum values detected in column {'.'.join(table_path)}.{col_name}, disabling Alphanum support."
+ )
+ else:
+ assert col_name in col_dict
+ lens = set(map(len, alphanum_samples))
+ if len(lens) > 1:
+ logger.warning(
+ f"Mixed Alphanum lengths detected in column {'.'.join(table_path)}.{col_name}, disabling Alphanum support."
+ )
+ else:
+ (length,) = lens
+ col_dict[col_name] = String_Alphanum(length=length)
+
  # @lru_cache()
  # def get_table_schema(self, path: DbPath) -> Dict[str, ColType]:
  # return self.query_table_schema(path)

data_diff/databases/database_types.py

@@ -5,11 +5,11 @@
 
 from runtype import dataclass
 
-from data_diff.utils import ArithUUID
+from data_diff.utils import ArithAlphanumeric, ArithUUID, ArithString
 
 
 DbPath = Tuple[str, ...]
-DbKey = Union[int, str, bytes, ArithUUID]
+DbKey = Union[int, str, bytes, ArithUUID, ArithAlphanumeric]
 DbTime = datetime
 
 
@@ -18,11 +18,6 @@ class ColType:
  pass
 
 
-class IKey(ABC):
- "Interface for ColType, for using a column as a key in data-diff"
- python_type: type
-
-
 @dataclass
 class PrecisionType(ColType):
  precision: int
@@ -63,6 +58,9 @@ class IKey(ABC):
  "Interface for ColType, for using a column as a key in data-diff"
  python_type: type
 
+ def make_value(self, value):
+ return self.python_type(value)
+
 
 class Decimal(FractionalType, IKey): # Snowflake may use Decimal as a key
  @property
@@ -80,6 +78,10 @@ class ColType_UUID(ColType, IKey):
  python_type = ArithUUID
 
 
+class ColType_Alphanum(ColType, IKey):
+ python_type = ArithAlphanumeric
+
+
 class Native_UUID(ColType_UUID):
  pass
 
@@ -88,6 +90,24 @@ class String_UUID(StringType, ColType_UUID):
  pass
 
 
+@dataclass
+class String_Alphanum(StringType, ColType_Alphanum):
+ length: int
+
+ @staticmethod
+ def test_value(value: str) -> bool:
+ try:
+ ArithAlphanumeric(value)
+ return True
+ except ValueError:
+ return False
+
+ def make_value(self, value):
+ if len(value) != self.length:
+ raise ValueError(f"Expected alphanumeric value of length {self.length}, but got '{value}'.")
+ return self.python_type(value, max_len=self.length)
+
+
 @dataclass
 class Text(StringType):
  supported = False

data_diff/diff_tables.py

@@ -16,7 +16,7 @@
 from .utils import safezip, split_space
 from .databases.base import Database
 from .databases.database_types import (
- ArithUUID,
+ ArithString,
  IKey,
  Native_UUID,
  NumericType,
@@ -175,10 +175,10 @@ def get_values(self) -> list:
  def choose_checkpoints(self, count: int) -> List[DbKey]:
  "Suggests a bunch of evenly-spaced checkpoints to split by (not including start, end)"
  assert self.is_bounded
- if isinstance(self.min_key, ArithUUID):
+ if isinstance(self.min_key, ArithString):
+ assert type(self.min_key) is type(self.max_key)
  checkpoints = split_space(self.min_key.int, self.max_key.int, count)
- assert isinstance(self.max_key, ArithUUID)
- return [ArithUUID(int=i) for i in checkpoints]
+ return [self.min_key.new(int=i) for i in checkpoints]
 
  return split_space(self.min_key, self.max_key, count)
 
@@ -363,7 +363,7 @@ def diff_tables(self, table1: TableSegment, table2: TableSegment) -> DiffResult:
 
  def _parse_key_range_result(self, key_type, key_range):
  mn, mx = key_range
- cls = key_type.python_type
+ cls = key_type.make_value
  # We add 1 because our ranges are exclusive of the end (like in Python)
  try:
  return cls(mn), cls(mx) + 1

data_diff/sql.py

@@ -8,7 +8,7 @@
 
 from .utils import join_iter
 
-from .databases.database_types import AbstractDatabase, DbPath, DbKey, DbTime, ArithUUID
+from .databases.database_types import AbstractDatabase, DbPath, DbKey, DbTime, ArithString
 
 
 class Sql:
@@ -69,7 +69,7 @@ def compile(self, c: Compiler):
  return "b'%s'" % self.value.decode()
  elif isinstance(self.value, str):
  return "'%s'" % self.value
- elif isinstance(self.value, ArithUUID):
+ elif isinstance(self.value, ArithString):
  return "'%s'" % self.value
  return str(self.value)
 

data_diff/utils.py

@@ -3,6 +3,9 @@
 
 from typing import Union, Any
 from uuid import UUID
+import string
+
+alphanums = string.digits + string.ascii_lowercase
 
 
 def safezip(*args):
@@ -16,24 +19,99 @@ def split_space(start, end, count):
  return list(range(start, end, (size + 1) // (count + 1)))[1 : count + 1]
 
 
-class ArithUUID(UUID):
+class ArithString:
+ @classmethod
+ def new(cls, *args, **kw):
+ return cls(*args, **kw)
+
+
+class ArithUUID(UUID, ArithString):
  "A UUID that supports basic arithmetic (add, sub)"
 
+ def __int__(self):
+ return self.int
+
  def __add__(self, other: Union[UUID, int]):
  if isinstance(other, int):
- return type(self)(int=self.int + other)
+ return self.new(int=self.int + other)
  return NotImplemented
 
  def __sub__(self, other: Union[UUID, int]):
  if isinstance(other, int):
- return type(self)(int=self.int - other)
+ return self.new(int=self.int - other)
  elif isinstance(other, UUID):
  return self.int - other.int
  return NotImplemented
 
+
+def numberToBase(num, base):
+ digits = []
+ while num > 0:
+ num, remainder = divmod(num, base)
+ digits.append(remainder)
+ return "".join(alphanums[i] for i in digits[::-1])
+
+
+class ArithAlphanumeric(ArithString):
+ def __init__(self, str: str = None, int: int = None, max_len=None):
+ if str is None:
+ str = numberToBase(int, len(alphanums))
+ else:
+ assert int is None
+
+ if max_len and len(str) > max_len:
+ raise ValueError(f"Length of alphanum value '{str}' is longer than the expected {max_len}")
+
+ self._str = str
+ self._max_len = max_len
+
+ @property
+ def int(self):
+ return int(self._str, len(alphanums))
+
+ def __str__(self):
+ s = self._str
+ if self._max_len:
+ s = s.rjust(self._max_len, "0")
+ return s
+
+ def __len__(self):
+ return len(self._str)
+
  def __int__(self):
  return self.int
 
+ def __repr__(self):
+ return f'alphanum"{self._str}"'
+
+ def __add__(self, other: "Union[ArithAlphanumeric, int]"):
+ if isinstance(other, int):
+ res = self.new(int=self.int + other)
+ if len(str(res)) != len(self):
+ raise ValueError("Overflow error when adding to alphanumeric")
+ return res
+ return NotImplemented
+
+ def __sub__(self, other: "Union[ArithAlphanumeric, int]"):
+ if isinstance(other, int):
+ return type(self)(int=self.int - other)
+ elif isinstance(other, ArithAlphanumeric):
+ return self.int - other.int
+ return NotImplemented
+
+ def __ge__(self, other):
+ if not isinstance(other, type(self)):
+ return NotImplemented
+ return self.int >= other.int
+
+ def __lt__(self, other):
+ if not isinstance(other, type(self)):
+ return NotImplemented
+ return self.int < other.int
+
+ def new(self, *args, **kw):
+ return type(self)(*args, **kw, max_len=self._max_len)
+
 
 def is_uuid(u):
  try:
@@ -57,23 +135,24 @@ def number_to_human(n):
 def _join_if_any(sym, args):
  args = list(args)
  if not args:
- return ''
+ return ""
  return sym.join(str(a) for a in args if a)
 
-def remove_password_from_url(url: str, replace_with: str="***") -> str:
+
+def remove_password_from_url(url: str, replace_with: str = "***") -> str:
  parsed = urlparse(url)
- account = parsed.username or ''
+ account = parsed.username or ""
  if parsed.password:
- account += ':' + replace_with
+ account += ":" + replace_with
  host = _join_if_any(":", filter(None, [parsed.hostname, parsed.port]))
  netloc = _join_if_any("@", filter(None, [account, host]))
  replaced = parsed._replace(netloc=netloc)
  return replaced.geturl()
 
+
 def join_iter(joiner: Any, iterable: iter) -> iter:
  it = iter(iterable)
  yield next(it)
  for i in it:
  yield joiner
  yield i
-

tests/test_config.py

@@ -46,16 +46,16 @@ def test_basic(self):
  def test_remove_password(self):
  replace_with = "*****"
  urls = [
- 'd://host/',
- 'd://host:123/',
- 'd://user@host:123/',
- 'd://user:PASS@host:123/',
- 'd://:PASS@host:123/',
- 'd://:PASS@host:123/path',
- 'd://:PASS@host:123/path?whatever#blabla',
+ "d://host/",
+ "d://host:123/",
+ "d://user@host:123/",
+ "d://user:PASS@host:123/",
+ "d://:PASS@host:123/",
+ "d://:PASS@host:123/path",
+ "d://:PASS@host:123/path?whatever#blabla",
  ]
  for url in urls:
  removed = remove_password_from_url(url, replace_with)
- expected = url.replace('PASS', replace_with)
+ expected = url.replace("PASS", replace_with)
  removed = remove_password_from_url(url, replace_with)
  self.assertEqual(removed, expected)

tests/test_diff_tables.py

@@ -9,6 +9,7 @@
 from data_diff.databases import connect_to_uri
 from data_diff.diff_tables import TableDiffer, TableSegment, split_space
 from data_diff import databases as db
+from data_diff.utils import ArithAlphanumeric
 
 from .common import (
  TEST_MYSQL_CONN_STRING,
@@ -369,7 +370,7 @@ def test_diff_sorted_by_key(self):
 
 
 @test_per_database
-class TestStringKeys(TestPerDatabase):
+class TestUUIDs(TestPerDatabase):
  def setUp(self):
  super().setUp()
 
@@ -408,6 +409,57 @@ def test_string_keys(self):
  self.assertRaises(ValueError, list, differ.diff_tables(self.a, self.b))
 
 
+@test_per_database
+class TestAlphanumericKeys(TestPerDatabase):
+ def setUp(self):
+ super().setUp()
+
+ queries = [
+ f"CREATE TABLE {self.table_src}(id varchar(100), text_comment varchar(1000))",
+ ]
+ for i in range(0, 10000, 1000):
+ queries.append(f"INSERT INTO {self.table_src} VALUES ('{ArithAlphanumeric(int=i, max_len=10)}', '{i}')")
+
+ queries += [
+ f"CREATE TABLE {self.table_dst} AS SELECT * FROM {self.table_src}",
+ ]
+
+ self.new_alphanum = "abcdefghij"
+ queries.append(f"INSERT INTO {self.table_src} VALUES ('{self.new_alphanum}', 'This one is different')")
+
+ # TODO test unexpected values?
+
+ for query in queries:
+ self.connection.query(query, None)
+
+ _commit(self.connection)
+
+ self.a = TableSegment(self.connection, self.table_src_path, "id", "text_comment", case_sensitive=False)
+ self.b = TableSegment(self.connection, self.table_dst_path, "id", "text_comment", case_sensitive=False)
+
+ def test_alphanum_keys(self):
+ # Test the class itself
+ assert str(ArithAlphanumeric(int=0, max_len=1)) == "0"
+ assert str(ArithAlphanumeric(int=0, max_len=10)) == "0" * 10
+ assert str(ArithAlphanumeric(int=1, max_len=10)) == "0" * 9 + "1"
+
+ # Test in the differ
+
+ differ = TableDiffer()
+ diff = list(differ.diff_tables(self.a, self.b))
+ self.assertEqual(diff, [("-", (str(self.new_alphanum), "This one is different"))])
+
+ self.connection.query(
+ f"INSERT INTO {self.table_src} VALUES ('@@@', '<-- this bad value should not break us')", None
+ )
+ _commit(self.connection)
+
+ self.a = TableSegment(self.connection, self.table_src_path, "id", "text_comment", case_sensitive=False)
+ self.b = TableSegment(self.connection, self.table_dst_path, "id", "text_comment", case_sensitive=False)
+
+ self.assertRaises(NotImplementedError, list, differ.diff_tables(self.a, self.b))
+
+
 @test_per_database
 class TestTableSegment(TestPerDatabase):
  def setUp(self) -> None: