Fix for snowflake, mysql, presto, bigquery

Author: erezsh

data_diff/databases/base.py

@@ -231,6 +231,15 @@ def _normalize_table_path(self, path: DbPath) -> DbPath:
     def parse_table_name(self, name: str) -> DbPath:
         return parse_table_name(name)
 
+    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}"
+
+    def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
+        return f"TRIM({value})"
+
 
 class ThreadedDatabase(Database):
     """Access the database through singleton threads.
@@ -267,14 +276,6 @@ 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}"
-
-    def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
-        return f"TRIM({value})"
 
 
 CHECKSUM_HEXDIGITS = 15  # Must be 15 or lower

data_diff/databases/bigquery.py

@@ -22,6 +22,8 @@ class BigQuery(Database):
         "BIGNUMERIC": Decimal,
         "FLOAT64": Float,
         "FLOAT32": Float,
+        # Text
+        "STRING": Text,
     }
     ROUNDS_ON_PREC_LOSS = False  # Technically BigQuery doesn't allow implicit rounding or truncation
 

data_diff/databases/database_types.py

@@ -1,3 +1,4 @@
+import decimal
 from abc import ABC, abstractmethod
 from typing import Sequence, Optional, Tuple, Union, Dict, Any
 from datetime import datetime
@@ -54,7 +55,12 @@ class Float(FractionalType):
 
 
 class Decimal(FractionalType):
-    pass
+    @property
+    def python_type(self) -> type:
+        if self.precision == 0:
+            return int
+        return decimal.Decimal
+
 
 
 class StringType(ColType):

data_diff/databases/mysql.py

@@ -22,6 +22,7 @@ class MySQL(ThreadedDatabase):
         "int": Integer,
         # Text
         "varchar": Text,
+        "char": Text,
     }
     ROUNDS_ON_PREC_LOSS = True
 

data_diff/databases/oracle.py

@@ -70,6 +70,7 @@ def normalize_number(self, value: str, coltype: FractionalType) -> str:
 
     def _parse_type(
         self,
+        table_name: DbPath,
         col_name: str,
         type_repr: str,
         datetime_precision: int = None,

data_diff/databases/presto.py

@@ -29,6 +29,8 @@ class Presto(Database):
         "integer": Integer,
         "real": Float,
         "double": Float,
+        # Text
+        "varchar": Text,
     }
     ROUNDS_ON_PREC_LOSS = True
 
@@ -108,4 +110,17 @@ def _parse_type(
                 prec, scale = map(int, m.groups())
                 return n_cls(scale)
 
+        string_regexps = {
+            r"varchar\((\d+)\)": Text,
+            r"char\((\d+)\)": Text
+        }
+        for regexp, n_cls in string_regexps.items():
+            m = re.match(regexp + "$", type_repr)
+            if m:
+                return n_cls()
+
         return super()._parse_type(table_path, col_name, type_repr, datetime_precision, numeric_precision)
+
+    def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
+        # Trim doesn't work on CHAR type
+        return f"TRIM(CAST({value} AS VARCHAR))"

data_diff/databases/snowflake.py

@@ -20,6 +20,8 @@ class Snowflake(Database):
         # Numbers
         "NUMBER": Decimal,
         "FLOAT": Float,
+        # Text
+        "TEXT": Text,
     }
     ROUNDS_ON_PREC_LOSS = False
 

data_diff/diff_tables.py

@@ -5,7 +5,7 @@
 import time
 from operator import attrgetter, methodcaller
 from collections import defaultdict
-from typing import List, Tuple, Iterator, Optional
+from typing import List, Tuple, Iterator, Optional, Type
 import logging
 from concurrent.futures import ThreadPoolExecutor
 
@@ -18,6 +18,7 @@
     ArithUUID,
     NumericType,
     PrecisionType,
+    StringType,
     UnknownColType,
     Schema,
     Schema_CaseInsensitive,
@@ -295,7 +296,8 @@ def diff_tables(self, table1: TableSegment, table2: TableSegment) -> DiffResult:
         mins, maxs = zip(*key_ranges)
 
         key_type = table1._schema["id"]
-        assert type(key_type) == type(table2._schema["id"])
+        key_type2 = table2._schema["id"]
+        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))
@@ -324,7 +326,7 @@ def _validate_and_adjust_columns(self, table1, table2):
             col2 = table2._schema[c]
             if isinstance(col1, PrecisionType):
                 if not isinstance(col2, PrecisionType):
-                    raise TypeError(f"Incompatible types for column {c}:  {col1} <-> {col2}")
+                    raise TypeError(f"Incompatible types for column '{c}':  {col1} <-> {col2}")
 
                 lowest = min(col1, col2, key=attrgetter("precision"))
 
@@ -336,7 +338,7 @@ def _validate_and_adjust_columns(self, table1, table2):
 
             elif isinstance(col1, NumericType):
                 if not isinstance(col2, NumericType):
-                    raise TypeError(f"Incompatible types for column {c}:  {col1} <-> {col2}")
+                    raise TypeError(f"Incompatible types for column '{c}':  {col1} <-> {col2}")
 
                 lowest = min(col1, col2, key=attrgetter("precision"))
 
@@ -346,6 +348,10 @@ def _validate_and_adjust_columns(self, table1, table2):
                 table1._schema[c] = col1.replace(precision=lowest.precision)
                 table2._schema[c] = col2.replace(precision=lowest.precision)
 
+            elif isinstance(col1, StringType):
+                if not isinstance(col2, StringType):
+                    raise TypeError(f"Incompatible types for column '{c}':  {col1} <-> {col2}")
+
         for t in [table1, table2]:
             for c in t._relevant_columns:
                 ctype = t._schema[c]

tests/test_database_types.py

@@ -245,7 +245,7 @@ def __iter__(self):
             "bignumeric",
         ],
         "uuid": [
-            "$uuid"
+            "STRING",
         ],
     },
     db.Snowflake: {
@@ -272,7 +272,8 @@ def __iter__(self):
             "numeric",
         ],
         "uuid": [
-            "$uuid"
+            "varchar",
+            "varchar(100)",
         ],
 
     },
@@ -331,7 +332,8 @@ def __iter__(self):
             "decimal(30,6)",
         ],
         "uuid": [
-            "$uuid"
+            "varchar",
+            "char(100)",
         ],
 
     },
@@ -435,7 +437,8 @@ def _drop_table_if_exists(conn, table):
             conn.query(f"DROP TABLE {table}", None)
         else:
             conn.query(f"DROP TABLE IF EXISTS {table}", None)
-            conn.query("COMMIT", None)
+            if not isinstance(conn, db.BigQuery):
+                conn.query("COMMIT", None)
 
 
 class TestDiffCrossDatabaseTables(unittest.TestCase):