Merge pull request datafold#757 from pik94/prevent-type-overflow

Prevent type overflow

Author: dlawin

data_diff/databases/base.py

@@ -199,10 +199,19 @@ def apply_query(callback: Callable[[str], Any], sql_code: Union[str, ThreadLocal
 class BaseDialect(abc.ABC):
  SUPPORTS_PRIMARY_KEY: ClassVar[bool] = False
  SUPPORTS_INDEXES: ClassVar[bool] = False
+ PREVENT_OVERFLOW_WHEN_CONCAT: ClassVar[bool] = False
  TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {}
 
  PLACEHOLDER_TABLE = None # Used for Oracle
 
+ # Some database do not support long string so concatenation might lead to type overflow
+
+ _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)
@@ -392,10 +401,19 @@ 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.md5_as_hex(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]
@@ -769,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 md5_as_hex(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.
@@ -885,13 +907,16 @@ class Database(abc.ABC):
  Instanciated using :meth:`~data_diff.connect`
  """
 
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = BaseDialect
+
  SUPPORTS_ALPHANUMS: ClassVar[bool] = True
  SUPPORTS_UNIQUE_CONSTAINT: ClassVar[bool] = False
  CONNECT_URI_KWPARAMS: ClassVar[List[str]] = []
 
  default_schema: Optional[str] = None
  _interactive: bool = False
  is_closed: bool = False
+ _dialect: BaseDialect = None
 
  @property
  def name(self):
@@ -1120,10 +1145,13 @@ def close(self):
  return super().close()
 
  @property
- @abstractmethod
  def dialect(self) -> BaseDialect:
  "The dialect of the database. Used internally by Database, and also available publicly."
 
+ if not self._dialect:
+ self._dialect = self.DIALECT_CLASS()
+ return self._dialect
+
  @property
  @abstractmethod
  def CONNECT_URI_HELP(self) -> str:

data_diff/databases/bigquery.py

@@ -1,5 +1,5 @@
 import re
-from typing import Any, List, Union
+from typing import Any, ClassVar, List, Union, Type
 
 import attrs
 
@@ -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 md5_as_hex(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))"
@@ -179,9 +182,9 @@ def normalize_struct(self, value: str, _coltype: Struct) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class BigQuery(Database):
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  CONNECT_URI_HELP = "bigquery://<project>/<dataset>"
  CONNECT_URI_PARAMS = ["dataset"]
- dialect = Dialect()
 
  project: str
  dataset: str

data_diff/databases/clickhouse.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional, Type
+from typing import Any, ClassVar, Dict, Optional, Type
 
 import attrs
 
@@ -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 md5_as_hex(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:
@@ -164,7 +167,7 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class Clickhouse(ThreadedDatabase):
- dialect = Dialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  CONNECT_URI_HELP = "clickhouse://<user>:<password>@<host>/<database>"
  CONNECT_URI_PARAMS = ["database?"]
 

data_diff/databases/databricks.py

@@ -1,5 +1,5 @@
 import math
-from typing import Any, Dict, Sequence
+from typing import Any, ClassVar, Dict, Sequence, Type
 import logging
 
 import attrs
@@ -82,6 +82,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 md5_as_hex(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"""
 
@@ -104,7 +107,7 @@ def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class Databricks(ThreadedDatabase):
- dialect = Dialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  CONNECT_URI_HELP = "databricks://:<access_token>@<server_hostname>/<http_path>"
  CONNECT_URI_PARAMS = ["catalog", "schema"]
 

data_diff/databases/duckdb.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Union
+from typing import Any, ClassVar, Dict, Union, Type
 
 import attrs
 
@@ -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 md5_as_hex(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:
@@ -116,7 +119,7 @@ def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class DuckDB(Database):
- dialect = Dialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  SUPPORTS_UNIQUE_CONSTAINT = False # Temporary, until we implement it
  CONNECT_URI_HELP = "duckdb://<dbname>@<filepath>"
  CONNECT_URI_PARAMS = ["database", "dbpath"]

data_diff/databases/mssql.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, ClassVar, Dict, Optional, Type
 
 import attrs
 
@@ -38,7 +38,7 @@ def import_mssql():
 class Dialect(BaseDialect):
  name = "MsSQL"
  ROUNDS_ON_PREC_LOSS = True
- SUPPORTS_PRIMARY_KEY = True
+ SUPPORTS_PRIMARY_KEY: ClassVar[bool] = True
  SUPPORTS_INDEXES = True
  TYPE_CLASSES = {
  # Timestamps
@@ -151,10 +151,13 @@ 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 md5_as_hex(self, s: str) -> str:
+ return f"HashBytes('MD5', {s})"
+
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class MsSQL(ThreadedDatabase):
- dialect = Dialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  CONNECT_URI_HELP = "mssql://<user>:<password>@<host>/<database>/<schema>"
  CONNECT_URI_PARAMS = ["database", "schema"]
 

data_diff/databases/mysql.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, ClassVar, Dict, Type
 
 import attrs
 
@@ -40,7 +40,7 @@ def import_mysql():
 class Dialect(BaseDialect):
  name = "MySQL"
  ROUNDS_ON_PREC_LOSS = True
- SUPPORTS_PRIMARY_KEY = True
+ SUPPORTS_PRIMARY_KEY: ClassVar[bool] = True
  SUPPORTS_INDEXES = True
  TYPE_CLASSES = {
  # Dates
@@ -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 md5_as_hex(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))")
@@ -117,7 +120,7 @@ def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class MySQL(ThreadedDatabase):
- dialect = Dialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  SUPPORTS_ALPHANUMS = False
  SUPPORTS_UNIQUE_CONSTAINT = True
  CONNECT_URI_HELP = "mysql://<user>:<password>@<host>/<database>"

data_diff/databases/oracle.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, ClassVar, Dict, List, Optional, Type
 
 import attrs
 
@@ -43,7 +43,7 @@ class Dialect(
  BaseDialect,
 ):
  name = "Oracle"
- SUPPORTS_PRIMARY_KEY = True
+ SUPPORTS_PRIMARY_KEY: ClassVar[bool] = True
  SUPPORTS_INDEXES = True
  TYPE_CLASSES: Dict[str, type] = {
  "NUMBER": Decimal,
@@ -137,6 +137,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 md5_as_hex(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))"
@@ -161,7 +164,7 @@ def normalize_number(self, value: str, coltype: FractionalType) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class Oracle(ThreadedDatabase):
- dialect = Dialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  CONNECT_URI_HELP = "oracle://<user>:<password>@<host>/<database>"
  CONNECT_URI_PARAMS = ["database?"]
 

data_diff/databases/postgresql.py

@@ -42,7 +42,7 @@ def import_postgresql():
 class PostgresqlDialect(BaseDialect):
  name = "PostgreSQL"
  ROUNDS_ON_PREC_LOSS = True
- SUPPORTS_PRIMARY_KEY = True
+ SUPPORTS_PRIMARY_KEY: ClassVar[bool] = True
  SUPPORTS_INDEXES = True
 
  TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {
@@ -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 md5_as_hex(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')"
@@ -119,7 +122,7 @@ def normalize_json(self, value: str, _coltype: JSON) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class PostgreSQL(ThreadedDatabase):
- dialect = PostgresqlDialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = PostgresqlDialect
  SUPPORTS_UNIQUE_CONSTAINT = True
  CONNECT_URI_HELP = "postgresql://<user>:<password>@<host>/<database>"
  CONNECT_URI_PARAMS = ["database?"]

data_diff/databases/presto.py

@@ -1,6 +1,6 @@
 from functools import partial
 import re
-from typing import Any
+from typing import Any, ClassVar, Type
 
 import attrs
 
@@ -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 md5_as_hex(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))"
@@ -150,7 +153,7 @@ def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
 
 @attrs.define(frozen=False, init=False, kw_only=True)
 class Presto(Database):
- dialect = Dialect()
+ DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
  CONNECT_URI_HELP = "presto://<user>@<host>/<catalog>/<schema>"
  CONNECT_URI_PARAMS = ["catalog", "schema"]