Added support for native UUIDs in postgresql.

Author: erezsh

data_diff/databases/base.py

@@ -9,12 +9,14 @@
 
 from data_diff.utils import is_uuid, safezip
 from .database_types import (
-    ColType_UUID,
     AbstractDatabase,
     ColType,
     Integer,
     Decimal,
     Float,
+    ColType_UUID,
+    Native_UUID,
+    String_UUID,
     TemporalType,
     UnknownColType,
     Text,
@@ -162,7 +164,7 @@ def _parse_type(
                 )
             )
 
-        elif issubclass(cls, Text):
+        elif issubclass(cls, (Text, Native_UUID)):
             return cls()
 
         raise TypeError(f"Parsing {type_repr} returned an unknown type '{cls}'.")
@@ -198,7 +200,7 @@ def _refine_coltypes(self, table_path: DbPath, col_dict: Dict[str, ColType]):
         if not text_columns:
             return
 
-        fields = [self.normalize_uuid(c, ColType_UUID()) for c in text_columns]
+        fields = [self.normalize_uuid(c, String_UUID()) for c in text_columns]
         samples_by_row = self.query(Select(fields, TableName(table_path), limit=16), list)
         if not samples_by_row:
             logger.warning(f"Table {table_path} is empty.")
@@ -216,7 +218,7 @@ def _refine_coltypes(self, table_path: DbPath, col_dict: Dict[str, ColType]):
                     )
                 else:
                     assert col_name in col_dict
-                    col_dict[col_name] = ColType_UUID()
+                    col_dict[col_name] = String_UUID()
 
     # @lru_cache()
     # def get_table_schema(self, path: DbPath) -> Dict[str, ColType]:
@@ -241,7 +243,9 @@ def offset_limit(self, offset: Optional[int] = None, limit: Optional[int] = None
         return f"LIMIT {limit}"
 
     def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
-        return f"TRIM({value})"
+        if isinstance(coltype, String_UUID):
+            return f"TRIM({value})"
+        return self.to_string(value)
 
 
 class ThreadedDatabase(Database):

data_diff/databases/database_types.py

@@ -53,11 +53,13 @@ class FractionalType(NumericType):
 class Float(FractionalType):
     pass
 
+
 class IKey(ABC):
     "Interface for ColType, for using a column as a key in data-diff"
     python_type: type
 
-class Decimal(FractionalType, IKey):    # Snowflake may use Decimal as a key
+
+class Decimal(FractionalType, IKey):  # Snowflake may use Decimal as a key
     @property
     def python_type(self) -> type:
         if self.precision == 0:
@@ -69,10 +71,18 @@ class StringType(ColType):
     pass
 
 
-class ColType_UUID(StringType, IKey):
+class ColType_UUID(ColType, IKey):
     python_type = ArithUUID
 
 
+class Native_UUID(ColType_UUID):
+    pass
+
+
+class String_UUID(StringType, ColType_UUID):
+    pass
+
+
 @dataclass
 class Text(StringType):
     supported = False

data_diff/databases/postgresql.py

@@ -30,6 +30,8 @@ class PostgreSQL(ThreadedDatabase):
         "character varying": Text,
         "varchar": Text,
         "text": Text,
+        # UUID
+        "uuid": Native_UUID,
     }
     ROUNDS_ON_PREC_LOSS = True
 

data_diff/diff_tables.py

@@ -19,10 +19,10 @@
 from .databases.database_types import (
     ArithUUID,
     IKey,
+    Native_UUID,
     NumericType,
     PrecisionType,
     StringType,
-    UnknownColType,
     Schema,
     Schema_CaseInsensitive,
     Schema_CaseSensitive,
@@ -96,11 +96,22 @@ def _normalize_column(self, name: str, template: str = None) -> str:
                 "Cannot compile query when the schema is unknown. Please use TableSegment.with_schema()."
             )
 
+        col_type = self._schema[name]
         col = self._quote_column(name)
+
+        if isinstance(col_type, Native_UUID):
+            # Normalize first, apply template after (for uuids)
+            # Needed because min/max(uuid) fails in postgresql
+            col = self.database.normalize_value_by_type(col, col_type)
+            if template is not None:
+                col = template % col  # Apply template using Python's string formatting
+            return col
+
+        # Apply template before normalizing (for ints)
         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])
+        return self.database.normalize_value_by_type(col, col_type)
 
     def with_schema(self) -> "TableSegment":
         "Queries the table schema from the database, and returns a new instance of TableSegmentWithSchema."
@@ -274,7 +285,7 @@ class TableDiffer:
     """
 
     bisection_factor: int = DEFAULT_BISECTION_FACTOR
-    bisection_threshold: Number = DEFAULT_BISECTION_THRESHOLD       # Accepts inf for tests
+    bisection_threshold: Number = DEFAULT_BISECTION_THRESHOLD  # Accepts inf for tests
     threaded: bool = True
     max_threadpool_size: Optional[int] = 1
 

tests/test_diff_tables.py

@@ -339,10 +339,8 @@ def setUp(self):
         self.null_uuid = uuid.uuid1(32132131)
         queries += [
             f"CREATE TABLE {self.table_dst} AS SELECT * FROM {self.table_src}",
-
             f"INSERT INTO {self.table_src} VALUES ('{self.null_uuid}', NULL)",
-
-            "COMMIT"
+            "COMMIT",
         ]
 
         for query in queries:
@@ -366,15 +364,11 @@ def setUp(self):
             f"DROP TABLE IF EXISTS {self.table_src}",
             f"DROP TABLE IF EXISTS {self.table_dst}",
             f"CREATE TABLE {self.table_src}(id varchar(100), comment varchar(1000))",
-
             f"INSERT INTO {self.table_src} VALUES ('{uuid.uuid1(1)}', '1')",
-
             f"CREATE TABLE {self.table_dst} AS SELECT * FROM {self.table_src}",
-
             # Add a row where a column has NULL value
             f"INSERT INTO {self.table_src} VALUES ('{self.null_uuid}', NULL)",
-
-            "COMMIT"
+            "COMMIT",
         ]
 
         for query in queries:
@@ -500,7 +494,7 @@ def test_left_table_empty(self):
         queries = [
             f"INSERT INTO {self.table_dst} SELECT id, comment FROM {self.table_src}",
             f"TRUNCATE {self.table_src}",
-            "COMMIT"
+            "COMMIT",
         ]
         for query in queries:
             self.connection.query(query, None)

tests/test_postgresql.py

@@ -0,0 +1,46 @@
+import unittest
+
+from data_diff.databases.connect import connect_to_uri
+from data_diff import TableSegment, TableDiffer
+from .common import TEST_POSTGRESQL_CONN_STRING, random_table_suffix
+
+
+class TestWithConnection(unittest.TestCase):
+    def setUp(self) -> None:
+        self.connection = connect_to_uri(TEST_POSTGRESQL_CONN_STRING)
+
+        self.connection.query('CREATE EXTENSION IF NOT EXISTS "uuid-ossp";', None)
+
+        table_suffix = random_table_suffix()
+
+        self.table_src = f"src{table_suffix}"
+        self.table_dst = f"dst{table_suffix}"
+
+    def test_uuid(self):
+        queries = [
+            f"DROP TABLE IF EXISTS {self.table_src}",
+            f"DROP TABLE IF EXISTS {self.table_dst}",
+            f"CREATE TABLE {self.table_src} (id uuid DEFAULT uuid_generate_v4 (), comment VARCHAR, PRIMARY KEY (id))",
+            "COMMIT",
+        ]
+        for i in range(100):
+            queries.append(f"INSERT INTO {self.table_src}(comment) VALUES ('{i}')")
+
+        queries += [
+            "COMMIT",
+            f"CREATE TABLE {self.table_dst} AS SELECT * FROM {self.table_src}",
+            "COMMIT",
+        ]
+
+        queries.append(f"INSERT INTO {self.table_src}(comment) VALUES ('This one is different')")
+
+        for query in queries:
+            self.connection.query(query, None)
+
+        a = TableSegment(self.connection, (self.table_src,), "id", "comment")
+        b = TableSegment(self.connection, (self.table_dst,), "id", "comment")
+
+        differ = TableDiffer()
+        diff = list(differ.diff_tables(a, b))
+        uuid = diff[0][1][0]
+        self.assertEqual(diff, [("-", (uuid, "This one is different"))])