dcs-sdk 1.6.5__py3-none-any.whl

data_diff/databases/presto.py

@@ -0,0 +1,222 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from functools import partial
+from typing import Any, ClassVar, Type
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    Boolean,
+    ColType,
+    ColType_UUID,
+    DbPath,
+    DbTime,
+    Decimal,
+    Float,
+    FractionalType,
+    Integer,
+    Native_UUID,
+    TemporalType,
+    Text,
+    Timestamp,
+    TimestampTZ,
+)
+from data_diff.databases.base import (
+    CHECKSUM_HEXDIGITS,
+    CHECKSUM_OFFSET,
+    MD5_HEXDIGITS,
+    TIMESTAMP_PRECISION_POS,
+    BaseDialect,
+    Database,
+    QueryResult,
+    ThreadLocalInterpreter,
+    import_helper,
+)
+from data_diff.schema import RawColumnInfo
+from data_diff.utils import match_regexps
+
+
+def query_cursor(c, sql_code):
+    try:
+        c.execute(sql_code)
+        if sql_code.lower().startswith(("select", "explain", "show")):
+            columns = c.description and [col[0] for col in c.description]
+            return QueryResult(c.fetchall(), columns)
+        elif sql_code.lower().startswith(("create", "drop")):
+            try:
+                c.connection.commit()
+            except AttributeError:
+                ...
+    except Exception as _e:
+        try:
+            c.connection.rollback()
+        except Exception as rollback_error:
+            print("Rollback failed:", rollback_error)
+        raise
+
+
+@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,
+        # UUID
+        "uuid": Native_UUID,
+    }
+
+    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, is_table: bool = False):
+        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:
+        raise NotImplementedError()
+
+    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,"
+            "NULL as collation_name, NULL as character_maximum_length "
+            "FROM INFORMATION_SCHEMA.COLUMNS "
+            f"WHERE table_name = '{table}' AND table_schema = '{schema}'"
+        )
+
+    @property
+    def is_autocommit(self) -> bool:
+        return False

data_diff/databases/redis.py

@@ -0,0 +1,93 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import os
+import threading
+import time
+
+from loguru import logger
+from redis import ConnectionPool, Redis, RedisError
+
+
+class RedisBackend:
+    _INSTANCE = None
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_instance(cls, force_recreate: bool = False):
+        with cls._lock:
+            if cls._INSTANCE is None or force_recreate:
+                logger.info("Creating a new RedisBackend instance (force_recreate=%s)", force_recreate)
+                cls._INSTANCE = cls._create_instance()
+            return cls._INSTANCE
+
+    @classmethod
+    def _create_instance(cls):
+        redis_url = (
+            os.getenv("DCS_REDIS_URL", None)
+            or os.getenv("DCS_RABBIT_URL", None)
+            or os.getenv("REDIS_URL", None)
+            or os.getenv("RABBIT_URL", None)
+        )
+        if not redis_url:
+            logger.warning("environment variable is not configured for redis")
+            return None
+        try:
+            pool = ConnectionPool.from_url(
+                redis_url,
+                max_connections=int(os.getenv("DCS_REDIS_MAX_CONNECTIONS", "100")),
+                health_check_interval=int(os.getenv("DCS_REDIS_HEALTH_CHECK_INTERVAL", "30")),
+                socket_connect_timeout=float(os.getenv("DCS_REDIS_SOCKET_CONNECT_TIMEOUT", "5")),
+                socket_timeout=float(os.getenv("DCS_REDIS_SOCKET_TIMEOUT", "5")),
+                socket_keepalive=True,
+            )
+            client = Redis(connection_pool=pool, decode_responses=False, retry_on_timeout=True)
+            # ping to ensure connection and raise early if config wrong
+            client.ping()
+            return cls(client)
+        except RedisError as exc:
+            logger.exception("Failed to create Redis client: %s", exc)
+            raise
+
+    def __init__(self, client: Redis):
+        self._client = client
+
+    @property
+    def client(self) -> Redis:
+        return self._client
+
+    def ensure_connected(self) -> bool:
+        try:
+            self._client.ping()
+            return True
+        except RedisError:
+            logger.warning("Redis ping failed; trying to recreate connection/pool")
+            try:
+                time.sleep(0.2)
+                new_instance = self.__class__._create_instance()
+                self.__class__._INSTANCE = new_instance
+                self._client = new_instance.client
+                self._client.ping()
+                logger.info("Recreated Redis client/pool successfully")
+                return True
+            except Exception as exc:
+                logger.exception("Failed to recreate Redis client: %s", exc)
+                return False
+
+    def close(self):
+        try:
+            self._client.connection_pool.disconnect()
+            logger.info("Closed Redis connection pool")
+        except Exception as e:
+            logger.exception("Error closing Redis connection pool: %s", e)

data_diff/databases/redshift.py

@@ -0,0 +1,233 @@
+#  Copyright 2022-present, the Waterdip Labs Pvt. Ltd.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from typing import Any, ClassVar, Dict, Iterable, List, Tuple, Type
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    JSON,
+    ColType,
+    DbPath,
+    Float,
+    FractionalType,
+    Integer,
+    TemporalType,
+    TimestampTZ,
+)
+from data_diff.databases.postgresql import (
+    CHECKSUM_HEXDIGITS,
+    CHECKSUM_OFFSET,
+    MD5_HEXDIGITS,
+    TIMESTAMP_PRECISION_POS,
+    BaseDialect,
+    PostgreSQL,
+    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"
+        )