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

data_diff/databases/mysql.py

@@ -0,0 +1,159 @@
+from typing import Any, ClassVar, Dict, Type, Union
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    Datetime,
+    Timestamp,
+    Float,
+    Decimal,
+    Integer,
+    Text,
+    TemporalType,
+    FractionalType,
+    ColType_UUID,
+    Boolean,
+    Date,
+)
+from data_diff.databases.base import (
+    ThreadedDatabase,
+    import_helper,
+    ConnectError,
+    BaseDialect,
+    ThreadLocalInterpreter,
+)
+from data_diff.databases.base import (
+    MD5_HEXDIGITS,
+    CHECKSUM_HEXDIGITS,
+    TIMESTAMP_PRECISION_POS,
+    CHECKSUM_OFFSET,
+)
+
+
+@import_helper("mysql")
+def import_mysql():
+    import mysql.connector
+
+    return mysql.connector
+
+
+@attrs.define(frozen=False)
+class Dialect(BaseDialect):
+    name = "MySQL"
+    ROUNDS_ON_PREC_LOSS = True
+    SUPPORTS_PRIMARY_KEY: ClassVar[bool] = True
+    SUPPORTS_INDEXES = True
+    TYPE_CLASSES = {
+        # Dates
+        "datetime": Datetime,
+        "timestamp": Timestamp,
+        "date": Date,
+        # Numbers
+        "double": Float,
+        "float": Float,
+        "decimal": Decimal,
+        "int": Integer,
+        "bigint": Integer,
+        "mediumint": Integer,
+        "smallint": Integer,
+        "tinyint": Integer,
+        # Text
+        "varchar": Text,
+        "char": Text,
+        "varbinary": Text,
+        "binary": Text,
+        "text": Text,
+        "mediumtext": Text,
+        "longtext": Text,
+        "tinytext": Text,
+        # Boolean
+        "boolean": Boolean,
+    }
+
+    def quote(self, s: str) -> str:
+        return f"`{s}`"
+
+    def to_string(self, s: str) -> str:
+        return f"cast({s} as char)"
+
+    def is_distinct_from(self, a: str, b: str) -> str:
+        return f"not ({a} <=> {b})"
+
+    def random(self) -> str:
+        return "RAND()"
+
+    def type_repr(self, t) -> str:
+        try:
+            return {
+                str: "VARCHAR(1024)",
+            }[t]
+        except KeyError:
+            return super().type_repr(t)
+
+    def explain_as_text(self, query: str) -> str:
+        return f"EXPLAIN FORMAT=TREE {query}"
+
+    def optimizer_hints(self, s: str):
+        return f"/*+ {s} */ "
+
+    def set_timezone_to_utc(self) -> str:
+        return "SET @@session.time_zone='+00:00'"
+
+    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))")
+
+        s = self.to_string(f"cast({value} as datetime(6))")
+        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_uuid(self, value: str, coltype: ColType_UUID) -> str:
+        return f"TRIM(CAST({value} AS char))"
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class MySQL(ThreadedDatabase):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
+    SUPPORTS_ALPHANUMS = False
+    SUPPORTS_UNIQUE_CONSTAINT = True
+    CONNECT_URI_HELP = "mysql://<user>:<password>@<host>/<database>"
+    CONNECT_URI_PARAMS = ["database?"]
+
+    _args: Dict[str, Any]
+
+    def __init__(self, *, thread_count, **kw) -> None:
+        super().__init__(thread_count=thread_count)
+        self._args = kw
+
+        # In MySQL schema and database are synonymous
+        try:
+            self.default_schema = kw["database"]
+        except KeyError:
+            raise ValueError("MySQL URL must specify a database")
+
+    def create_connection(self):
+        mysql = import_mysql()
+        try:
+            return mysql.connect(charset="utf8", use_unicode=True, **self._args)
+        except mysql.Error as e:
+            if e.errno == mysql.errorcode.ER_ACCESS_DENIED_ERROR:
+                raise ConnectError("Bad user name or password") from e
+            elif e.errno == mysql.errorcode.ER_BAD_DB_ERROR:
+                raise ConnectError("Database does not exist") from e
+            raise ConnectError(*e.args) from e
+
+    def _query_in_worker(self, sql_code: Union[str, ThreadLocalInterpreter]):
+        "This method runs in a worker thread"
+        if self._init_error:
+            raise self._init_error
+        if not self.thread_local.conn.is_connected():
+            self.thread_local.conn.ping(reconnect=True, attempts=3, delay=5)
+        return self._query_conn(self.thread_local.conn, sql_code)

data_diff/databases/oracle.py

@@ -0,0 +1,195 @@
+from typing import Any, ClassVar, Dict, List, Optional, Type
+
+import attrs
+
+from data_diff.schema import RawColumnInfo
+from data_diff.utils import match_regexps
+from data_diff.abcs.database_types import (
+    Decimal,
+    Float,
+    Text,
+    DbPath,
+    TemporalType,
+    ColType,
+    DbTime,
+    ColType_UUID,
+    Timestamp,
+    TimestampTZ,
+    FractionalType,
+)
+from data_diff.databases.base import (
+    BaseDialect,
+    ThreadedDatabase,
+    import_helper,
+    ConnectError,
+    QueryError,
+    CHECKSUM_OFFSET,
+    CHECKSUM_HEXDIGITS,
+    MD5_HEXDIGITS,
+)
+from data_diff.databases.base import TIMESTAMP_PRECISION_POS
+
+SESSION_TIME_ZONE = None  # Changed by the tests
+
+
+@import_helper("oracle")
+def import_oracle():
+    import oracledb
+
+    return oracledb
+
+
+@attrs.define(frozen=False)
+class Dialect(
+    BaseDialect,
+):
+    name = "Oracle"
+    SUPPORTS_PRIMARY_KEY: ClassVar[bool] = True
+    SUPPORTS_INDEXES = True
+    TYPE_CLASSES: Dict[str, type] = {
+        "NUMBER": Decimal,
+        "FLOAT": Float,
+        # Text
+        "CHAR": Text,
+        "NCHAR": Text,
+        "NVARCHAR2": Text,
+        "VARCHAR2": Text,
+        "DATE": Timestamp,
+    }
+    ROUNDS_ON_PREC_LOSS = True
+    PLACEHOLDER_TABLE = "DUAL"
+
+    def quote(self, s: str) -> str:
+        return f'"{s}"'
+
+    def to_string(self, s: str) -> str:
+        return f"cast({s} as varchar(1024))"
+
+    def limit_select(
+        self,
+        select_query: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        has_order_by: Optional[bool] = None,
+    ) -> str:
+        if offset:
+            raise NotImplementedError("No support for OFFSET in query")
+
+        return f"SELECT * FROM ({select_query}) FETCH NEXT {limit} ROWS ONLY"
+
+    def concat(self, items: List[str]) -> str:
+        joined_exprs = " || ".join(items)
+        return f"({joined_exprs})"
+
+    def timestamp_value(self, t: DbTime) -> str:
+        return "timestamp '%s'" % t.isoformat(" ")
+
+    def random(self) -> str:
+        return "dbms_random.value"
+
+    def is_distinct_from(self, a: str, b: str) -> str:
+        return f"DECODE({a}, {b}, 1, 0) = 0"
+
+    def type_repr(self, t) -> str:
+        try:
+            return {
+                str: "VARCHAR(1024)",
+            }[t]
+        except KeyError:
+            return super().type_repr(t)
+
+    def constant_values(self, rows) -> str:
+        return " UNION ALL ".join(
+            "SELECT %s FROM DUAL" % ", ".join(self._constant_value(v) for v in row) for row in rows
+        )
+
+    def explain_as_text(self, query: str) -> str:
+        raise NotImplementedError("Explain not yet implemented in Oracle")
+
+    def parse_type(self, table_path: DbPath, info: RawColumnInfo) -> ColType:
+        regexps = {
+            r"TIMESTAMP\((\d)\) WITH LOCAL TIME ZONE": Timestamp,
+            r"TIMESTAMP\((\d)\) WITH TIME ZONE": TimestampTZ,
+            r"TIMESTAMP\((\d)\)": Timestamp,
+        }
+
+        for m, t_cls in match_regexps(regexps, info.data_type):
+            precision = int(m.group(1))
+            return t_cls(precision=precision, rounds=self.ROUNDS_ON_PREC_LOSS)
+
+        return super().parse_type(table_path, info)
+
+    def set_timezone_to_utc(self) -> str:
+        return "ALTER SESSION SET TIME_ZONE = 'UTC'"
+
+    def current_timestamp(self) -> str:
+        return "LOCALTIMESTAMP"
+
+    def md5_as_int(self, s: str) -> str:
+        # standard_hash is faster than DBMS_CRYPTO.Hash
+        # 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))"
+
+    def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
+        if coltype.rounds:
+            return f"to_char(cast({value} as timestamp({coltype.precision})), 'YYYY-MM-DD HH24:MI:SS.FF6')"
+
+        if coltype.precision > 0:
+            truncated = f"to_char({value}, 'YYYY-MM-DD HH24:MI:SS.FF{coltype.precision}')"
+        else:
+            truncated = f"to_char({value}, 'YYYY-MM-DD HH24:MI:SS.')"
+        return f"RPAD({truncated}, {TIMESTAMP_PRECISION_POS+6}, '0')"
+
+    def normalize_number(self, value: str, coltype: FractionalType) -> str:
+        # FM999.9990
+        format_str = "FM" + "9" * (38 - coltype.precision)
+        if coltype.precision:
+            format_str += "0." + "9" * (coltype.precision - 1) + "0"
+        return f"to_char({value}, '{format_str}')"
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class Oracle(ThreadedDatabase):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
+    CONNECT_URI_HELP = "oracle://<user>:<password>@<host>/<database>"
+    CONNECT_URI_PARAMS = ["database?"]
+
+    kwargs: Dict[str, Any]
+    _oracle: Any
+
+    def __init__(self, *, host, database, thread_count, **kw) -> None:
+        super().__init__(thread_count=thread_count)
+        self.kwargs = dict(dsn=f"{host}/{database}" if database else host, **kw)
+        self.default_schema = kw.get("user").upper()
+        self._oracle = None
+
+    def create_connection(self):
+        self._oracle = import_oracle()
+        try:
+            c = self._oracle.connect(**self.kwargs)
+            if SESSION_TIME_ZONE:
+                c.cursor().execute(f"ALTER SESSION SET TIME_ZONE = '{SESSION_TIME_ZONE}'")
+            return c
+        except Exception as e:
+            raise ConnectError(*e.args) from e
+
+    def _query_cursor(self, c, sql_code: str):
+        try:
+            return super()._query_cursor(c, sql_code)
+        except self._oracle.DatabaseError as e:
+            raise QueryError(e)
+
+    def select_table_schema(self, path: DbPath) -> str:
+        schema, name = self._normalize_table_path(path)
+
+        return (
+            f"SELECT column_name, data_type, 6 as datetime_precision, data_precision as numeric_precision, data_scale as numeric_scale"
+            f" FROM ALL_TAB_COLUMNS WHERE table_name = '{name}' AND owner = '{schema}'"
+        )

data_diff/databases/postgresql.py

@@ -0,0 +1,258 @@
+from typing import Any, ClassVar, Dict, List, Type
+from urllib.parse import unquote
+import attrs
+
+from data_diff.abcs.database_types import (
+    ColType,
+    DbPath,
+    JSON,
+    Timestamp,
+    TimestampTZ,
+    Float,
+    Decimal,
+    Integer,
+    TemporalType,
+    Native_UUID,
+    Text,
+    FractionalType,
+    Boolean,
+    Date,
+    Time,
+)
+from data_diff.databases.base import BaseDialect, ThreadedDatabase, import_helper, ConnectError
+from data_diff.databases.base import (
+    MD5_HEXDIGITS,
+    CHECKSUM_HEXDIGITS,
+    _CHECKSUM_BITSIZE,
+    TIMESTAMP_PRECISION_POS,
+    CHECKSUM_OFFSET,
+)
+
+SESSION_TIME_ZONE = None  # Changed by the tests
+
+
+@import_helper("postgresql")
+def import_postgresql():
+    import psycopg2.extras
+
+    psycopg2.extensions.set_wait_callback(psycopg2.extras.wait_select)
+    return psycopg2
+
+
+@attrs.define(frozen=False)
+class PostgresqlDialect(BaseDialect):
+    name = "PostgreSQL"
+    ROUNDS_ON_PREC_LOSS = True
+    SUPPORTS_PRIMARY_KEY: ClassVar[bool] = True
+    SUPPORTS_INDEXES = True
+
+    # https://www.postgresql.org/docs/current/datatype-numeric.html#DATATYPE-NUMERIC-DECIMAL
+    # without any precision or scale creates an “unconstrained numeric” column
+    # in which numeric values of any length can be stored, up to the implementation limits.
+    # https://www.postgresql.org/docs/current/datatype-numeric.html#DATATYPE-NUMERIC-TABLE
+    DEFAULT_NUMERIC_PRECISION = 16383
+
+    TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {
+        # Timestamps
+        "timestamp with time zone": TimestampTZ,
+        "timestamp without time zone": Timestamp,
+        "timestamp": Timestamp,
+        "date": Date,
+        "time with time zone": Time,
+        "time without time zone": Time,
+        # Numbers
+        "double precision": Float,
+        "real": Float,
+        "decimal": Decimal,
+        "smallint": Integer,
+        "integer": Integer,
+        "numeric": Decimal,
+        "bigint": Integer,
+        # Text
+        "character": Text,
+        "character varying": Text,
+        "varchar": Text,
+        "text": Text,
+        "json": JSON,
+        "jsonb": JSON,
+        "uuid": Native_UUID,
+        "boolean": Boolean,
+    }
+
+    def quote(self, s: str):
+        return f'"{s}"'
+
+    def to_string(self, s: str):
+        return f"{s}::varchar"
+
+    def concat(self, items: List[str]) -> str:
+        joined_exprs = " || ".join(items)
+        return f"({joined_exprs})"
+
+    def _convert_db_precision_to_digits(self, p: int) -> int:
+        # Subtracting 2 due to wierd precision issues in PostgreSQL
+        return super()._convert_db_precision_to_digits(p) - 2
+
+    def set_timezone_to_utc(self) -> str:
+        return "SET TIME ZONE 'UTC'"
+
+    def current_timestamp(self) -> str:
+        return "current_timestamp"
+
+    def type_repr(self, t) -> str:
+        if isinstance(t, TimestampTZ):
+            return f"timestamp ({t.precision}) with time zone"
+        return super().type_repr(t)
+
+    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:
+        def _add_padding(coltype: TemporalType, timestamp6: str):
+            return f"RPAD(LEFT({timestamp6}, {TIMESTAMP_PRECISION_POS+coltype.precision}), {TIMESTAMP_PRECISION_POS+6}, '0')"
+
+        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"cast({value} as varchar)"
+
+        if isinstance(coltype, Time) or is_time:
+            seconds = f"EXTRACT( epoch from {value})"
+            rounded = f"ROUND({seconds},  {coltype.precision})"
+            time_value = f"CAST('00:00:00' as time) + make_interval(0, 0, 0, 0, 0, 0, {rounded})"  # 6th arg = seconds
+            converted = f"to_char({time_value}, 'hh24:mi:ss.ff6')"
+            return converted
+
+        if coltype.rounds:
+            # NULL value expected to return NULL after normalization
+            null_case_begin = f"CASE WHEN {value} IS NULL THEN NULL ELSE "
+            null_case_end = "END"
+
+            # 294277 or 4714 BC would be out of range, make sure we can't round to that
+            # TODO test timezones for overflow?
+            max_timestamp = "294276-12-31 23:59:59.0000"
+            min_timestamp = "4713-01-01 00:00:00.00 BC"
+            timestamp = f"least('{max_timestamp}'::timestamp(6), {value}::timestamp(6))"
+            timestamp = f"greatest('{min_timestamp}'::timestamp(6), {timestamp})"
+
+            interval = format((0.5 * (10 ** (-coltype.precision))), f".{coltype.precision+1}f")
+
+            rounded_timestamp = (
+                f"left(to_char(least('{max_timestamp}'::timestamp, {timestamp})"
+                f"+ interval '{interval}', 'YYYY-mm-dd HH24:MI:SS.US'),"
+                f"length(to_char(least('{max_timestamp}'::timestamp, {timestamp})"
+                f"+ interval '{interval}', 'YYYY-mm-dd HH24:MI:SS.US')) - (6-{coltype.precision}))"
+            )
+
+            padded = _add_padding(coltype, rounded_timestamp)
+            return f"{null_case_begin} {padded} {null_case_end}"
+
+            # TODO years with > 4 digits not padded correctly
+            # current w/ precision 6: 294276-12-31 23:59:59.0000
+            # should be 294276-12-31 23:59:59.000000
+        else:
+            rounded_timestamp = f"to_char({value}::timestamp(6), 'YYYY-mm-dd HH24:MI:SS.US')"
+            padded = _add_padding(coltype, rounded_timestamp)
+            return padded
+
+    def normalize_number(self, value: str, coltype: FractionalType) -> str:
+        return self.to_string(f"{value}::decimal(38, {coltype.precision})")
+
+    def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
+        return self.to_string(f"{value}::int")
+
+    def normalize_json(self, value: str, _coltype: JSON) -> str:
+        return f"{value}::text"
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class PostgreSQL(ThreadedDatabase):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = PostgresqlDialect
+    SUPPORTS_UNIQUE_CONSTAINT = True
+    CONNECT_URI_HELP = "postgresql://<user>:<password>@<host>/<database>"
+    CONNECT_URI_PARAMS = ["database?"]
+
+    _args: Dict[str, Any]
+    _conn: Any
+
+    def __init__(self, *, thread_count, **kw) -> None:
+        super().__init__(thread_count=thread_count)
+        self._args = kw
+        self.default_schema = "public"
+
+    def create_connection(self):
+        if not self._args:
+            self._args["host"] = None  # psycopg2 requires 1+ arguments
+
+        pg = import_postgresql()
+        try:
+            self._args["password"] = unquote(self._args["password"])
+            self._conn = pg.connect(
+                **self._args, keepalives=1, keepalives_idle=5, keepalives_interval=2, keepalives_count=2
+            )
+            if SESSION_TIME_ZONE:
+                self._conn.cursor().execute(f"SET TIME ZONE '{SESSION_TIME_ZONE}'")
+            return self._conn
+        except pg.OperationalError as e:
+            raise ConnectError(*e.args) from e
+
+    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,
+                    -- see comment for DEFAULT_NUMERIC_PRECISION
+                    CASE
+                        WHEN data_type = 'numeric'
+                            THEN coalesce(numeric_precision, 131072 + {self.dialect.DEFAULT_NUMERIC_PRECISION})
+                        ELSE numeric_precision
+                    END AS numeric_precision,
+                    CASE
+                        WHEN data_type = 'numeric'
+                            THEN coalesce(numeric_scale, {self.dialect.DEFAULT_NUMERIC_PRECISION})
+                        ELSE numeric_scale
+                    END AS numeric_scale
+                    FROM {'.'.join(info_schema_path)}
+                    WHERE table_name = '{table}' AND table_schema = '{schema}'
+            """
+
+    def select_table_unique_columns(self, path: DbPath) -> str:
+        database, schema, table = self._normalize_table_path(path)
+
+        info_schema_path = ["information_schema", "key_column_usage"]
+        if database:
+            info_schema_path.insert(0, database)
+
+        return (
+            "SELECT column_name "
+            f"FROM {'.'.join(info_schema_path)} "
+            f"WHERE table_name = '{table}' 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"
+        )
+
+    def close(self):
+        super().close()
+        if self._conn is not None:
+            self._conn.close()