openmetadata-ingestion 1.4.7.2__py3-none-any.whl → 1.5.0.0__py3-none-any.whl

metadata/data_quality/validations/runtime_param_setter/table_diff_params_setter.py

@@ -0,0 +1,190 @@
+#  Copyright 2024 Collate
+#  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.
+"""Module that defines the TableDiffParamsSetter class."""
+from ast import literal_eval
+from typing import List, Optional
+from urllib.parse import urlparse
+
+from sqlalchemy.engine import Engine
+
+from metadata.data_quality.validations.models import (
+    Column,
+    TableDiffRuntimeParameters,
+    TableParameter,
+)
+from metadata.data_quality.validations.runtime_param_setter.param_setter import (
+    RuntimeParameterSetter,
+)
+from metadata.generated.schema.entity.data.table import Constraint, Table
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.tests.testCase import TestCase
+from metadata.ingestion.source.connections import get_connection
+from metadata.utils import fqn
+
+
+class TableDiffParamsSetter(RuntimeParameterSetter):
+    """
+    Set runtime parameters for a the table diff test.
+    Sets the following variables:
+    - service1Url: The url of the first service (data diff compliant)
+    - service2Url: The url of the second service (data diff compliant)
+    - table1: The table path for the first service (only schema and table name)
+    - table2: The table path for the second service (only schema and table name)
+    - keyColumns: If not defined, construct the key columns based on primary key or unique constraint.
+    - extraColumns: If not defined, construct the extra columns as all columns except the key columns.
+    - whereClause: Exrtact where clause based on partitioning and user input
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._table_columns = {
+            column.name.root: column for column in self.table_entity.columns
+        }
+
+    def get_parameters(self, test_case) -> TableDiffRuntimeParameters:
+        service1: Engine = get_connection(self.service_connection_config)
+        table2_fqn = self.get_parameter(test_case, "table2")
+        table2: Table = self.ometa_client.get_by_name(
+            Table, fqn=table2_fqn, nullable=False
+        )
+        service2 = self.get_service2_url(service1, table2, test_case)
+        key_columns = self.get_key_columns(test_case)
+        extra_columns = self.get_extra_columns(key_columns, test_case)
+        return TableDiffRuntimeParameters(
+            table1=TableParameter(
+                path=self.get_data_diff_table_path(
+                    self.table_entity.fullyQualifiedName.root
+                ),
+                serviceUrl=self.get_data_diff_url(
+                    str(service1.url), self.table_entity.fullyQualifiedName.root
+                ),
+                columns=self.filter_relevant_columns(
+                    self.table_entity.columns, key_columns, extra_columns
+                ),
+            ),
+            table2=TableParameter(
+                path=self.get_data_diff_table_path(table2_fqn),
+                serviceUrl=self.get_data_diff_url(service2, table2_fqn),
+                columns=self.filter_relevant_columns(
+                    table2.columns, key_columns, extra_columns
+                ),
+            ),
+            keyColumns=key_columns,
+            extraColumns=extra_columns,
+            whereClause=self.build_where_clause(test_case),
+        )
+
+    # pylint: disable=protected-access
+    def build_where_clause(self, test_case) -> Optional[str]:
+        param_where_clause = self.get_parameter(test_case, "where", None)
+        partition_where_clause = (
+            None
+            if not self.sampler._partition_details
+            or not self.sampler._partition_details.enablePartitioning
+            else self.sampler.get_partitioned_query().whereclause.compile(
+                compile_kwargs={"literal_binds": True}
+            )
+        )
+        where_clauses = [param_where_clause, partition_where_clause]
+        where_clauses = [x for x in where_clauses if x]
+        where_clauses = [f"({x})" for x in where_clauses]
+        return " AND ".join(where_clauses)
+
+    def get_service2_url(self, service1, table2, test_case):
+        service2 = self.get_parameter(test_case, "service2Url")
+        if service2 is not None:
+            pass
+        elif self.table_entity.service.id == table2.service.id:
+            service2 = str(service1.url)
+        else:
+            table2_service = self.ometa_client.get_by_id(
+                DatabaseService, table2.service.id
+            )
+            service2 = str(get_connection(table2_service.connection.config).url)
+        return service2
+
+    def get_extra_columns(
+        self, key_columns: List[str], test_case
+    ) -> Optional[List[str]]:
+        extra_columns_param = self.get_parameter(test_case, "useColumns", None)
+        if extra_columns_param is not None:
+            extra_columns: List[str] = literal_eval(extra_columns_param)
+            self.validate_columns(extra_columns)
+            return extra_columns
+        if extra_columns_param is None:
+            extra_columns_param = []
+            for column in self.table_entity.columns:
+                if column.name.root not in key_columns:
+                    extra_columns_param.insert(0, column.name.root)
+        return extra_columns_param
+
+    def get_key_columns(self, test_case) -> List[str]:
+        key_columns_param = self.get_parameter(test_case, "keyColumns", "[]")
+        key_columns: List[str] = literal_eval(key_columns_param)
+        if key_columns:
+            self.validate_columns(key_columns)
+        if not key_columns:
+            for column in self.table_entity.columns:
+                if column.constraint == Constraint.PRIMARY_KEY:
+                    key_columns.append(column.name.root)
+        if not key_columns:
+            for column in self.table_entity.columns:
+                if column.constraint == Constraint.UNIQUE:
+                    key_columns.append(column.name.root)
+        if not key_columns:
+            raise ValueError(
+                "Failed to resolve key columns for table diff.\n",
+                "Could not find primary key or unique constraint columns.\n",
+                "Specify 'keyColumns' to explicitly set the columns to use as keys.",
+            )
+        return key_columns
+
+    @staticmethod
+    def filter_relevant_columns(
+        columns: List[Column], key_columns: List[str], extra_columns: List[str]
+    ) -> List[Column]:
+        return [c for c in columns if c.name.root in [*key_columns, *extra_columns]]
+
+    @staticmethod
+    def get_parameter(test_case: TestCase, key: str, default=None):
+        return next(
+            (p.value for p in test_case.parameterValues if p.name == key), default
+        )
+
+    @staticmethod
+    def get_data_diff_url(service_url: str, table_fqn) -> str:
+        url = urlparse(service_url)
+        # remove the driver name from the url because table-diff doesn't support it
+        kwargs = {"scheme": url.scheme.split("+")[0]}
+        service, database, schema, table = fqn.split(  # pylint: disable=unused-variable
+            table_fqn
+        )
+        # path needs to include the database AND schema in some of the connectors
+        if kwargs["scheme"] in ["mssql"]:
+            kwargs["path"] = f"/{database}/{schema}"
+        return url._replace(**kwargs).geturl()
+
+    @staticmethod
+    def get_data_diff_table_path(table_fqn: str) -> str:
+        service, database, schema, table = fqn.split(  # pylint: disable=unused-variable
+            table_fqn
+        )
+        return fqn._build(  # pylint: disable=protected-access
+            "___SERVICE___", "__DATABASE__", schema, table
+        ).replace("___SERVICE___.__DATABASE__.", "")
+
+    def validate_columns(self, column_names: List[str]):
+        for column in column_names:
+            if not self._table_columns.get(column):
+                raise ValueError(
+                    f"Failed to resolve key columns for table diff.\n"
+                    f"Column '{column}' not found in table '{self.table_entity.name.root}'.\n"
+                )

metadata/data_quality/validations/table/base/tableRowCountToBeBetween.py

@@ -52,25 +52,16 @@ class BaseTableRowCountToBeBetweenValidator(BaseTestValidator):
                 [TestResultValue(name=ROW_COUNT, value=None)],
             )
 
-        min_bound = self.get_test_case_param_value(
-            self.test_case.parameterValues,  # type: ignore
-            "minValue",
-            float,
-            default=float("-inf"),
-        )
-
-        max_bound = self.get_test_case_param_value(
-            self.test_case.parameterValues,  # type: ignore
-            "maxValue",
-            float,
-            default=float("inf"),
-        )
+        min_bound = self.get_min_bound("minValue")
+        max_bound = self.get_max_bound("maxValue")
 
         return self.get_test_case_result_object(
             self.execution_date,
             self.get_test_case_status(min_bound <= res <= max_bound),
             f"Found rowCount={res} vs. the expected min={min_bound},  max={max_bound}.",
             [TestResultValue(name=ROW_COUNT, value=str(res))],
+            min_bound=min_bound,
+            max_bound=max_bound,
         )
 
     @abstractmethod

metadata/data_quality/validations/table/pandas/tableRowInsertedCountToBeBetween.py

@@ -13,7 +13,7 @@
 Validator for table row inserted count to be between test case
 """
 
-from datetime import datetime
+from datetime import datetime, timezone
 
 from dateutil.relativedelta import relativedelta
 
@@ -47,7 +47,7 @@ class TableRowInsertedCountToBeBetweenValidator(
             "MONTH": relativedelta(months=range_interval),
             "YEAR": relativedelta(years=range_interval),
         }
-        utc_now = datetime.utcnow()
+        utc_now = datetime.now(timezone.utc)
         threshold_date = utc_now - interval_type_matching_table[range_type]
         if range_type == "HOUR":
             threshold_date = threshold_date.replace(minute=0, second=0, microsecond=0)

metadata/data_quality/validations/table/sqlalchemy/tableDiff.py

@@ -0,0 +1,434 @@
+#  Copyright 2024 Collate
+#  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.
+# pylint: disable=missing-module-docstring
+import logging
+import traceback
+from itertools import islice
+from typing import Dict, Iterable, List, Optional, Tuple
+from urllib.parse import urlparse
+
+import data_diff
+import sqlalchemy.types
+from data_diff.diff_tables import DiffResultWrapper
+from data_diff.errors import DataDiffMismatchingKeyTypesError
+from data_diff.utils import ArithAlphanumeric
+from sqlalchemy import Column as SAColumn
+
+from metadata.data_quality.validations.base_test_handler import BaseTestValidator
+from metadata.data_quality.validations.mixins.sqa_validator_mixin import (
+    SQAValidatorMixin,
+)
+from metadata.data_quality.validations.models import TableDiffRuntimeParameters
+from metadata.generated.schema.entity.data.table import Column
+from metadata.generated.schema.entity.services.connections.database.sapHanaConnection import (
+    SapHanaScheme,
+)
+from metadata.generated.schema.tests.basic import (
+    TestCaseResult,
+    TestCaseStatus,
+    TestResultValue,
+)
+from metadata.profiler.orm.registry import Dialects
+from metadata.utils.logger import test_suite_logger
+
+logger = test_suite_logger()
+
+SUPPORTED_DIALECTS = [
+    Dialects.Snowflake,
+    Dialects.BigQuery,
+    Dialects.Athena,
+    Dialects.Redshift,
+    Dialects.Postgres,
+    Dialects.MySQL,
+    Dialects.MSSQL,
+    Dialects.Oracle,
+    Dialects.Trino,
+    SapHanaScheme.hana.value,
+]
+
+
+class UnsupportedDialectError(Exception):
+    def __init__(self, param: str, dialect: str):
+        super().__init__(f"Unsupported dialect in param {param}: {dialect}")
+
+
+def masked(s: str, mask: bool = True) -> str:
+    """Mask a string if masked is True otherwise return the string.
+    Only for development purposes, do not use in production.
+    Change it False if you want to see the data in the logs.
+
+    Args:
+        s: string to mask
+        mask: boolean to mask the string
+
+    Returns:
+        masked string if mask is True otherwise return the string
+    """
+    return "***" if mask else s
+
+
+class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
+    """
+    Compare two tables and fail if the number of differences exceeds a threshold
+    """
+
+    runtime_params: TableDiffRuntimeParameters
+
+    def run_validation(self) -> TestCaseResult:
+        self.runtime_params = self.get_runtime_params()
+        try:
+            self._validate_dialects()
+            return self._run()
+        except DataDiffMismatchingKeyTypesError as e:
+            result = TestCaseResult(
+                timestamp=self.execution_date,  # type: ignore
+                testCaseStatus=TestCaseStatus.Failed,
+                result=str(e),
+            )
+            return result
+        except UnsupportedDialectError as e:
+            result = TestCaseResult(
+                timestamp=self.execution_date,  # type: ignore
+                testCaseStatus=TestCaseStatus.Aborted,
+                result=str(e),
+            )
+            return result
+        except Exception as e:
+            logger.error(
+                f"Unexpected error while running the table diff test: {str(e)}\n{traceback.format_exc()}"
+            )
+            result = TestCaseResult(
+                timestamp=self.execution_date,  # type: ignore
+                testCaseStatus=TestCaseStatus.Aborted,
+                result=f"ERROR: Unexpected error while running the table diff test: {str(e)}",
+            )
+            logger.debug(result.result)
+            return result
+
+    def _run(self) -> TestCaseResult:
+        result = self.get_column_diff()
+        if result:
+            return result
+        threshold = self.get_test_case_param_value(
+            self.test_case.parameterValues, "threshold", int, default=0
+        )
+        table_diff_iter = self.get_table_diff()
+
+        if not threshold or self.test_case.computePassedFailedRowCount:
+            stats = table_diff_iter.get_stats_dict()
+            if stats["total"] > 0:
+                logger.debug("Sample of failed rows:")
+                # depending on the data, this require scanning a lot of data
+                # so we only log the sample in debug mode. data can be sensitive
+                # so it is masked by default
+                for s in islice(
+                    self.safe_table_diff_iterator(),
+                    10 if logger.level <= logging.DEBUG else 0,
+                ):
+                    logger.debug("%s", str([s[0]] + [masked(st) for st in s[1]]))
+            test_case_result = self.get_row_diff_test_case_result(
+                threshold,
+                stats["total"],
+                stats["updated"],
+                stats["exclusive_A"],
+                stats["exclusive_B"],
+            )
+            count = self._compute_row_count(self.runner, None)  # type: ignore
+            test_case_result.passedRows = stats["unchanged"]
+            test_case_result.passedRowsPercentage = (
+                test_case_result.passedRows / count * 100
+            )
+            test_case_result.failedRowsPercentage = (
+                test_case_result.failedRows / count * 100
+            )
+            return test_case_result
+        return self.get_row_diff_test_case_result(
+            threshold,
+            self.calculate_diffs_with_limit(table_diff_iter, threshold),
+        )
+
+    def get_incomparable_columns(self) -> List[str]:
+        """Get the columns that have types that are not comparable between the two tables. For example
+        a column that is a string in one table and an integer in the other.
+
+        Returns:
+            List[str]: A list of column names that have incomparable types
+        """
+        table1 = data_diff.connect_to_table(
+            self.runtime_params.table1.serviceUrl,
+            self.runtime_params.table1.path,
+            self.runtime_params.keyColumns,
+            extra_columns=self.runtime_params.extraColumns,
+        ).with_schema()
+        table2 = data_diff.connect_to_table(
+            self.runtime_params.table2.serviceUrl,
+            self.runtime_params.table2.path,
+            self.runtime_params.keyColumns,
+            extra_columns=self.runtime_params.extraColumns,
+        ).with_schema()
+        result = []
+        for column in table1.key_columns + table1.extra_columns:
+            col1_type = self._get_column_python_type(
+                table1._schema[column]  # pylint: disable=protected-access
+            )
+            # Skip columns that are not in the second table. We cover this case in get_changed_added_columns.
+            if table2._schema.get(column) is None:  # pylint: disable=protected-access
+                continue
+            col2_type = self._get_column_python_type(
+                table2._schema[column]  # pylint: disable=protected-access
+            )
+
+            if col1_type != col2_type:
+                result.append(column)
+        return result
+
+    @staticmethod
+    def _get_column_python_type(column: SAColumn):
+        """Try to resolve the python_type of a column by cascading through different SQLAlchemy types.
+        If no type is found, return the name of the column type. This is usually undesirable since it can
+        be very database specific, but it is better than nothing.
+
+        Args:
+            column: An SQLAlchemy column object
+        """
+        result = None
+        try:
+            result = column.python_type
+        except AttributeError:
+            pass
+        try:
+            result = getattr(sqlalchemy.types, type(column).__name__)().python_type
+        except AttributeError:
+            pass
+        try:
+            result = getattr(
+                sqlalchemy.types, type(column).__name__.upper()
+            )().python_type
+        except AttributeError:
+            pass
+        if result == ArithAlphanumeric:
+            result = str
+        elif result == bool:
+            result = int
+        elif result is None:
+            return type(result)
+        return result
+
+    def get_table_diff(self) -> DiffResultWrapper:
+        """Calls data_diff.diff_tables with the parameters from the test case."""
+        table1 = data_diff.connect_to_table(
+            self.runtime_params.table1.serviceUrl,
+            self.runtime_params.table1.path,
+            self.runtime_params.keyColumns,  # type: ignore
+        )
+        table2 = data_diff.connect_to_table(
+            self.runtime_params.table2.serviceUrl,
+            self.runtime_params.table2.path,
+            self.runtime_params.keyColumns,  # type: ignore
+        )
+        data_diff_kwargs = {
+            "key_columns": self.runtime_params.keyColumns,
+            "extra_columns": self.runtime_params.extraColumns,
+            "where": self.get_where(),
+        }
+        logger.debug(
+            "Calling table diff with parameters: table1=%s, table2=%s, kwargs=%s",
+            table1.table_path,
+            table2.table_path,
+            ",".join(f"{k}={v}" for k, v in data_diff_kwargs.items()),
+        )
+        return data_diff.diff_tables(table1, table2, **data_diff_kwargs)  # type: ignore
+
+    def get_where(self) -> Optional[str]:
+        """Returns the where clause from the test case parameters or None if it is a blank string."""
+        return self.runtime_params.whereClause or None
+
+    def get_runtime_params(self) -> TableDiffRuntimeParameters:
+        raw = self.get_test_case_param_value(
+            self.test_case.parameterValues, "runtimeParams", str
+        )
+        runtime_params = TableDiffRuntimeParameters.model_validate_json(raw)
+        return runtime_params
+
+    def get_row_diff_test_case_result(
+        self,
+        threshold: int,
+        total_diffs: int,
+        changed: Optional[int] = None,
+        removed: Optional[int] = None,
+        added: Optional[int] = None,
+    ) -> TestCaseResult:
+        """Build a test case result for a row diff test. If the number of differences is less than the threshold,
+        the test will pass, otherwise it will fail. The result will contain the number of added, removed, and changed
+        rows, as well as the total number of differences.
+
+        Args:
+            threshold: The maximum number of differences allowed before the test fails
+            total_diffs: The total number of differences between the tables
+            changed: The number of rows that have been changed
+            removed: The number of rows that have been removed
+            added: The number of rows that have been added
+
+        Returns:
+            TestCaseResult: The result of the row diff test
+        """
+        return TestCaseResult(
+            timestamp=self.execution_date,  # type: ignore
+            testCaseStatus=self.get_test_case_status(
+                (threshold or total_diffs) == 0 or total_diffs < threshold
+            ),
+            result=f"Found {total_diffs} different rows which is more than the threshold of {threshold}",
+            failedRows=total_diffs,
+            validateColumns=False,
+            testResultValue=[
+                TestResultValue(name="removedRows", value=str(removed)),
+                TestResultValue(name="addedRows", value=str(added)),
+                TestResultValue(name="changedRows", value=str(changed)),
+                TestResultValue(name="diffCount", value=str(total_diffs)),
+            ],
+        )
+
+    def _validate_dialects(self):
+        for name, param in [
+            ("table1.serviceUrl", self.runtime_params.table1.serviceUrl),
+            ("table2.serviceUrl", self.runtime_params.table2.serviceUrl),
+        ]:
+            dialect = urlparse(param).scheme
+            if dialect not in SUPPORTED_DIALECTS:
+                raise UnsupportedDialectError(name, dialect)
+
+    def get_column_diff(self) -> Optional[TestCaseResult]:
+        """Get the column diff between the two tables. If there are no differences, return None."""
+        removed, added = self.get_changed_added_columns(
+            self.runtime_params.table1.columns, self.runtime_params.table2.columns
+        )
+        changed = self.get_incomparable_columns()
+        if removed or added or changed:
+            return self.column_validation_result(
+                removed,
+                added,
+                changed,
+            )
+        return None
+
+    @staticmethod
+    def get_changed_added_columns(
+        left: List[Column], right: List[Column]
+    ) -> Optional[Tuple[List[str], List[str]]]:
+        """Given a list of columns from two tables, return the columns that are removed and added.
+
+        Args:
+            left: List of columns from the first table
+            right: List of columns from the second table
+
+        Returns:
+            A tuple of lists containing the removed and added columns or None if there are no differences
+        """
+        removed: List[str] = []
+        added: List[str] = []
+        right_columns_dict: Dict[str, Column] = {c.name.root: c for c in right}
+        for column in left:
+            table2_column = right_columns_dict.get(column.name.root)
+            if table2_column is None:
+                removed.append(column.name.root)
+                continue
+            del right_columns_dict[column.name.root]
+        added.extend(right_columns_dict.keys())
+        return removed, added
+
+    def column_validation_result(
+        self, removed: List[str], added: List[str], changed: List[str]
+    ) -> TestCaseResult:
+        """Build the result for a column validation result. Messages will only be added
+        for non-empty categories. Values will be populated reported for all categories.
+
+        Args:
+            removed: List of removed columns
+            added: List of added columns
+            changed: List of changed columns
+
+        Returns:
+            TestCaseResult: The result of the column validation with a meaningful message
+        """
+        message = (
+            f"Tables have {sum(map(len, [removed, added, changed]))} different columns:"
+        )
+        if removed:
+            message += f"\n  Removed columns: {','.join(removed)}\n"
+        if added:
+            message += f"\n  Added columns: {','.join(added)}\n"
+        if changed:
+            message += "\n  Changed columns:"
+            for col in changed:
+                col1 = next(
+                    c for c in self.runtime_params.table1.columns if c.name.root == col
+                )
+                col2 = next(
+                    c for c in self.runtime_params.table2.columns if c.name.root == col
+                )
+                message += (
+                    f"\n    {col}: {col1.dataType.value} -> {col2.dataType.value}"
+                )
+        return TestCaseResult(
+            timestamp=self.execution_date,  # type: ignore
+            testCaseStatus=TestCaseStatus.Failed,
+            result=message,
+            testResultValue=[
+                TestResultValue(name="removedColumns", value=str(len(removed))),
+                TestResultValue(name="addedColumns", value=str(len(added))),
+                TestResultValue(name="changedColumns", value=str(len(changed))),
+            ],
+        )
+
+    def calculate_diffs_with_limit(
+        self, diff_iter: Iterable[Tuple[str, Tuple[str, ...]]], limit: int
+    ) -> int:
+        """Given an iterator of diffs like
+        - ('+', (...))
+        - ('-', (...))
+        ...
+        Calculate the total diffs by combining diffs for the same key. This gives an accurate count of the total diffs
+        as opposed to self.calculate_diff_num(diff_list)just counting the number of diffs in the list.
+
+        Args:
+            diff_iter: iterator returned from the data_diff algorithm
+
+        Returns:
+            int: accurate count of the total diffs up to the limit
+
+        """
+        len_key_columns = len(self.runtime_params.keyColumns)
+        key_set = set()
+        # combine diffs on same key to "!"
+        for _, values in diff_iter:
+            k = values[:len_key_columns]
+            if k in key_set:
+                continue
+            key_set.add(k)
+            if len(key_set) > limit:
+                len(key_set)
+        return len(key_set)
+
+    def safe_table_diff_iterator(self) -> DiffResultWrapper:
+        """A safe iterator object which properly closes the diff object when the generator is exhausted.
+        Otherwise the data_diff library will continue to hold the connection open and eventually
+        raise a KeyError.
+        """
+        gen = self.get_table_diff()
+        try:
+            yield from gen
+        finally:
+            try:
+                gen.diff.close()
+            except KeyError as ex:
+                if str(ex) == "2":
+                    # This is a known issue in data_diff where the diff object is closed
+                    pass