dcs-sdk 1.6.5__py3-none-any.whl

data_diff/format.py

@@ -0,0 +1,369 @@
+#  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 collections
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple, Type
+
+import attrs
+
+from data_diff.abcs.database_types import (
+    JSON,
+    Array,
+    Boolean,
+    ColType,
+    ColType_Alphanum,
+    ColType_UUID,
+    Date,
+    FractionalType,
+    NumericType,
+    String_Alphanum,
+    Struct,
+    TemporalType,
+)
+from data_diff.diff_tables import DiffResultWrapper
+
+
+def jsonify_error(table1: List[str], table2: List[str], error: str) -> "FailedDiff":
+    return attrs.asdict(
+        FailedDiff(
+            status="failed",
+            dataset1=table1,
+            dataset2=table2,
+            error=error,
+        )
+    )
+
+
+Columns = List[Tuple[str, str, ColType]]
+
+
+def jsonify(
+    diff: DiffResultWrapper,
+    dataset1_columns: Columns,
+    dataset2_columns: Columns,
+    columns_diff: Dict[str, List[str]],
+    with_summary: bool = False,
+    stats_only: bool = False,
+) -> "JsonDiff":
+    """
+    Converts the diff result into a JSON-serializable format.
+    Optionally add stats summary and schema diff.
+    """
+    diff_info = diff.info_tree.info
+    table1 = diff_info.tables[0]
+    table2 = diff_info.tables[1]
+    key_columns = table1.key_columns
+
+    t1_exclusive_rows = []
+    t2_exclusive_rows = []
+    diff_rows = []
+    rows = None
+    schema = [field for field, _ in diff_info.diff_schema]
+
+    t1_exclusive_rows, t2_exclusive_rows, diff_rows = _group_rows(diff_info, schema)
+
+    if not stats_only:
+        rows = _make_rows_diff(t1_exclusive_rows, t2_exclusive_rows, diff_rows, key_columns)
+
+    summary = None
+    if with_summary:
+        summary = _jsonify_diff_summary(diff.get_stats_dict())
+
+    columns = _jsonify_columns_diff(dataset1_columns, dataset2_columns, columns_diff, list(key_columns))
+
+    is_different = bool(
+        t1_exclusive_rows
+        or t2_exclusive_rows
+        or diff_rows
+        or (columns_diff["added"] or columns_diff["removed"] or columns_diff["changed"])
+    )
+    return attrs.asdict(
+        JsonDiff(
+            status="success",
+            result="different" if is_different else "identical",
+            dataset1=list(table1.table_path),
+            dataset2=list(table2.table_path),
+            rows=rows,
+            summary=summary,
+            columns=columns,
+        )
+    )
+
+
+@attrs.define(frozen=True)
+class JsonExclusiveRowValue:
+    """
+    Value of a single column in a row
+    """
+
+    isPK: bool
+    value: Any
+
+
+@attrs.define(frozen=True)
+class JsonDiffRowValue:
+    """
+    Pair of diffed values for 2 rows with equal PKs
+    """
+
+    dataset1: Any
+    dataset2: Any
+    isDiff: bool
+    isPK: bool
+
+
+@attrs.define(frozen=True)
+class Total:
+    dataset1: int
+    dataset2: int
+
+
+@attrs.define(frozen=True)
+class ExclusiveRows:
+    dataset1: int
+    dataset2: int
+
+
+@attrs.define(frozen=True)
+class Rows:
+    total: Total
+    exclusive: ExclusiveRows
+    updated: int
+    unchanged: int
+
+
+@attrs.define(frozen=True)
+class Stats:
+    diffCounts: Dict[str, int]
+
+
+@attrs.define(frozen=True)
+class JsonDiffSummary:
+    rows: Rows
+    stats: Stats
+
+
+@attrs.define(frozen=True)
+class ExclusiveColumns:
+    dataset1: List[str]
+    dataset2: List[str]
+
+
+class ColumnKind(Enum):
+    INTEGER = "integer"
+    FLOAT = "float"
+    STRING = "string"
+    DATE = "date"
+    TIME = "time"
+    DATETIME = "datetime"
+    BOOL = "boolean"
+    UNSUPPORTED = "unsupported"
+
+
+KIND_MAPPING: List[Tuple[Type[ColType], ColumnKind]] = [
+    (Boolean, ColumnKind.BOOL),
+    (Date, ColumnKind.DATE),
+    (TemporalType, ColumnKind.DATETIME),
+    (FractionalType, ColumnKind.FLOAT),
+    (NumericType, ColumnKind.INTEGER),
+    (ColType_UUID, ColumnKind.STRING),
+    (ColType_Alphanum, ColumnKind.STRING),
+    (String_Alphanum, ColumnKind.STRING),
+    (JSON, ColumnKind.STRING),
+    (Array, ColumnKind.STRING),
+    (Struct, ColumnKind.STRING),
+    (ColType, ColumnKind.UNSUPPORTED),
+]
+
+
+@attrs.define(frozen=True)
+class Column:
+    name: str
+    type: str
+    kind: str
+
+
+@attrs.define(frozen=True)
+class JsonColumnsSummary:
+    dataset1: List[Column]
+    dataset2: List[Column]
+    primaryKey: List[str]
+    exclusive: ExclusiveColumns
+    typeChanged: List[str]
+
+
+@attrs.define(frozen=True)
+class ExclusiveDiff:
+    dataset1: List[Dict[str, JsonExclusiveRowValue]]
+    dataset2: List[Dict[str, JsonExclusiveRowValue]]
+
+
+@attrs.define(frozen=True)
+class RowsDiff:
+    exclusive: ExclusiveDiff
+    diff: List[Dict[str, JsonDiffRowValue]]
+
+
+@attrs.define(frozen=True)
+class FailedDiff:
+    status: str  # Literal ["failed"]
+    dataset1: List[str]
+    dataset2: List[str]
+    error: str
+
+    version: str = "1.0.0"
+
+
+@attrs.define(frozen=True)
+class JsonDiff:
+    status: str  # Literal ["success"]
+    result: str  # Literal ["different", "identical"]
+    dataset1: List[str]
+    dataset2: List[str]
+    rows: Optional[RowsDiff]
+    summary: Optional[JsonDiffSummary]
+    columns: Optional[JsonColumnsSummary]
+
+    version: str = "1.1.0"
+
+
+def _group_rows(
+    diff_info: DiffResultWrapper, schema: List[str]
+) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]:
+    t1_exclusive_rows = []
+    t2_exclusive_rows = []
+    diff_rows = []
+
+    for row in diff_info.diff:
+        row_w_schema = dict(zip(schema, row))
+        is_t1_exclusive = row_w_schema["is_exclusive_a"]
+        is_t2_exclusive = row_w_schema["is_exclusive_b"]
+
+        if is_t1_exclusive:
+            t1_exclusive_rows.append(row_w_schema)
+
+        elif is_t2_exclusive:
+            t2_exclusive_rows.append(row_w_schema)
+
+        else:
+            diff_rows.append(row_w_schema)
+
+    return t1_exclusive_rows, t2_exclusive_rows, diff_rows
+
+
+def _make_rows_diff(
+    t1_exclusive_rows: List[Dict[str, Any]],
+    t2_exclusive_rows: List[Dict[str, Any]],
+    diff_rows: List[Dict[str, Any]],
+    key_columns: List[str],
+) -> RowsDiff:
+    diff_rows_jsonified = []
+    for row in diff_rows:
+        diff_rows_jsonified.append(_jsonify_diff(row, key_columns))
+
+    t1_exclusive_rows_jsonified = []
+    for row in t1_exclusive_rows:
+        t1_exclusive_rows_jsonified.append(_jsonify_exclusive(row, key_columns))
+
+    t2_exclusive_rows_jsonified = []
+    for row in t2_exclusive_rows:
+        t2_exclusive_rows_jsonified.append(_jsonify_exclusive(row, key_columns))
+
+    return RowsDiff(
+        exclusive=ExclusiveDiff(dataset1=t1_exclusive_rows_jsonified, dataset2=t2_exclusive_rows_jsonified),
+        diff=diff_rows_jsonified,
+    )
+
+
+def _jsonify_diff(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, JsonDiffRowValue]:
+    columns = collections.defaultdict(dict)
+    for field, value in row.items():
+        if field in ("is_exclusive_a", "is_exclusive_b"):
+            continue
+
+        if field.startswith("is_diff_"):
+            column_name = field[len("is_diff_") :]
+            columns[column_name]["isDiff"] = bool(value)
+
+        elif field.endswith("_a"):
+            column_name = field[: -len("_a")]
+            columns[column_name]["dataset1"] = value
+            columns[column_name]["isPK"] = column_name in key_columns
+
+        elif field.endswith("_b"):
+            column_name = field[: -len("_b")]
+            columns[column_name]["dataset2"] = value
+            columns[column_name]["isPK"] = column_name in key_columns
+
+    return {column: JsonDiffRowValue(**data) for column, data in columns.items()}
+
+
+def _jsonify_exclusive(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, JsonExclusiveRowValue]:
+    columns = collections.defaultdict(dict)
+    for field, value in row.items():
+        if field in ("is_exclusive_a", "is_exclusive_b"):
+            continue
+        if field.startswith("is_diff_"):
+            continue
+        if field.endswith("_b") and row["is_exclusive_b"]:
+            column_name = field[: -len("_b")]
+            columns[column_name]["isPK"] = column_name in key_columns
+            columns[column_name]["value"] = value
+        elif field.endswith("_a") and row["is_exclusive_a"]:
+            column_name = field[: -len("_a")]
+            columns[column_name]["isPK"] = column_name in key_columns
+            columns[column_name]["value"] = value
+    return {column: JsonExclusiveRowValue(**data) for column, data in columns.items()}
+
+
+def _jsonify_diff_summary(stats_dict: dict) -> JsonDiffSummary:
+    return JsonDiffSummary(
+        rows=Rows(
+            total=Total(dataset1=stats_dict["rows_A"], dataset2=stats_dict["rows_B"]),
+            exclusive=ExclusiveRows(
+                dataset1=stats_dict["exclusive_A"],
+                dataset2=stats_dict["exclusive_B"],
+            ),
+            updated=stats_dict["updated"],
+            unchanged=stats_dict["unchanged"],
+        ),
+        stats=Stats(diffCounts=stats_dict["values"]),
+    )
+
+
+def _jsonify_columns_diff(
+    dataset1_columns: Columns, dataset2_columns: Columns, columns_diff: Dict[str, List[str]], key_columns: List[str]
+) -> JsonColumnsSummary:
+    return JsonColumnsSummary(
+        dataset1=[
+            Column(name=name, type=type_, kind=_map_kind(kind).value) for (name, type_, kind) in dataset1_columns
+        ],
+        dataset2=[
+            Column(name=name, type=type_, kind=_map_kind(kind).value) for (name, type_, kind) in dataset2_columns
+        ],
+        primaryKey=key_columns,
+        exclusive=ExclusiveColumns(
+            dataset2=list(columns_diff.get("added", [])),
+            dataset1=list(columns_diff.get("removed", [])),
+        ),
+        typeChanged=list(columns_diff.get("changed", [])),
+    )
+
+
+def _map_kind(kind: ColType) -> ColumnKind:
+    for raw_kind, json_kind in KIND_MAPPING:
+        if isinstance(kind, raw_kind):
+            return json_kind
+    return ColumnKind.UNSUPPORTED