feat: prevent type overflow when long string concatenating

Author: pik94

data_diff/databases/base.py

@@ -204,6 +204,15 @@ class BaseDialect(abc.ABC):
     PLACEHOLDER_TABLE = None  # Used for Oracle
     USE_TOP_INSTEAD_LIMIT: bool = False  # True for MsSQL or Teradata
 
+    # Some database do not support long string so concatenation might lead to type overflow
+    PREVENT_OVERFLOW_WHEN_CONCAT = False
+
+    _prevent_overflow_when_concat: bool = False
+
+    def enable_preventing_type_overflow(self) -> None:
+        logger.info("Preventing type overflow when concatenation is enabled")
+        self._prevent_overflow_when_concat = True
+
     def parse_table_name(self, name: str) -> DbPath:
         "Parse the given table name into a DbPath"
         return parse_table_name(name)
@@ -393,10 +402,18 @@ def render_checksum(self, c: Compiler, elem: Checksum) -> str:
         return f"sum({md5})"
 
     def render_concat(self, c: Compiler, elem: Concat) -> str:
+        if self._prevent_overflow_when_concat:
+            items = [
+                f"{self.compile(c, Code(self.to_md5(self.to_string(self.compile(c, expr)))))}" for expr in elem.exprs
+            ]
+
         # We coalesce because on some DBs (e.g. MySQL) concat('a', NULL) is NULL
-        items = [
-            f"coalesce({self.compile(c, Code(self.to_string(self.compile(c, expr))))}, '<null>')" for expr in elem.exprs
-        ]
+        else:
+            items = [
+                f"coalesce({self.compile(c, Code(self.to_string(self.compile(c, expr))))}, '<null>')"
+                for expr in elem.exprs
+            ]
+
         assert items
         if len(items) == 1:
             return items[0]
@@ -770,6 +787,10 @@ def set_timezone_to_utc(self) -> str:
     def md5_as_int(self, s: str) -> str:
         "Provide SQL for computing md5 and returning an int"
 
+    @abstractmethod
+    def to_md5(self, s: str) -> str:
+        """Method to calculate MD5"""
+
     @abstractmethod
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         """Creates an SQL expression, that converts 'value' to a normalized timestamp.

data_diff/databases/bigquery.py

@@ -134,6 +134,9 @@ def parse_table_name(self, name: str) -> DbPath:
     def md5_as_int(self, s: str) -> str:
         return f"cast(cast( ('0x' || substr(TO_HEX(md5({s})), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS})) as int64) as numeric) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"md5({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             timestamp = f"timestamp_micros(cast(round(unix_micros(cast({value} as timestamp))/1000000, {coltype.precision})*1000000 as int))"

data_diff/databases/clickhouse.py

@@ -105,6 +105,9 @@ def md5_as_int(self, s: str) -> str:
             f"reinterpretAsUInt128(reverse(unhex(lowerUTF8(substr(hex(MD5({s})), {substr_idx}))))) - {CHECKSUM_OFFSET}"
         )
 
+    def to_md5(self, s: str) -> str:
+        return f"hex(MD5({s}))"
+
     def normalize_number(self, value: str, coltype: FractionalType) -> str:
         # If a decimal value has trailing zeros in a fractional part, when casting to string they are dropped.
         # For example:

data_diff/databases/databricks.py

@@ -76,6 +76,9 @@ def parse_table_name(self, name: str) -> DbPath:
     def md5_as_int(self, s: str) -> str:
         return f"cast(conv(substr(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16, 10) as decimal(38, 0)) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"md5({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         """Databricks timestamp contains no more than 6 digits in precision"""
 

data_diff/databases/duckdb.py

@@ -100,6 +100,9 @@ def current_timestamp(self) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"('0x' || SUBSTRING(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS},{CHECKSUM_HEXDIGITS}))::BIGINT - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"md5({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         # It's precision 6 by default. If precision is less than 6 -> we remove the trailing numbers.
         if coltype.rounds and coltype.precision > 0:

data_diff/databases/mssql.py

@@ -147,6 +147,9 @@ def normalize_number(self, value: str, coltype: NumericType) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"convert(bigint, convert(varbinary, '0x' + RIGHT(CONVERT(NVARCHAR(32), HashBytes('MD5', {s}), 2), {CHECKSUM_HEXDIGITS}), 1)) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"HashBytes('MD5', {s})"
+
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class MsSQL(ThreadedDatabase):

data_diff/databases/mysql.py

@@ -101,6 +101,9 @@ def set_timezone_to_utc(self) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"conv(substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16, 10) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"md5({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             return self.to_string(f"cast( cast({value} as datetime({coltype.precision})) as datetime(6))")

data_diff/databases/oracle.py

@@ -133,6 +133,9 @@ def md5_as_int(self, s: str) -> str:
         # TODO: Find a way to use UTL_RAW.CAST_TO_BINARY_INTEGER ?
         return f"to_number(substr(standard_hash({s}, 'MD5'), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 'xxxxxxxxxxxxxxx') - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"standard_hash({s}, 'MD5'"
+
     def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
         # Cast is necessary for correct MD5 (trimming not enough)
         return f"CAST(TRIM({value}) AS VARCHAR(36))"

data_diff/databases/postgresql.py

@@ -98,6 +98,9 @@ def type_repr(self, t) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"('x' || substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}))::bit({_CHECKSUM_BITSIZE})::bigint - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"md5({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             return f"to_char({value}::timestamp({coltype.precision}), 'YYYY-mm-dd HH24:MI:SS.US')"

data_diff/databases/presto.py

@@ -128,6 +128,9 @@ def current_timestamp(self) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"cast(from_base(substr(to_hex(md5(to_utf8({s}))), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16) as decimal(38, 0)) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"to_hex(md5(to_utf8({s})))"
+
     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/redshift.py

@@ -47,6 +47,9 @@ def type_repr(self, t) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"strtol(substring(md5({s}), {1+MD5_HEXDIGITS-CHECKSUM_HEXDIGITS}), 16)::decimal(38) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"md5({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             timestamp = f"{value}::timestamp(6)"

data_diff/databases/snowflake.py

@@ -76,6 +76,9 @@ def type_repr(self, t) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"BITAND(md5_number_lower64({s}), {CHECKSUM_MASK}) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"md5_number_lower64({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             timestamp = f"to_timestamp(round(date_part(epoch_nanosecond, convert_timezone('UTC', {value})::timestamp(9))/1000000000, {coltype.precision}))"

data_diff/databases/vertica.py

@@ -36,6 +36,7 @@ def import_vertica():
     return vertica_python
 
 
+@attrs.define(frozen=False)
 class Dialect(BaseDialect):
     name = "Vertica"
     ROUNDS_ON_PREC_LOSS = True
@@ -109,6 +110,9 @@ def current_timestamp(self) -> str:
     def md5_as_int(self, s: str) -> str:
         return f"CAST(HEX_TO_INTEGER(SUBSTRING(MD5({s}), {1 + MD5_HEXDIGITS - CHECKSUM_HEXDIGITS})) AS NUMERIC(38, 0)) - {CHECKSUM_OFFSET}"
 
+    def to_md5(self, s: str) -> str:
+        return f"MD5({s})"
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             return f"TO_CHAR({value}::TIMESTAMP({coltype.precision}), 'YYYY-MM-DD HH24:MI:SS.US')"

data_diff/diff_tables.py

@@ -208,6 +208,10 @@ def _diff_tables_wrapper(self, table1: TableSegment, table2: TableSegment, info_
             event_json = create_start_event_json(options)
             run_as_daemon(send_event_json, event_json)
 
+        if table1.database.dialect.PREVENT_OVERFLOW_WHEN_CONCAT or table2.database.dialect.PREVENT_OVERFLOW_WHEN_CONCAT:
+            table1.database.dialect.enable_preventing_type_overflow()
+            table2.database.dialect.enable_preventing_type_overflow()
+
         start = time.monotonic()
         error = None
         try:

tests/test_query.py

@@ -71,6 +71,9 @@ def optimizer_hints(self, s: str):
     def md5_as_int(self, s: str) -> str:
         raise NotImplementedError
 
+    def to_md5(self, s: str) -> str:
+        raise NotImplementedError
+
     def normalize_number(self, value: str, coltype: FractionalType) -> str:
         raise NotImplementedError