dcs-sdk 1.6.5__py3-none-any.whl

data_diff/databases/oracle.py

@@ -0,0 +1,238 @@
+#  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 secrets
+import string
+import time
+from typing import Any, ClassVar, Dict, List, Optional, Type
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    ColType,
+    ColType_UUID,
+    DbPath,
+    DbTime,
+    Decimal,
+    Float,
+    FractionalType,
+    TemporalType,
+    Text,
+    Timestamp,
+    TimestampTZ,
+)
+from data_diff.databases.base import (
+    CHECKSUM_HEXDIGITS,
+    CHECKSUM_OFFSET,
+    MD5_HEXDIGITS,
+    TIMESTAMP_PRECISION_POS,
+    BaseDialect,
+    ConnectError,
+    QueryError,
+    ThreadedDatabase,
+    import_helper,
+)
+from data_diff.schema import RawColumnInfo
+from data_diff.utils import match_regexps
+
+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, is_table: bool = False) -> str:
+        if s in self.TABLE_NAMES and self.default_schema and is_table:
+            return f'"{self.default_schema}"."{s}"'
+        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}')"
+
+    def generate_view_name(self, view_name: str | None = None) -> str:
+        if view_name is not None:
+            return view_name.upper()
+        random_string = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(8))
+        timestamp = int(time.time())
+        return f"view_{timestamp}_{random_string.lower()}".upper()
+
+    def parse_table_name(self, name: str) -> DbPath:
+        "Parse the given table name into a DbPath"
+        self.TABLE_NAMES.append(name.split(".")[-1])
+        return tuple(name.split("."))
+
+
+@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
+    _conn: 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("schema", None) or kw.get("user").upper()
+        self.dialect.default_schema = self.default_schema
+        self.kwargs = {k: v for k, v in self.kwargs.items() if v}
+        self._oracle = None
+        if "schema" in self.kwargs:
+            del self.kwargs["schema"]
+        self._conn = self.create_connection()
+
+    def create_connection(self):
+        self._oracle = import_oracle()
+        try:
+            self._conn = self._oracle.connect(**self.kwargs)
+            if SESSION_TIME_ZONE:
+                self._conn.cursor().execute(f"ALTER SESSION SET TIME_ZONE = '{SESSION_TIME_ZONE}'")
+            return self._conn
+        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, "
+            f"data_scale as numeric_scale, NULL as collation_name, char_length as character_maximum_length "
+            f"FROM ALL_TAB_COLUMNS WHERE table_name = '{name}' AND owner = '{schema}'"
+        )
+
+    def close(self):
+        super().close()
+        if self._conn is not None:
+            self._conn.close()

data_diff/databases/postgresql.py

@@ -0,0 +1,293 @@
+#  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, List, Tuple, Type
+from urllib.parse import unquote
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    JSON,
+    Boolean,
+    ColType,
+    Date,
+    DbPath,
+    Decimal,
+    Float,
+    FractionalType,
+    Integer,
+    Native_UUID,
+    TemporalType,
+    Text,
+    Time,
+    Timestamp,
+    TimestampTZ,
+)
+from data_diff.databases.base import (
+    _CHECKSUM_BITSIZE,
+    CHECKSUM_HEXDIGITS,
+    CHECKSUM_OFFSET,
+    MD5_HEXDIGITS,
+    TIMESTAMP_PRECISION_POS,
+    BaseDialect,
+    ConnectError,
+    QueryResult,
+    ThreadedDatabase,
+    import_helper,
+)
+
+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, is_table: bool = False) -> str:
+        if s in self.TABLE_NAMES and self.default_schema and is_table:
+            return f'"{self.default_schema}"."{s}"'
+        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:
+        precision = min(coltype.precision, 10)
+        return self.to_string(f"{value}::decimal(38, {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"
+
+    def parse_table_name(self, name: str) -> DbPath:
+        "Parse the given table name into a DbPath"
+        self.TABLE_NAMES.append(name.split(".")[-1])
+        return tuple(name.split("."))
+
+
+@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 = self._args.get("schema", "public")
+        self.dialect.default_schema = self.default_schema
+
+    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(
+                database=self._args.get("database"),
+                user=self._args.get("user"),
+                password=self._args.get("password"),
+                host=self._args.get("host"),
+                port=self._args.get("port"),
+                keepalives=1,
+                keepalives_idle=5,
+                keepalives_interval=2,
+                keepalives_count=2,
+                options="-c search_path={}".format(self.default_schema),
+            )
+            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, "
+            f"CASE WHEN data_type = 'numeric' "
+            f"THEN coalesce(numeric_precision, 131072 + {self.dialect.DEFAULT_NUMERIC_PRECISION}) "
+            f"ELSE numeric_precision END AS numeric_precision, "
+            f"CASE WHEN data_type = 'numeric' "
+            f"THEN coalesce(numeric_scale, {self.dialect.DEFAULT_NUMERIC_PRECISION}) "
+            f"ELSE numeric_scale END AS numeric_scale, "
+            f"COALESCE(collation_name, NULL) AS collation_name, "
+            f"CASE WHEN data_type = 'character varying' "
+            f"THEN character_maximum_length END AS character_maximum_length "
+            f"FROM {'.'.join(info_schema_path)} "
+            f"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()