dcs-sdk 1.6.5__py3-none-any.whl

data_diff/databases/databricks.py

@@ -0,0 +1,262 @@
+#  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 logging
+import math
+from typing import Any, ClassVar, Dict, Sequence, Type
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    Boolean,
+    ColType,
+    Date,
+    DbPath,
+    Decimal,
+    Float,
+    Integer,
+    NumericType,
+    TemporalType,
+    Text,
+    Timestamp,
+    UnknownColType,
+)
+from data_diff.databases.base import (
+    CHECKSUM_HEXDIGITS,
+    CHECKSUM_OFFSET,
+    MD5_HEXDIGITS,
+    BaseDialect,
+    ThreadedDatabase,
+    import_helper,
+    parse_table_name,
+)
+from data_diff.schema import RawColumnInfo
+
+
+@import_helper(text="You can install it using 'pip install databricks-sql-connector'")
+def import_databricks():
+    import databricks.sql
+
+    return databricks
+
+
+@attrs.define(frozen=False)
+class Dialect(BaseDialect):
+    name = "Databricks"
+    ROUNDS_ON_PREC_LOSS = True
+    TYPE_CLASSES = {
+        # Numbers
+        "INT": Integer,
+        "SMALLINT": Integer,
+        "TINYINT": Integer,
+        "BIGINT": Integer,
+        "FLOAT": Float,
+        "DOUBLE": Float,
+        "DECIMAL": Decimal,
+        # Timestamps
+        "TIMESTAMP": Timestamp,
+        "TIMESTAMP_NTZ": Timestamp,
+        "DATE": Date,
+        # Text
+        "STRING": Text,
+        "VARCHAR": Text,
+        # Boolean
+        "BOOLEAN": Boolean,
+    }
+
+    def type_repr(self, t) -> str:
+        try:
+            return {str: "STRING"}[t]
+        except KeyError:
+            return super().type_repr(t)
+
+    def quote(self, s: str, is_table: bool = False) -> str:
+        return f"`{s}`"
+
+    def to_string(self, s: str) -> str:
+        return f"cast({s} as string)"
+
+    def _convert_db_precision_to_digits(self, p: int) -> int:
+        # Subtracting 2 due to wierd precision issues
+        return max(super()._convert_db_precision_to_digits(p) - 2, 0)
+
+    def set_timezone_to_utc(self) -> str:
+        return "SET TIME ZONE 'UTC'"
+
+    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(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"""
+        try:
+            is_date = coltype.is_date
+        except:
+            is_date = False
+        if isinstance(coltype, Date) or is_date:
+            return f"date_format({value}, 'yyyy-MM-dd')"
+        if coltype.rounds:
+            # cast to timestamp due to unix_micros() requiring timestamp
+            timestamp = f"cast(round(unix_micros(cast({value} as timestamp)) / 1000000, {coltype.precision}) * 1000000 as bigint)"
+            return f"date_format(timestamp_micros({timestamp}), 'yyyy-MM-dd HH:mm:ss.SSSSSS')"
+
+        precision_format = "S" * coltype.precision + "0" * (6 - coltype.precision)
+        return f"date_format({value}, 'yyyy-MM-dd HH:mm:ss.{precision_format}')"
+
+    def normalize_number(self, value: str, coltype: NumericType) -> str:
+        value = f"cast({value} as decimal(38, {coltype.precision}))"
+        if coltype.precision > 0:
+            value = f"format_number({value}, {coltype.precision})"
+        return f"replace({self.to_string(value)}, ',', '')"
+
+    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 Databricks(ThreadedDatabase):
+    DIALECT_CLASS: ClassVar[Type[BaseDialect]] = Dialect
+    CONNECT_URI_HELP = "databricks://:<access_token>@<host>/<http_path>"
+    CONNECT_URI_PARAMS = ["catalog", "schema"]
+
+    catalog: str
+    _args: Dict[str, Any]
+
+    def __init__(self, *, thread_count, **kw) -> None:
+        super().__init__(thread_count=thread_count)
+        logging.getLogger("databricks.sql").setLevel(logging.WARNING)
+
+        self._args = kw
+        self.default_schema = kw.get("schema", "default")
+        self.catalog = kw.get("catalog", "hive_metastore")
+
+    def create_connection(self):
+        databricks = import_databricks()
+
+        try:
+            return databricks.sql.connect(
+                server_hostname=self._args["host"],
+                http_path=self._args["http_path"],
+                access_token=self._args["access_token"],
+                catalog=self.catalog,
+            )
+        except databricks.sql.exc.Error as e:
+            raise ConnectionError(*e.args) from e
+
+    def query_table_schema(self, path: DbPath) -> Dict[str, RawColumnInfo]:
+        # Databricks has INFORMATION_SCHEMA only for Databricks Runtime, not for Databricks SQL.
+        # https://docs.databricks.com/spark/latest/spark-sql/language-manual/information-schema/columns.html
+        # So, to obtain information about schema, we should use another approach.
+
+        conn = self.create_connection()
+
+        catalog, schema, table = self._normalize_table_path(path)
+        with conn.cursor() as cursor:
+            cursor.columns(catalog_name=catalog, schema_name=schema, table_name=table)
+            try:
+                rows = cursor.fetchall()
+            finally:
+                conn.close()
+            if not rows:
+                raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns")
+
+            d = {
+                r.COLUMN_NAME: RawColumnInfo(
+                    column_name=r.COLUMN_NAME, data_type=r.TYPE_NAME, numeric_precision=r.DECIMAL_DIGITS
+                )
+                for r in rows
+            }
+            assert len(d) == len(rows)
+            return d
+
+    # 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 "
+    #         f"FROM {'.'.join(info_schema_path)} "
+    #         f"WHERE table_name = '{name}' AND table_schema = '{schema}'"
+    #     )
+
+    def _process_table_schema(
+        self, path: DbPath, raw_schema: Dict[str, RawColumnInfo], filter_columns: Sequence[str], where: str = None
+    ):
+        accept = {i.lower() for i in filter_columns}
+        col_infos = [row for name, row in raw_schema.items() if name.lower() in accept]
+
+        resulted_rows = []
+        for info in col_infos:
+            raw_data_type = info.data_type
+            row_type = info.data_type.split("(")[0]
+            info = attrs.evolve(info, data_type=row_type)
+            type_cls = self.dialect.TYPE_CLASSES.get(row_type, UnknownColType)
+
+            if issubclass(type_cls, Integer):
+                info = attrs.evolve(info, numeric_scale=0)
+
+            elif issubclass(type_cls, Float):
+                numeric_precision = math.ceil(info.numeric_precision / math.log(2, 10))
+                info = attrs.evolve(info, numeric_precision=numeric_precision)
+
+            elif issubclass(type_cls, Decimal):
+                items = raw_data_type[8:].rstrip(")").split(",")
+                numeric_precision, numeric_scale = int(items[0]), int(items[1])
+                info = attrs.evolve(
+                    info,
+                    numeric_precision=numeric_precision,
+                    numeric_scale=numeric_scale,
+                )
+
+            elif issubclass(type_cls, Timestamp):
+                info = attrs.evolve(
+                    info,
+                    datetime_precision=info.numeric_precision,
+                    numeric_precision=None,
+                )
+
+            else:
+                info = attrs.evolve(info, numeric_precision=None)
+
+            resulted_rows.append(info)
+
+        col_dict: Dict[str, ColType] = {info.column_name: self.dialect.parse_type(path, info) for info in resulted_rows}
+
+        self._refine_coltypes(path, col_dict, where)
+        return col_dict
+
+    @property
+    def is_autocommit(self) -> bool:
+        return True
+
+    def _normalize_table_path(self, path: DbPath) -> DbPath:
+        if len(path) == 1:
+            return self.catalog, self.default_schema, path[0]
+        elif len(path) == 2:
+            return self.catalog, 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 catalog.schema.table"
+        )

data_diff/databases/duckdb.py

@@ -0,0 +1,207 @@
+#  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, Type, Union
+
+import attrs
+from packaging.version import parse as parse_version
+
+from data_diff.abcs.database_types import (
+    Boolean,
+    ColType,
+    DbPath,
+    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,
+    ConnectError,
+    Database,
+    ThreadLocalInterpreter,
+    import_helper,
+)
+from data_diff.schema import RawColumnInfo
+from data_diff.utils import match_regexps
+from data_diff.version import __version__
+
+
+@import_helper("duckdb")
+def import_duckdb():
+    import duckdb
+
+    return duckdb
+
+
+@attrs.define(frozen=False)
+class Dialect(BaseDialect):
+    name = "DuckDB"
+    ROUNDS_ON_PREC_LOSS = False
+    SUPPORTS_PRIMARY_KEY = True
+    SUPPORTS_INDEXES = True
+
+    # https://duckdb.org/docs/sql/data_types/numeric#fixed-point-decimals
+    # The default WIDTH and SCALE is DECIMAL(18, 3), if none are specified.
+    DEFAULT_NUMERIC_PRECISION = 3
+
+    TYPE_CLASSES = {
+        # Timestamps
+        "TIMESTAMP WITH TIME ZONE": TimestampTZ,
+        "TIMESTAMP": Timestamp,
+        # Numbers
+        "DOUBLE": Float,
+        "FLOAT": Float,
+        "DECIMAL": Decimal,
+        "INTEGER": Integer,
+        "BIGINT": Integer,
+        # Text
+        "VARCHAR": Text,
+        "TEXT": Text,
+        # UUID
+        "UUID": Native_UUID,
+        # Bool
+        "BOOLEAN": Boolean,
+    }
+
+    def quote(self, s: str, is_table: bool = False):
+        return f'"{s}"'
+
+    def to_string(self, s: str):
+        return f"{s}::VARCHAR"
+
+    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 parse_type(self, table_path: DbPath, info: RawColumnInfo) -> ColType:
+        regexps = {
+            r"DECIMAL\((\d+),(\d+)\)": Decimal,
+        }
+
+        for m, t_cls in match_regexps(regexps, info.data_type):
+            precision = int(m.group(2))
+            return t_cls(precision=precision)
+
+        return super().parse_type(table_path, info)
+
+    def set_timezone_to_utc(self) -> str:
+        return "SET GLOBAL TimeZone='UTC'"
+
+    def current_timestamp(self) -> str:
+        return "current_timestamp"
+
+    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:
+            return f"CONCAT(SUBSTRING(STRFTIME({value}::TIMESTAMP, '%Y-%m-%d %H:%M:%S.'),1,23), LPAD(((ROUND(strftime({value}::timestamp, '%f')::DECIMAL(15,7)/100000,{coltype.precision-1})*100000)::INT)::VARCHAR,6,'0'))"
+
+        return f"rpad(substring(strftime({value}::timestamp, '%Y-%m-%d %H:%M:%S.%f'),1,{TIMESTAMP_PRECISION_POS+coltype.precision}),26,'0')"
+
+    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}::INTEGER")
+
+
+@attrs.define(frozen=False, init=False, kw_only=True)
+class DuckDB(Database):
+    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"]
+
+    _args: Dict[str, Any] = attrs.field(init=False)
+    _conn: Any = attrs.field(init=False)
+
+    def __init__(self, **kw) -> None:
+        super().__init__()
+        self._args = kw
+        self._conn = self.create_connection()
+        self.default_schema = "main"
+
+    @property
+    def is_autocommit(self) -> bool:
+        return True
+
+    def _query(self, sql_code: Union[str, ThreadLocalInterpreter]):
+        "Uses the standard SQL cursor interface"
+        return self._query_conn(self._conn, sql_code)
+
+    def close(self):
+        super().close()
+        self._conn.close()
+
+    def create_connection(self):
+        ddb = import_duckdb()
+        try:
+            # custom_user_agent is only available in duckdb >= 0.9.2
+            if parse_version(ddb.__version__) >= parse_version("0.9.2"):
+                custom_user_agent = f"dcs-diff/v{__version__}"
+                config = {"custom_user_agent": custom_user_agent}
+                connection = ddb.connect(database=self._args["filepath"], config=config)
+                custom_user_agent_results = connection.sql("PRAGMA USER_AGENT;").fetchall()
+                custom_user_agent_filtered = custom_user_agent_results[0][0]
+                assert custom_user_agent in custom_user_agent_filtered
+            else:
+                connection = ddb.connect(database=self._args["filepath"])
+            return connection
+        except ddb.OperationalError as e:
+            raise ConnectError(*e.args) from e
+        except AssertionError:
+            raise ConnectError("Assertion failed: Custom user agent is invalid.") from None
+
+    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)
+            dynamic_database_clause = f"'{database}'"
+        else:
+            dynamic_database_clause = "current_catalog()"
+
+        return (
+            f"SELECT column_name, data_type, datetime_precision, numeric_precision, numeric_scale, collation_name, character_maximum_length FROM {'.'.join(info_schema_path)} "
+            f"WHERE table_name = '{table}' AND table_schema = '{schema}' and table_catalog = {dynamic_database_clause}"
+        )
+
+    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"
+        )