collate-data-diff 0.11.2__py3-none-any.whl

data_diff/databases/presto.py

@@ -0,0 +1,197 @@
+from functools import partial
+import re
+from typing import Any, ClassVar, Type
+
+import attrs
+
+from data_diff.schema import RawColumnInfo
+from data_diff.utils import match_regexps
+
+from data_diff.abcs.database_types import (
+    Timestamp,
+    TimestampTZ,
+    Integer,
+    Float,
+    Text,
+    FractionalType,
+    DbPath,
+    DbTime,
+    Decimal,
+    ColType,
+    ColType_UUID,
+    TemporalType,
+    Boolean,
+)
+from data_diff.databases.base import (
+    BaseDialect,
+    Database,
+    import_helper,
+    ThreadLocalInterpreter,
+)
+from data_diff.databases.base import (
+    MD5_HEXDIGITS,
+    CHECKSUM_HEXDIGITS,
+    CHECKSUM_OFFSET,
+    TIMESTAMP_PRECISION_POS,
+)
+
+
+def query_cursor(c, sql_code):
+    c.execute(sql_code)
+    if sql_code.lower().startswith("select"):
+        return c.fetchall()
+    # Required for the query to actually run 🤯
+    if re.match(r"(insert|create|truncate|drop|explain)", sql_code, re.IGNORECASE):
+        return c.fetchone()
+
+
+@import_helper("presto")
+def import_presto():
+    import prestodb
+
+    return prestodb
+
+
+class Dialect(BaseDialect):
+    name = "Presto"
+    ROUNDS_ON_PREC_LOSS = True
+    TYPE_CLASSES = {
+        # Timestamps
+        "timestamp with time zone": TimestampTZ,
+        "timestamp without time zone": Timestamp,
+        "timestamp": Timestamp,
+        # Numbers
+        "integer": Integer,
+        "bigint": Integer,
+        "real": Float,
+        "double": Float,
+        # Text
+        "varchar": Text,
+        # Boolean
+        "boolean": Boolean,
+    }
+
+    def explain_as_text(self, query: str) -> str:
+        return f"EXPLAIN (FORMAT TEXT) {query}"
+
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamp with time zone"
+
+        try:
+            return {float: "REAL"}[t]
+        except KeyError:
+            return super().type_repr(t)
+
+    def timestamp_value(self, t: DbTime) -> str:
+        return f"timestamp '{t.isoformat(' ')}'"
+
+    def quote(self, s: str):
+        return f'"{s}"'
+
+    def to_string(self, s: str):
+        return f"cast({s} as varchar)"
+
+    def parse_type(self, table_path: DbPath, info: RawColumnInfo) -> ColType:
+        timestamp_regexps = {
+            r"timestamp\((\d)\)": Timestamp,
+            r"timestamp\((\d)\) with time zone": TimestampTZ,
+        }
+        for m, t_cls in match_regexps(timestamp_regexps, info.data_type):
+            precision = int(m.group(1))
+            return t_cls(precision=precision, rounds=self.ROUNDS_ON_PREC_LOSS)
+
+        number_regexps = {r"decimal\((\d+),(\d+)\)": Decimal}
+        for m, n_cls in match_regexps(number_regexps, info.data_type):
+            _prec, scale = map(int, m.groups())
+            return n_cls(scale)
+
+        string_regexps = {r"varchar\((\d+)\)": Text, r"char\((\d+)\)": Text}
+        for m, n_cls in match_regexps(string_regexps, info.data_type):
+            return n_cls()
+
+        return super().parse_type(table_path, info)
+
+    def set_timezone_to_utc(self) -> str:
+        return "SET TIME ZONE '+00:00'"
+
+    def current_timestamp(self) -> str:
+        return "current_timestamp"
+
+    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))"
+
+    def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
+        # TODO rounds
+        if coltype.rounds:
+            s = f"date_format(cast({value} as timestamp(6)), '%Y-%m-%d %H:%i:%S.%f')"
+        else:
+            s = f"date_format(cast({value} as timestamp(6)), '%Y-%m-%d %H:%i:%S.%f')"
+
+        return f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS+coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS+6}, '0')"
+
+    def normalize_number(self, value: str, coltype: FractionalType) -> str:
+        return self.to_string(f"cast({value} as decimal(38,{coltype.precision}))")
+
+    def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
+        return self.to_string(f"cast ({value} as int)")
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class Presto(Database):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
+    CONNECT_URI_HELP = "presto://<user>@<host>/<catalog>/<schema>"
+    CONNECT_URI_PARAMS = ["catalog", "schema"]
+
+    _conn: Any
+
+    def __init__(self, **kw) -> None:
+        super().__init__()
+        self.default_schema = "public"
+        prestodb = import_presto()
+
+        if kw.get("schema"):
+            self.default_schema = kw.get("schema")
+
+        if kw.get("auth") == "basic":  # if auth=basic, add basic authenticator for Presto
+            kw["auth"] = prestodb.auth.BasicAuthentication(kw["user"], kw.pop("password"))
+
+        if "cert" in kw:  # if a certificate was specified in URI, verify session with cert
+            cert = kw.pop("cert")
+            self._conn = prestodb.dbapi.connect(**kw)
+            self._conn._http_session.verify = cert
+        else:
+            self._conn = prestodb.dbapi.connect(**kw)
+
+    def _query(self, sql_code: str) -> list:
+        "Uses the standard SQL cursor interface"
+        c = self._conn.cursor()
+
+        if isinstance(sql_code, ThreadLocalInterpreter):
+            return sql_code.apply_queries(partial(query_cursor, c))
+
+        return query_cursor(c, sql_code)
+
+    def close(self):
+        super().close()
+        self._conn.close()
+
+    def select_table_schema(self, path: DbPath) -> str:
+        schema, table = self._normalize_table_path(path)
+
+        return (
+            "SELECT column_name, data_type, 3 as datetime_precision, 3 as numeric_precision, NULL as numeric_scale "
+            "FROM INFORMATION_SCHEMA.COLUMNS "
+            f"WHERE table_name = '{table}' AND table_schema = '{schema}'"
+        )
+
+    @property
+    def is_autocommit(self) -> bool:
+        return False

data_diff/databases/redshift.py

@@ -0,0 +1,217 @@
+from typing import Any, ClassVar, Iterable, List, Dict, Tuple, Type
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    ColType,
+    Float,
+    JSON,
+    TemporalType,
+    FractionalType,
+    DbPath,
+    TimestampTZ,
+    Integer,
+)
+from data_diff.databases.postgresql import (
+    BaseDialect,
+    PostgreSQL,
+    MD5_HEXDIGITS,
+    CHECKSUM_HEXDIGITS,
+    CHECKSUM_OFFSET,
+    TIMESTAMP_PRECISION_POS,
+    PostgresqlDialect,
+)
+from data_diff.schema import RawColumnInfo
+
+
+@attrs.define(frozen=False)
+class Dialect(PostgresqlDialect):
+    name = "Redshift"
+    TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {
+        **PostgresqlDialect.TYPE_CLASSES,
+        "double": Float,
+        "real": Float,
+        "super": JSON,
+        "int": Integer,  # Redshift Spectrum
+        "float": Float,  # Redshift Spectrum
+    }
+    SUPPORTS_INDEXES = False
+
+    def concat(self, items: List[str]) -> str:
+        joined_exprs = " || ".join(items)
+        return f"({joined_exprs})"
+
+    def is_distinct_from(self, a: str, b: str) -> str:
+        return f"({a} IS NULL != {b} IS NULL) OR ({a}!={b})"
+
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamptz"
+        return super().type_repr(t)
+
+    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 md5_as_hex(self, s: str) -> str:
+        return f"md5({s})"
+
+    def normalize_number(self, value: str, coltype: FractionalType) -> str:
+        return self.to_string(f"{value}::decimal(38,{coltype.precision})")
+
+    def normalize_json(self, value: str, _coltype: JSON) -> str:
+        return f"nvl2({value}, json_serialize({value}), NULL)"
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class Redshift(PostgreSQL):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
+    CONNECT_URI_HELP = "redshift://<user>:<password>@<host>/<database>"
+    CONNECT_URI_PARAMS = ["database?"]
+
+    def select_table_schema(self, path: DbPath) -> str:
+        database, schema, table = self._normalize_table_path(path)
+
+        info_schema_path = ["information_schema", "columns"]
+        if database:
+            info_schema_path.insert(0, database)
+
+        return (
+            f"SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale FROM {'.'.join(info_schema_path)} "
+            f"WHERE table_name = '{table.lower()}' AND table_schema = '{schema.lower()}'"
+        )
+
+    def select_external_table_schema(self, path: DbPath) -> str:
+        database, schema, table = self._normalize_table_path(path)
+
+        db_clause = ""
+        if database:
+            db_clause = f" AND redshift_database_name = '{database.lower()}'"
+
+        return (
+            f"""SELECT
+                columnname AS column_name
+                , CASE WHEN external_type = 'string' THEN 'varchar' ELSE external_type END AS data_type
+                , NULL AS datetime_precision
+                , NULL AS numeric_precision
+                , NULL AS numeric_scale
+            FROM svv_external_columns
+                WHERE tablename = '{table.lower()}' AND schemaname = '{schema.lower()}'
+            """
+            + db_clause
+        )
+
+    def query_external_table_schema(self, path: DbPath) -> Dict[str, RawColumnInfo]:
+        rows = self.query(self.select_external_table_schema(path), list)
+        if not rows:
+            raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns")
+
+        schema_dict = self._normalize_schema_info(rows)
+        return schema_dict
+
+    def select_view_columns(self, path: DbPath) -> str:
+        _, schema, table = self._normalize_table_path(path)
+
+        return """select * from pg_get_cols('{}.{}')
+                cols(col_name name, col_type varchar)
+            """.format(schema, table)
+
+    def query_pg_get_cols(self, path: DbPath) -> Dict[str, RawColumnInfo]:
+        rows = self.query(self.select_view_columns(path), list)
+        if not rows:
+            raise RuntimeError(f"{self.name}: View '{'.'.join(path)}' does not exist, or has no columns")
+
+        schema_dict = self._normalize_schema_info(rows)
+        return schema_dict
+
+    def select_svv_columns_schema(self, path: DbPath) -> Dict[str, tuple]:
+        database, schema, table = self._normalize_table_path(path)
+
+        db_clause = ""
+        if database:
+            db_clause = f" AND table_catalog = '{database.lower()}'"
+
+        return (
+            f"""
+            select
+            	distinct
+                column_name,
+            	data_type,
+            	datetime_precision,
+            	numeric_precision,
+            	numeric_scale
+            from
+            	svv_columns
+            where table_name = '{table.lower()}' and table_schema = '{schema.lower()}'
+            """
+            + db_clause
+        )
+
+    def query_svv_columns(self, path: DbPath) -> Dict[str, RawColumnInfo]:
+        rows = self.query(self.select_svv_columns_schema(path), list)
+        if not rows:
+            raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns")
+
+        d = {
+            r[0]: RawColumnInfo(
+                column_name=r[0],
+                data_type=r[1],
+                datetime_precision=r[2],
+                numeric_precision=r[3],
+                numeric_scale=r[4],
+                collation_name=r[5] if len(r) > 5 else None,
+            )
+            for r in rows
+        }
+        assert len(d) == len(rows)
+        return d
+
+    # when using a non-information_schema source, strip (N) from type(N) etc. to match
+    # typical information_schema output
+    def _normalize_schema_info(self, rows: Iterable[Tuple[Any]]) -> Dict[str, RawColumnInfo]:
+        schema_dict: Dict[str, RawColumnInfo] = {}
+        for r in rows:
+            col_name = r[0]
+            type_info = r[1].split("(")
+            base_type = type_info[0]
+            precision = None
+            scale = None
+
+            if len(type_info) > 1:
+                if base_type == "numeric":
+                    precision, scale = type_info[1][:-1].split(",")
+                    precision = int(precision)
+                    scale = int(scale)
+
+            schema_dict[col_name] = RawColumnInfo(
+                column_name=col_name,
+                data_type=base_type,
+                datetime_precision=None,
+                numeric_precision=precision,
+                numeric_scale=scale,
+                collation_name=None,
+            )
+        return schema_dict
+
+    def query_table_schema(self, path: DbPath) -> Dict[str, RawColumnInfo]:
+        try:
+            return super().query_table_schema(path)
+        except RuntimeError:
+            try:
+                return self.query_external_table_schema(path)
+            except RuntimeError:
+                try:
+                    return self.query_pg_get_cols(path)
+                except Exception:
+                    return self.query_svv_columns(path)
+
+    def _normalize_table_path(self, path: DbPath) -> DbPath:
+        if len(path) == 1:
+            return None, self.default_schema, path[0]
+        elif len(path) == 2:
+            return None, path[0], path[1]
+        elif len(path) == 3:
+            return path
+
+        raise ValueError(
+            f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected format: table, schema.table, or database.schema.table"
+        )

data_diff/databases/snowflake.py

@@ -0,0 +1,207 @@
+import base64
+from typing import Any, ClassVar, Union, List, Type, Optional
+import logging
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    Timestamp,
+    TimestampTZ,
+    Decimal,
+    Float,
+    Text,
+    FractionalType,
+    TemporalType,
+    DbPath,
+    Boolean,
+    Date,
+    Time,
+)
+from data_diff.databases.base import (
+    BaseDialect,
+    ConnectError,
+    Database,
+    import_helper,
+    CHECKSUM_MASK,
+    ThreadLocalInterpreter,
+    CHECKSUM_OFFSET,
+)
+
+
+@import_helper("snowflake")
+def import_snowflake():
+    import snowflake.connector
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.backends import default_backend
+
+    return snowflake, serialization, default_backend
+
+
+class Dialect(BaseDialect):
+    name = "Snowflake"
+    ROUNDS_ON_PREC_LOSS = False
+    TYPE_CLASSES = {
+        # Timestamps
+        "TIMESTAMP_NTZ": Timestamp,
+        "TIMESTAMP_LTZ": Timestamp,
+        "TIMESTAMP_TZ": TimestampTZ,
+        "DATE": Date,
+        "TIME": Time,
+        # Numbers
+        "NUMBER": Decimal,
+        "FLOAT": Float,
+        # Text
+        "TEXT": Text,
+        # Boolean
+        "BOOLEAN": Boolean,
+    }
+
+    def explain_as_text(self, query: str) -> str:
+        return f"EXPLAIN USING TEXT {query}"
+
+    def quote(self, s: str):
+        return f'"{s}"'
+
+    def to_string(self, s: str):
+        return f"cast({s} as string)"
+
+    def set_timezone_to_utc(self) -> str:
+        return "ALTER SESSION SET TIMEZONE = 'UTC'"
+
+    def optimizer_hints(self, hints: str) -> str:
+        raise NotImplementedError("Optimizer hints not yet implemented in snowflake")
+
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamp_tz({t.precision})"
+        return super().type_repr(t)
+
+    def md5_as_int(self, s: str) -> str:
+        return f"BITAND(md5_number_lower64({s}), {CHECKSUM_MASK}) - {CHECKSUM_OFFSET}"
+
+    def md5_as_hex(self, s: str) -> str:
+        return f"md5({s})"
+
+    def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
+        try:
+            is_date = coltype.is_date
+            is_time = coltype.is_time
+        except:
+            is_date = False
+            is_time = False
+        if isinstance(coltype, Date) or is_date:
+            return f"({value}::varchar)"
+        elif isinstance(coltype, Time) or is_time:
+            microseconds = f"TIMEDIFF(microsecond, cast('00:00:00' as time), {value})"
+            rounded = f"round({microseconds}, -6 + {coltype.precision})"
+            time_value = f"TIMEADD(microsecond, {rounded}, cast('00:00:00' as time))"
+            converted = f"TO_VARCHAR({time_value}, 'HH24:MI:SS.FF6')"
+            return converted
+
+        if coltype.rounds:
+            timestamp = f"to_timestamp(round(date_part(epoch_nanosecond, convert_timezone('UTC', {value})::timestamp(9))/1000000000, {coltype.precision}))"
+        else:
+            timestamp = f"cast(convert_timezone('UTC', {value}) as timestamp({coltype.precision}))"
+
+        return f"to_char({timestamp}, 'YYYY-MM-DD HH24:MI:SS.FF6')"
+
+    def normalize_number(self, value: str, coltype: FractionalType) -> str:
+        return self.to_string(f"cast({value} as decimal(38, {coltype.precision}))")
+
+    def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
+        return self.to_string(f"{value}::int")
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class Snowflake(Database):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
+    CONNECT_URI_HELP = "snowflake://<user>:<password>@<account>/<database>/<SCHEMA>?warehouse=<WAREHOUSE>"
+    CONNECT_URI_PARAMS = ["database", "schema"]
+    CONNECT_URI_KWPARAMS = ["warehouse"]
+
+    _conn: Any
+
+    def __init__(self, *, schema: str, key: Optional[str] = None, key_content: Optional[str] = None, **kw) -> None:
+        super().__init__()
+        snowflake, serialization, default_backend = import_snowflake()
+        logging.getLogger("snowflake.connector").setLevel(logging.WARNING)
+
+        # Ignore the error: snowflake.connector.network.RetryRequest: could not find io module state
+        # It's a known issue: https://github.com/snowflakedb/snowflake-connector-python/issues/145
+        logging.getLogger("snowflake.connector.network").disabled = True
+
+        assert '"' not in schema, "Schema name should not contain quotes!"
+        if key_content and key:
+            raise ConnectError("Only key value or key file path can be specified, not both")
+
+        key_bytes = None
+        if key:
+            with open(key, "rb") as f:
+                key_bytes = f.read()
+        if key_content:
+            key_bytes = base64.b64decode(key_content)
+
+        # If a private key is used, read it from the specified path and pass it as "private_key" to the connector.
+        if key_bytes:
+            if "password" in kw:
+                raise ConnectError("Cannot use password and key at the same time")
+            if kw.get("private_key_passphrase"):
+                encoded_passphrase = kw.get("private_key_passphrase").encode()
+            else:
+                encoded_passphrase = None
+            p_key = serialization.load_pem_private_key(
+                key_bytes,
+                password=encoded_passphrase,
+                backend=default_backend(),
+            )
+
+            kw["private_key"] = p_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+
+        self._conn = snowflake.connector.connect(schema=f'"{schema}"', **kw)
+
+        self.default_schema = schema
+
+    def close(self):
+        super().close()
+        self._conn.close()
+
+    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]):
+        "Uses the standard SQL cursor interface"
+        return self._query_conn(self._conn, sql_code)
+
+    def select_table_schema(self, path: DbPath) -> str:
+        """Provide SQL for selecting the table schema as (name, type, date_prec, num_prec)"""
+        database, schema, name = self._normalize_table_path(path)
+        info_schema_path = ["information_schema", "columns"]
+        if database:
+            info_schema_path.insert(0, database)
+
+        return (
+            "SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale"
+            "   , coalesce(collation_name, 'utf8') "
+            f"FROM {'.'.join(info_schema_path)} "
+            f"WHERE table_name = '{name}' AND table_schema = '{schema}'"
+        )
+
+    def _normalize_table_path(self, path: DbPath) -> DbPath:
+        if len(path) == 1:
+            return None, self.default_schema, path[0]
+        elif len(path) == 2:
+            return None, path[0], path[1]
+        elif len(path) == 3:
+            return path
+
+        raise ValueError(
+            f"{self.name}: Bad table path for {self}: '{'.'.join(path)}'. Expected format: table, schema.table, or database.schema.table"
+        )
+
+    @property
+    def is_autocommit(self) -> bool:
+        return True
+
+    def query_table_unique_columns(self, path: DbPath) -> List[str]:
+        return []

data_diff/databases/trino.py

@@ -0,0 +1,50 @@
+from typing import Any, ClassVar, Type
+
+import attrs
+
+from data_diff.abcs.database_types import TemporalType, ColType_UUID
+from data_diff.databases import presto
+from data_diff.databases.base import import_helper
+from data_diff.databases.base import TIMESTAMP_PRECISION_POS, BaseDialect
+
+
+@import_helper("trino")
+def import_trino():
+    import trino
+
+    return trino
+
+
+class Dialect(presto.Dialect):
+    name = "Trino"
+
+    def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
+        if coltype.rounds:
+            s = f"date_format(cast({value} as timestamp({coltype.precision})), '%Y-%m-%d %H:%i:%S.%f')"
+        else:
+            s = f"date_format(cast({value} as timestamp(6)), '%Y-%m-%d %H:%i:%S.%f')"
+
+        return (
+            f"RPAD(RPAD({s}, {TIMESTAMP_PRECISION_POS + coltype.precision}, '.'), {TIMESTAMP_PRECISION_POS + 6}, '0')"
+        )
+
+    def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
+        return f"TRIM({value})"
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class Trino(presto.Presto):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
+    CONNECT_URI_HELP = "trino://<user>@<host>/<catalog>/<schema>"
+    CONNECT_URI_PARAMS = ["catalog", "schema"]
+
+    _conn: Any
+
+    def __init__(self, **kw) -> None:
+        super().__init__()
+        trino = import_trino()
+
+        if kw.get("schema"):
+            self.default_schema = kw.get("schema")
+
+        self._conn = trino.dbapi.connect(**kw)