Now samples all TEXT columns at once; Added utils.py

Author: erezsh

data_diff/databases/base.py

@@ -8,6 +8,7 @@
 import threading
 from abc import abstractmethod
 
+from data_diff.utils import is_uuid, safezip
 from .database_types import (
     ColType_UUID,
     AbstractDatabase,
@@ -29,14 +30,6 @@ 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)
@@ -172,18 +165,7 @@ 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()
+            return cls()
 
         raise TypeError(f"Parsing {type_repr} returned an unknown type '{cls}'.")
 
@@ -204,8 +186,31 @@ def query_table_schema(self, path: DbPath, filter_columns: Optional[Sequence[str
             accept = {i.lower() for i in filter_columns}
             rows = [r for r in rows if r[0].lower() in accept]
 
+        col_dict: Dict[str, ColType] = {row[0]: self._parse_type(path, *row) for row in rows}
+
+        self._refine_coltypes(path, col_dict)
+
         # Return a dict of form {name: type} after normalization
-        return {row[0]: self._parse_type(path, *row) for row in rows}
+        return col_dict
+
+    def _refine_coltypes(self, table_path: DbPath, col_dict: Dict[str, ColType]):
+        "Refine the types in the column dict, by querying the database for a sample of their values"
+
+        text_columns = [k for k, v in col_dict.items() if isinstance(v, Text)]
+
+        samples_by_row = self.query(Select(text_columns, TableName(table_path), limit=16), list)
+        samples_by_col = list(zip(*samples_by_row))
+        for col_name, samples in safezip(text_columns, samples_by_col):
+            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:
+                    assert col_name in col_dict
+                    col_dict[col_name] = ColType_UUID()
 
     # @lru_cache()
     # def get_table_schema(self, path: DbPath) -> Dict[str, ColType]:

data_diff/databases/database_types.py

@@ -1,28 +1,10 @@
-from uuid import UUID
 from abc import ABC, abstractmethod
 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
+from data_diff.utils import ArithUUID
 
 
 DbPath = Tuple[str, ...]
@@ -31,6 +13,7 @@ def __int__(self):
 
 
 class ColType:
+    supported = True
     pass
 
 
@@ -87,8 +70,9 @@ class ColType_UUID(StringType, IKey):
     python_type = ArithUUID
 
 
+@dataclass
 class Text(StringType):
-    pass
+    supported = False
 
 
 @dataclass
@@ -104,6 +88,8 @@ def __post_init__(self):
 class UnknownColType(ColType):
     text: str
 
+    supported = False
+
 
 class AbstractDatabase(ABC):
     @abstractmethod

data_diff/diff_tables.py

@@ -12,6 +12,7 @@
 from runtype import dataclass
 
 from .sql import Select, Checksum, Compare, DbPath, DbKey, DbTime, Count, TableName, Time, Min, Max, Value
+from .utils import safezip, split_space
 from .databases.base import Database
 from .databases.database_types import (
     ArithUUID,
@@ -31,17 +32,6 @@
 DEFAULT_BISECTION_FACTOR = 32
 
 
-def safezip(*args):
-    "zip but makes sure all sequences are the same length"
-    assert len(set(map(len, args))) == 1
-    return zip(*args)
-
-
-def split_space(start, end, count):
-    size = end - start
-    return list(range(start, end, (size + 1) // (count + 1)))[1 : count + 1]
-
-
 @dataclass(frozen=False)
 class TableSegment:
     """Signifies a segment of rows (and selected columns) within a table
@@ -359,9 +349,9 @@ def _validate_and_adjust_columns(self, table1, table2):
         for t in [table1, table2]:
             for c in t._relevant_columns:
                 ctype = t._schema[c]
-                if isinstance(ctype, UnknownColType):
-                    logger.warn(
-                        f"[{t.database.name}] Column '{c}' of type '{ctype.text}' has no compatibility handling. "
+                if not ctype.supported:
+                    logger.warning(
+                        f"[{t.database.name}] Column '{c}' of type '{ctype}' has no compatibility handling. "
                         "If encoding/formatting differs between databases, it may result in false positives."
                     )
 

data_diff/utils.py

@@ -0,0 +1,40 @@
+from typing import Sequence, Optional, Tuple, Union, Dict, Any
+from uuid import UUID
+
+
+def safezip(*args):
+    "zip but makes sure all sequences are the same length"
+    assert len(set(map(len, args))) == 1
+    return zip(*args)
+
+
+def split_space(start, end, count):
+    size = end - start
+    return list(range(start, end, (size + 1) // (count + 1)))[1 : count + 1]
+
+
+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
+
+
+def is_uuid(u):
+    try:
+        UUID(u)
+    except ValueError:
+        return False
+    return True