Merge pull request #133 from datafold/fix_uuids

Fix UUIDs + small fix for presto

Author: erezsh

data_diff/databases/base.py

@@ -1,4 +1,3 @@
-from uuid import UUID
 import math
 import sys
 import logging
@@ -16,7 +15,6 @@
     Integer,
     Decimal,
     Float,
-    PrecisionType,
     TemporalType,
     UnknownColType,
     Text,

data_diff/databases/mysql.py

@@ -70,4 +70,4 @@ def normalize_number(self, value: str, coltype: FractionalType) -> str:
         return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))")
 
     def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
-        return f"CAST(TRIM({value}) AS char)"
+        return f"TRIM(CAST({value} AS char))"

data_diff/databases/presto.py

@@ -27,6 +27,7 @@ class Presto(Database):
         "timestamp": Timestamp,
         # Numbers
         "integer": Integer,
+        "bigint": Integer,
         "real": Float,
         "double": Float,
         # Text

data_diff/diff_tables.py

@@ -79,19 +79,27 @@ def __post_init__(self):
         if self.min_update is not None and self.max_update is not None and self.min_update >= self.max_update:
             raise ValueError("Error: min_update expected to be smaller than max_update!")
 
-    @property
-    def _key_column(self):
-        return self._quote_column(self.key_column)
-
     @property
     def _update_column(self):
         return self._quote_column(self.update_column)
 
-    def _quote_column(self, c):
+    def _quote_column(self, c: str) -> str:
         if self._schema:
             c = self._schema.get_key(c)  # Get the actual name. Might be case-insensitive.
         return self.database.quote(c)
 
+    def _normalize_column(self, name: str, template: str = None) -> str:
+        if not self._schema:
+            raise RuntimeError(
+                "Cannot compile query when the schema is unknown. Please use TableSegment.with_schema()."
+            )
+
+        col = self._quote_column(name)
+        if template is not None:
+            col = template % col  # Apply template using Python's string formatting
+
+        return self.database.normalize_value_by_type(col, self._schema[name])
+
     def with_schema(self) -> "TableSegment":
         "Queries the table schema from the database, and returns a new instance of TableSegmentWithSchema."
         if self._schema:
@@ -115,9 +123,9 @@ def with_schema(self) -> "TableSegment":
 
     def _make_key_range(self):
         if self.min_key is not None:
-            yield Compare("<=", Value(self.min_key), self._key_column)
+            yield Compare("<=", Value(self.min_key), self._quote_column(self.key_column))
         if self.max_key is not None:
-            yield Compare("<", self._key_column, Value(self.max_key))
+            yield Compare("<", self._quote_column(self.key_column), Value(self.max_key))
 
     def _make_update_range(self):
         if self.min_update is not None:
@@ -127,7 +135,7 @@ def _make_update_range(self):
 
     def _make_select(self, *, table=None, columns=None, where=None, group_by=None, order_by=None):
         if columns is None:
-            columns = [self._key_column]
+            columns = [self._normalize_column(self.key_column)]
         where = list(self._make_key_range()) + list(self._make_update_range()) + ([] if where is None else [where])
         order_by = None if order_by is None else [order_by]
         return Select(
@@ -184,14 +192,7 @@ def _relevant_columns(self) -> List[str]:
 
     @property
     def _relevant_columns_repr(self) -> List[str]:
-        if not self._schema:
-            raise RuntimeError(
-                "Cannot compile query when the schema is unknown. Please use TableSegment.with_schema()."
-            )
-        return [
-            self.database.normalize_value_by_type(self._quote_column(c), self._schema[c])
-            for c in self._relevant_columns
-        ]
+        return [self._normalize_column(c) for c in self._relevant_columns]
 
     def count(self) -> Tuple[int, int]:
         """Count how many rows are in the segment, in one pass."""
@@ -214,7 +215,13 @@ def count_and_checksum(self) -> Tuple[int, int]:
 
     def query_key_range(self) -> Tuple[int, int]:
         """Query database for minimum and maximum key. This is used for setting the initial bounds."""
-        select = self._make_select(columns=[Min(self._key_column), Max(self._key_column)])
+        # Normalizes the result (needed for UUIDs) after the min/max computation
+        select = self._make_select(
+            columns=[
+                self._normalize_column(self.key_column, "min(%s)"),
+                self._normalize_column(self.key_column, "max(%s)"),
+            ]
+        )
         min_key, max_key = self.database.query(select, tuple)
 
         if min_key is None or max_key is None:
@@ -296,13 +303,16 @@ def diff_tables(self, table1: TableSegment, table2: TableSegment) -> DiffResult:
         key_ranges = self._threaded_call("query_key_range", [table1, table2])
         mins, maxs = zip(*key_ranges)
 
-        key_type = table1._schema["id"]
-        key_type2 = table2._schema["id"]
+        key_type = table1._schema[table1.key_column]
+        key_type2 = table2._schema[table2.key_column]
         assert key_type.python_type is key_type2.python_type
 
         # We add 1 because our ranges are exclusive of the end (like in Python)
-        min_key = min(map(key_type.python_type, mins))
-        max_key = max(map(key_type.python_type, maxs)) + 1
+        try:
+            min_key = min(map(key_type.python_type, mins))
+            max_key = max(map(key_type.python_type, maxs)) + 1
+        except (TypeError, ValueError) as e:
+            raise type(e)(f"Cannot apply {key_type} to {mins}, {maxs}.") from e
 
         table1 = table1.new(min_key=min_key, max_key=max_key)
         table2 = table2.new(min_key=min_key, max_key=max_key)

data_diff/sql.py

@@ -124,7 +124,7 @@ def compile(self, c: Compiler):
             compiled_exprs = ", ".join(map(c.compile, self.exprs))
             expr = f"concat({compiled_exprs})"
         else:
-            expr ,= self.exprs
+            (expr,) = self.exprs
             expr = c.compile(expr)
         md5 = c.database.md5_to_int(expr)
         return f"sum({md5})"

tests/common.py

@@ -17,10 +17,12 @@
 N_SAMPLES = int(os.environ.get("N_SAMPLES", DEFAULT_N_SAMPLES))
 BENCHMARK = os.environ.get("BENCHMARK", False)
 
+
 def get_git_revision_short_hash() -> str:
-    return subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).decode('ascii').strip()
+    return subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
+
 
-GIT_REVISION=get_git_revision_short_hash()
+GIT_REVISION = get_git_revision_short_hash()
 
 level = logging.ERROR
 if os.environ.get("LOG_LEVEL", False):

tests/test_database_types.py

@@ -186,7 +186,7 @@ def __iter__(self):
         "int": [
             # "smallint",  # 2 bytes
             "int",  # 4 bytes
-            "bigint", # 8 bytes
+            "bigint",  # 8 bytes
         ],
         # https://www.postgresql.org/docs/current/datatype-datetime.html
         "datetime": [
@@ -214,7 +214,7 @@ def __iter__(self):
             # "smallint", # 2 bytes
             # "mediumint", # 3 bytes
             "int",  # 4 bytes
-            "bigint", # 8 bytes
+            "bigint",  # 8 bytes
         ],
         # https://dev.mysql.com/doc/refman/8.0/en/datetime.html
         "datetime": [
@@ -327,7 +327,7 @@ def __iter__(self):
             # "smallint", # 2 bytes
             # "mediumint", # 3 bytes
             "int",  # 4 bytes
-            "bigint", # 8 bytes
+            "bigint",  # 8 bytes
         ],
         "datetime": [
             "timestamp",
@@ -548,8 +548,12 @@ def test_types(self, source_db, target_db, source_type, target_type, type_catego
         _insert_to_table(dst_conn, dst_table, values_in_source, target_type)
         insertion_target_duration = time.time() - start
 
-        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)
+        if type_category == "uuid":
+            self.table = TableSegment(self.src_conn, src_table_path, "col", None, ("id",), case_sensitive=False)
+            self.table2 = TableSegment(self.dst_conn, dst_table_path, "col", None, ("id",), case_sensitive=False)
+        else:
+            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)
 
         start = time.time()
         self.assertEqual(N_SAMPLES, self.table.count())
@@ -595,7 +599,10 @@ def test_types(self, source_db, target_db, source_type, target_type, type_catego
         download_duration = time.time() - start
         expected = []
         self.assertEqual(expected, diff)
-        self.assertEqual(len(sample_values), differ.stats.get("rows_downloaded", 0))
+        if type_category == "uuid":
+            pass  # UUIDs aren't serial, so they mess with the first max_rows estimation.
+        else:
+            self.assertEqual(len(sample_values), differ.stats.get("rows_downloaded", 0))
 
         result = {
             "test": self._testMethodName,