Merge pull request #751 from datafold/squash-specialised-mixins

Simplify: Squash database-specialised mixins into their database-specialised dialects

Author: nolar

data_diff/databases/bigquery.py

@@ -40,7 +40,7 @@
     CHECKSUM_HEXDIGITS,
     MD5_HEXDIGITS,
 )
-from data_diff.databases.base import TIMESTAMP_PRECISION_POS, ThreadLocalInterpreter, Mixin_RandomSample
+from data_diff.databases.base import TIMESTAMP_PRECISION_POS, ThreadLocalInterpreter
 
 
 @import_helper(text="Please install BigQuery and configure your google-cloud access.")
@@ -63,13 +63,87 @@ def import_bigquery_service_account_impersonation():
 
 
 @attrs.define(frozen=False)
-class Mixin_MD5(AbstractMixin_MD5):
+class Dialect(
+    BaseDialect, AbstractMixin_Schema, AbstractMixin_MD5, AbstractMixin_NormalizeValue, AbstractMixin_TimeTravel
+):
+    name = "BigQuery"
+    ROUNDS_ON_PREC_LOSS = False  # Technically BigQuery doesn't allow implicit rounding or truncation
+    TYPE_CLASSES = {
+        # Dates
+        "TIMESTAMP": Timestamp,
+        "DATETIME": Datetime,
+        # Numbers
+        "INT64": Integer,
+        "INT32": Integer,
+        "NUMERIC": Decimal,
+        "BIGNUMERIC": Decimal,
+        "FLOAT64": Float,
+        "FLOAT32": Float,
+        "STRING": Text,
+        "BOOL": Boolean,
+        "JSON": JSON,
+    }
+    TYPE_ARRAY_RE = re.compile(r"ARRAY<(.+)>")
+    TYPE_STRUCT_RE = re.compile(r"STRUCT<(.+)>")
+
+    def random(self) -> str:
+        return "RAND()"
+
+    def quote(self, s: str):
+        return f"`{s}`"
+
+    def to_string(self, s: str):
+        return f"cast({s} as string)"
+
+    def type_repr(self, t) -> str:
+        try:
+            return {str: "STRING", float: "FLOAT64"}[t]
+        except KeyError:
+            return super().type_repr(t)
+
+    def parse_type(
+        self,
+        table_path: DbPath,
+        col_name: str,
+        type_repr: str,
+        *args: Any,  # pass-through args
+        **kwargs: Any,  # pass-through args
+    ) -> ColType:
+        col_type = super().parse_type(table_path, col_name, type_repr, *args, **kwargs)
+        if isinstance(col_type, UnknownColType):
+            m = self.TYPE_ARRAY_RE.fullmatch(type_repr)
+            if m:
+                item_type = self.parse_type(table_path, col_name, m.group(1), *args, **kwargs)
+                col_type = Array(item_type=item_type)
+
+            # We currently ignore structs' structure, but later can parse it too. Examples:
+            # - STRUCT<INT64, STRING(10)> (unnamed)
+            # - STRUCT<foo INT64, bar STRING(10)> (named)
+            # - STRUCT<foo INT64, bar ARRAY<INT64>> (with complex fields)
+            # - STRUCT<foo INT64, bar STRUCT<a INT64, b INT64>> (nested)
+            m = self.TYPE_STRUCT_RE.fullmatch(type_repr)
+            if m:
+                col_type = Struct()
+
+        return col_type
+
+    def to_comparable(self, value: str, coltype: ColType) -> str:
+        """Ensure that the expression is comparable in ``IS DISTINCT FROM``."""
+        if isinstance(coltype, (JSON, Array, Struct)):
+            return self.normalize_value_by_type(value, coltype)
+        else:
+            return super().to_comparable(value, coltype)
+
+    def set_timezone_to_utc(self) -> str:
+        raise NotImplementedError()
+
+    def parse_table_name(self, name: str) -> DbPath:
+        path = parse_table_name(name)
+        return tuple(i for i in path if i is not None)
+
     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}"
 
-
-@attrs.define(frozen=False)
-class Mixin_NormalizeValue(AbstractMixin_NormalizeValue):
     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))"
@@ -112,9 +186,6 @@ def normalize_struct(self, value: str, _coltype: Struct) -> str:
         # match on both sides: i.e. have properly ordered keys, same spacing, same quotes, etc.
         return f"to_json_string({value})"
 
-
-@attrs.define(frozen=False)
-class Mixin_Schema(AbstractMixin_Schema):
     def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable:
         return (
             table(table_schema, "INFORMATION_SCHEMA", "TABLES")
@@ -126,9 +197,6 @@ def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable:
             .select(this.table_name)
         )
 
-
-@attrs.define(frozen=False)
-class Mixin_TimeTravel(AbstractMixin_TimeTravel):
     def time_travel(
         self,
         table: Compilable,
@@ -155,86 +223,6 @@ def time_travel(
         )
 
 
-@attrs.define(frozen=False)
-class Dialect(
-    BaseDialect, Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, AbstractMixin_MD5, AbstractMixin_NormalizeValue
-):
-    name = "BigQuery"
-    ROUNDS_ON_PREC_LOSS = False  # Technically BigQuery doesn't allow implicit rounding or truncation
-    TYPE_CLASSES = {
-        # Dates
-        "TIMESTAMP": Timestamp,
-        "DATETIME": Datetime,
-        # Numbers
-        "INT64": Integer,
-        "INT32": Integer,
-        "NUMERIC": Decimal,
-        "BIGNUMERIC": Decimal,
-        "FLOAT64": Float,
-        "FLOAT32": Float,
-        "STRING": Text,
-        "BOOL": Boolean,
-        "JSON": JSON,
-    }
-    TYPE_ARRAY_RE = re.compile(r"ARRAY<(.+)>")
-    TYPE_STRUCT_RE = re.compile(r"STRUCT<(.+)>")
-
-    def random(self) -> str:
-        return "RAND()"
-
-    def quote(self, s: str):
-        return f"`{s}`"
-
-    def to_string(self, s: str):
-        return f"cast({s} as string)"
-
-    def type_repr(self, t) -> str:
-        try:
-            return {str: "STRING", float: "FLOAT64"}[t]
-        except KeyError:
-            return super().type_repr(t)
-
-    def parse_type(
-        self,
-        table_path: DbPath,
-        col_name: str,
-        type_repr: str,
-        *args: Any,  # pass-through args
-        **kwargs: Any,  # pass-through args
-    ) -> ColType:
-        col_type = super().parse_type(table_path, col_name, type_repr, *args, **kwargs)
-        if isinstance(col_type, UnknownColType):
-            m = self.TYPE_ARRAY_RE.fullmatch(type_repr)
-            if m:
-                item_type = self.parse_type(table_path, col_name, m.group(1), *args, **kwargs)
-                col_type = Array(item_type=item_type)
-
-            # We currently ignore structs' structure, but later can parse it too. Examples:
-            # - STRUCT<INT64, STRING(10)> (unnamed)
-            # - STRUCT<foo INT64, bar STRING(10)> (named)
-            # - STRUCT<foo INT64, bar ARRAY<INT64>> (with complex fields)
-            # - STRUCT<foo INT64, bar STRUCT<a INT64, b INT64>> (nested)
-            m = self.TYPE_STRUCT_RE.fullmatch(type_repr)
-            if m:
-                col_type = Struct()
-
-        return col_type
-
-    def to_comparable(self, value: str, coltype: ColType) -> str:
-        """Ensure that the expression is comparable in ``IS DISTINCT FROM``."""
-        if isinstance(coltype, (JSON, Array, Struct)):
-            return self.normalize_value_by_type(value, coltype)
-        else:
-            return super().to_comparable(value, coltype)
-
-    def set_timezone_to_utc(self) -> str:
-        raise NotImplementedError()
-
-    def parse_table_name(self, name: str) -> DbPath:
-        path = parse_table_name(name)
-        return tuple(i for i in path if i is not None)
-
-
 @attrs.define(frozen=False, init=False, kw_only=True)
 class BigQuery(Database):
     CONNECT_URI_HELP = "bigquery://<project>/<dataset>"

data_diff/databases/clickhouse.py

@@ -11,7 +11,6 @@
     ThreadedDatabase,
     import_helper,
     ConnectError,
-    Mixin_RandomSample,
 )
 from data_diff.abcs.database_types import (
     ColType,
@@ -39,16 +38,74 @@ def import_clickhouse():
 
 
 @attrs.define(frozen=False)
-class Mixin_MD5(AbstractMixin_MD5):
+class Dialect(BaseDialect, AbstractMixin_MD5, AbstractMixin_NormalizeValue):
+    name = "Clickhouse"
+    ROUNDS_ON_PREC_LOSS = False
+    TYPE_CLASSES = {
+        "Int8": Integer,
+        "Int16": Integer,
+        "Int32": Integer,
+        "Int64": Integer,
+        "Int128": Integer,
+        "Int256": Integer,
+        "UInt8": Integer,
+        "UInt16": Integer,
+        "UInt32": Integer,
+        "UInt64": Integer,
+        "UInt128": Integer,
+        "UInt256": Integer,
+        "Float32": Float,
+        "Float64": Float,
+        "Decimal": Decimal,
+        "UUID": Native_UUID,
+        "String": Text,
+        "FixedString": Text,
+        "DateTime": Timestamp,
+        "DateTime64": Timestamp,
+        "Bool": Boolean,
+    }
+
+    def quote(self, s: str) -> str:
+        return f'"{s}"'
+
+    def to_string(self, s: str) -> str:
+        return f"toString({s})"
+
+    def _convert_db_precision_to_digits(self, p: int) -> int:
+        # Done the same as for PostgreSQL but need to rewrite in another way
+        # because it does not help for float with a big integer part.
+        return super()._convert_db_precision_to_digits(p) - 2
+
+    def _parse_type_repr(self, type_repr: str) -> Optional[Type[ColType]]:
+        nullable_prefix = "Nullable("
+        if type_repr.startswith(nullable_prefix):
+            type_repr = type_repr[len(nullable_prefix) :].rstrip(")")
+
+        if type_repr.startswith("Decimal"):
+            type_repr = "Decimal"
+        elif type_repr.startswith("FixedString"):
+            type_repr = "FixedString"
+        elif type_repr.startswith("DateTime64"):
+            type_repr = "DateTime64"
+
+        return self.TYPE_CLASSES.get(type_repr)
+
+    # def timestamp_value(self, t: DbTime) -> str:
+    #     # return f"'{t}'"
+    #     return f"'{str(t)[:19]}'"
+
+    def set_timezone_to_utc(self) -> str:
+        raise NotImplementedError()
+
+    def current_timestamp(self) -> str:
+        return "now()"
+
     def md5_as_int(self, s: str) -> str:
         substr_idx = 1 + MD5_HEXDIGITS - CHECKSUM_HEXDIGITS
         return (
             f"reinterpretAsUInt128(reverse(unhex(lowerUTF8(substr(hex(MD5({s})), {substr_idx}))))) - {CHECKSUM_OFFSET}"
         )
 
-
-@attrs.define(frozen=False)
-class Mixin_NormalizeValue(AbstractMixin_NormalizeValue):
     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:
@@ -106,70 +163,6 @@ def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         return f"rpad({value}, {TIMESTAMP_PRECISION_POS + 6}, '0')"
 
 
-@attrs.define(frozen=False)
-class Dialect(BaseDialect, Mixin_MD5, Mixin_NormalizeValue, AbstractMixin_MD5, AbstractMixin_NormalizeValue):
-    name = "Clickhouse"
-    ROUNDS_ON_PREC_LOSS = False
-    TYPE_CLASSES = {
-        "Int8": Integer,
-        "Int16": Integer,
-        "Int32": Integer,
-        "Int64": Integer,
-        "Int128": Integer,
-        "Int256": Integer,
-        "UInt8": Integer,
-        "UInt16": Integer,
-        "UInt32": Integer,
-        "UInt64": Integer,
-        "UInt128": Integer,
-        "UInt256": Integer,
-        "Float32": Float,
-        "Float64": Float,
-        "Decimal": Decimal,
-        "UUID": Native_UUID,
-        "String": Text,
-        "FixedString": Text,
-        "DateTime": Timestamp,
-        "DateTime64": Timestamp,
-        "Bool": Boolean,
-    }
-
-    def quote(self, s: str) -> str:
-        return f'"{s}"'
-
-    def to_string(self, s: str) -> str:
-        return f"toString({s})"
-
-    def _convert_db_precision_to_digits(self, p: int) -> int:
-        # Done the same as for PostgreSQL but need to rewrite in another way
-        # because it does not help for float with a big integer part.
-        return super()._convert_db_precision_to_digits(p) - 2
-
-    def _parse_type_repr(self, type_repr: str) -> Optional[Type[ColType]]:
-        nullable_prefix = "Nullable("
-        if type_repr.startswith(nullable_prefix):
-            type_repr = type_repr[len(nullable_prefix) :].rstrip(")")
-
-        if type_repr.startswith("Decimal"):
-            type_repr = "Decimal"
-        elif type_repr.startswith("FixedString"):
-            type_repr = "FixedString"
-        elif type_repr.startswith("DateTime64"):
-            type_repr = "DateTime64"
-
-        return self.TYPE_CLASSES.get(type_repr)
-
-    # def timestamp_value(self, t: DbTime) -> str:
-    #     # return f"'{t}'"
-    #     return f"'{str(t)[:19]}'"
-
-    def set_timezone_to_utc(self) -> str:
-        raise NotImplementedError()
-
-    def current_timestamp(self) -> str:
-        return "now()"
-
-
 @attrs.define(frozen=False, init=False, kw_only=True)
 class Clickhouse(ThreadedDatabase):
     dialect = Dialect()