UUID field support, initial (implemented only for MySQL)

Author: erezsh

data_diff/databases/base.py

@@ -1,13 +1,15 @@
+from uuid import UUID
 import math
 import sys
 import logging
-from typing import Dict, Tuple, Optional, Sequence, Type
+from typing import Dict, Tuple, Optional, Sequence, Type, List
 from functools import lru_cache, wraps
 from concurrent.futures import ThreadPoolExecutor
 import threading
 from abc import abstractmethod
 
 from .database_types import (
+    ColType_UUID,
     AbstractDatabase,
     ColType,
     Integer,
@@ -16,8 +18,9 @@
     PrecisionType,
     TemporalType,
     UnknownColType,
+    Text,
 )
-from data_diff.sql import DbPath, SqlOrStr, Compiler, Explain, Select
+from data_diff.sql import DbPath, SqlOrStr, Compiler, Explain, Select, TableName
 
 logger = logging.getLogger("database")
 
@@ -26,6 +29,14 @@ def parse_table_name(t):
     return tuple(t.split("."))
 
 
+def is_uuid(u):
+    try:
+        UUID(u)
+    except ValueError:
+        return False
+    return True
+
+
 def import_helper(package: str = None, text=""):
     def dec(f):
         @wraps(f)
@@ -102,7 +113,7 @@ def query(self, sql_ast: SqlOrStr, res_type: type):
             assert len(res) == 1, (sql_code, res)
             return res[0]
         elif getattr(res_type, "__origin__", None) is list and len(res_type.__args__) == 1:
-            if res_type.__args__ == (int,):
+            if res_type.__args__ == (int,) or res_type.__args__ == (str,):
                 return [_one(row) for row in res]
             elif res_type.__args__ == (Tuple,):
                 return [tuple(row) for row in res]
@@ -123,6 +134,7 @@ def _parse_type_repr(self, type_repr: str) -> Optional[Type[ColType]]:
 
     def _parse_type(
         self,
+        table_path: DbPath,
         col_name: str,
         type_repr: str,
         datetime_precision: int = None,
@@ -147,7 +159,7 @@ def _parse_type(
         elif issubclass(cls, Decimal):
             if numeric_scale is None:
                 raise ValueError(
-                    f"{self.name}: Unexpected numeric_scale is NULL, for column {col_name} of type {type_repr}."
+                    f"{self.name}: Unexpected numeric_scale is NULL, for column {'.'.join(table_path)}.{col_name} of type {type_repr}."
                 )
             return cls(precision=numeric_scale)
 
@@ -159,6 +171,20 @@ def _parse_type(
                 )
             )
 
+        elif issubclass(cls, Text):
+            samples = self.query(Select([col_name], TableName(table_path), limit=16), List[str])
+            uuid_samples = list(filter(is_uuid, samples))
+
+            if uuid_samples:
+                if len(uuid_samples) != len(samples):
+                    logger.warning(
+                        f"Mixed UUID/Non-UUID values detected in column {'.'.join(table_path)}.{col_name}, disabling UUID support."
+                    )
+                else:
+                    return ColType_UUID()
+
+            return Text()
+
         raise TypeError(f"Parsing {type_repr} returned an unknown type '{cls}'.")
 
     def select_table_schema(self, path: DbPath) -> str:
@@ -179,7 +205,7 @@ def query_table_schema(self, path: DbPath, filter_columns: Optional[Sequence[str
             rows = [r for r in rows if r[0].lower() in accept]
 
         # Return a dict of form {name: type} after normalization
-        return {row[0]: self._parse_type(*row) for row in rows}
+        return {row[0]: self._parse_type(path, *row) for row in rows}
 
     # @lru_cache()
     # def get_table_schema(self, path: DbPath) -> Dict[str, ColType]:
@@ -233,6 +259,12 @@ def create_connection(self):
     def close(self):
         self._queue.shutdown()
 
+    def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None):
+        if offset:
+            raise NotImplementedError("No support for OFFSET in query")
+
+        return f"LIMIT {limit}"
+
 
 CHECKSUM_HEXDIGITS = 15  # Must be 15 or lower
 MD5_HEXDIGITS = 32

data_diff/databases/database_types.py

@@ -1,11 +1,32 @@
+from uuid import UUID
 from abc import ABC, abstractmethod
-from typing import Sequence, Optional, Tuple, Union, Dict
+from typing import Sequence, Optional, Tuple, Union, Dict, Any
 from datetime import datetime
 
 from runtype import dataclass
 
+
+class ArithUUID(UUID):
+    "A UUID that supports basic arithmetic (add, sub)"
+
+    def __add__(self, other: Union[UUID, int]):
+        if isinstance(other, int):
+            return type(self)(int=self.int + other)
+        return NotImplemented
+
+    def __sub__(self, other: Union[UUID, int]):
+        if isinstance(other, int):
+            return type(self)(int=self.int - other)
+        elif isinstance(other, UUID):
+            return self.int - other.int
+        return NotImplemented
+
+    def __int__(self):
+        return self.int
+
+
 DbPath = Tuple[str, ...]
-DbKey = Union[int, str, bytes]
+DbKey = Union[int, str, bytes, ArithUUID]
 DbTime = datetime
 
 
@@ -57,12 +78,24 @@ class StringType(ColType):
     pass
 
 
-class UUID(StringType):
+class IKey(ABC):
+    "Interface for ColType, for using a column as a key in data-diff"
+    python_type: type
+
+
+class ColType_UUID(StringType, IKey):
+    python_type = ArithUUID
+
+
+class Text(StringType):
     pass
 
 
 @dataclass
-class Integer(NumericType):
+class Integer(NumericType, IKey):
+    precision: int = 0
+    python_type: type = int
+
     def __post_init__(self):
         assert self.precision == 0
 
@@ -88,6 +121,10 @@ def md5_to_int(self, s: str) -> str:
         "Provide SQL for computing md5 and returning an int"
         ...
 
+    @abstractmethod
+    def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None):
+        ...
+
     @abstractmethod
     def _query(self, sql_code: str) -> list:
         "Send query to database and return result"

data_diff/databases/mysql.py

@@ -20,6 +20,8 @@ class MySQL(ThreadedDatabase):
         "float": Float,
         "decimal": Decimal,
         "int": Integer,
+        # Text
+        "varchar": Text,
     }
     ROUNDS_ON_PREC_LOSS = True
 

data_diff/databases/oracle.py

@@ -90,3 +90,9 @@ def _parse_type(
                 )
 
         return super()._parse_type(type_repr, col_name, type_repr, datetime_precision, numeric_precision, numeric_scale)
+
+    def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None):
+        if offset:
+            raise NotImplementedError("No support for OFFSET in query")
+
+        return f"FETCH NEXT {limit} ROWS ONLY"

data_diff/databases/postgresql.py

@@ -25,6 +25,9 @@ class PostgreSQL(ThreadedDatabase):
         "integer": Integer,
         "numeric": Decimal,
         "bigint": Integer,
+        # Text
+        "varchar": Text,
+        "text": Text,
     }
     ROUNDS_ON_PREC_LOSS = True
 

data_diff/databases/presto.py

@@ -81,7 +81,12 @@ def select_table_schema(self, path: DbPath) -> str:
         )
 
     def _parse_type(
-        self, col_name: str, type_repr: str, datetime_precision: int = None, numeric_precision: int = None
+        self,
+        table_path: DbPath,
+        col_name: str,
+        type_repr: str,
+        datetime_precision: int = None,
+        numeric_precision: int = None,
     ) -> ColType:
         timestamp_regexps = {
             r"timestamp\((\d)\)": Timestamp,
@@ -103,4 +108,4 @@ def _parse_type(
                 prec, scale = map(int, m.groups())
                 return n_cls(scale)
 
-        return super()._parse_type(type_repr)
+        return super()._parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision)

data_diff/diff_tables.py

@@ -11,9 +11,10 @@
 
 from runtype import dataclass
 
-from .sql import Select, Checksum, Compare, DbPath, DbKey, DbTime, Count, TableName, Time, Min, Max
+from .sql import Select, Checksum, Compare, DbPath, DbKey, DbTime, Count, TableName, Time, Min, Max, Value
 from .databases.base import Database
 from .databases.database_types import (
+    ArithUUID,
     NumericType,
     PrecisionType,
     UnknownColType,
@@ -121,9 +122,9 @@ def with_schema(self) -> "TableSegment":
 
     def _make_key_range(self):
         if self.min_key is not None:
-            yield Compare("<=", str(self.min_key), self._key_column)
+            yield Compare("<=", Value(self.min_key), self._key_column)
         if self.max_key is not None:
-            yield Compare("<", self._key_column, str(self.max_key))
+            yield Compare("<", self._key_column, Value(self.max_key))
 
     def _make_update_range(self):
         if self.min_update is not None:
@@ -152,6 +153,11 @@ 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):
+            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 split_space(self.min_key, self.max_key, count)
 
     def segment_by_checkpoints(self, checkpoints: List[DbKey]) -> List["TableSegment"]:
@@ -297,9 +303,12 @@ 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"]
+        assert type(key_type) == type(table2._schema["id"])
+
         # We add 1 because our ranges are exclusive of the end (like in Python)
-        min_key = min(map(int, mins))
-        max_key = max(map(int, maxs)) + 1
+        min_key = min(map(key_type.python_type, mins))
+        max_key = max(map(key_type.python_type, maxs)) + 1
 
         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

@@ -6,7 +6,7 @@
 
 from runtype import dataclass
 
-from .databases.database_types import AbstractDatabase, DbPath, DbKey, DbTime
+from .databases.database_types import AbstractDatabase, DbPath, DbKey, DbTime, ArithUUID
 
 
 class Sql:
@@ -65,6 +65,8 @@ 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):
+            return "'%s'" % self.value
         return str(self.value)
 
 
@@ -75,6 +77,7 @@ class Select(Sql):
     where: Sequence[SqlOrStr] = None
     order_by: Sequence[SqlOrStr] = None
     group_by: Sequence[SqlOrStr] = None
+    limit: int = None
 
     def compile(self, parent_c: Compiler):
         c = parent_c.replace(in_select=True)
@@ -93,6 +96,9 @@ def compile(self, parent_c: Compiler):
         if self.order_by:
             select += " ORDER BY " + ", ".join(map(c.compile, self.order_by))
 
+        if self.limit is not None:
+            select += " " + c.database.offset_limit(0, self.limit)
+
         if parent_c.in_select:
             select = "(%s)" % select
         return select

tests/test_diff_tables.py

@@ -259,7 +259,11 @@ def setUp(self):
             "COMMIT",
         ]
 
-        queries.append(f"INSERT INTO a VALUES ('unexpected', '<-- this bad value should not break us')")
+        self.new_uuid = uuid.uuid1(32132131)
+        queries.append(f"INSERT INTO a VALUES ('{self.new_uuid}', 'This one is different')")
+
+        # TODO test unexpected values?
+        # queries.append(f"INSERT INTO a VALUES ('unexpected', '<-- this bad value should not break us')")
 
         for query in queries:
             self.connection.query(query, None)
@@ -270,7 +274,7 @@ def setUp(self):
     def test_string_keys(self):
         differ = TableDiffer()
         diff = list(differ.diff_tables(self.a, self.b))
-        breakpoint()
+        self.assertEqual(diff, [("-", (str(self.new_uuid), "This one is different"))])
 
 
 class TestTableSegment(TestWithConnection):