openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

airflow_provider_openmetadata/lineage/backend.py

@@ -76,16 +76,18 @@ class OpenMetadataLineageBackend(LineageBackend):
                     "timeout": config.timeout,
                     "retry": config.retry,
                     "retry_wait": config.retry_wait,
+                    "retry_codes": config.retry_codes,
                 }.items()
                 if value
             }
 
             if additional_client_config_arguments:
                 dag.log.info(
-                    f"Using custom timeout={config.timeout}, retry={config.retry}, retry_wait={config.retry_wait}"
+                    f"Using custom timeout={config.timeout}, retry={config.retry}, retry_wait={config.retry_wait}, retry_codes={config.retry_codes}"
                 )
             metadata = OpenMetadata(
-                config.metadata_config, additional_client_config_arguments
+                config.metadata_config,
+                additional_client_config_arguments=additional_client_config_arguments,
             )
 
             runner = AirflowLineageRunner(

airflow_provider_openmetadata/lineage/config/loader.py

@@ -14,7 +14,7 @@ OpenMetadata Airflow Lineage Backend
 """
 import json
 import os
-from typing import Optional
+from typing import List, Optional
 
 from airflow.configuration import AirflowConfigParser
 from pydantic import BaseModel
@@ -37,6 +37,7 @@ class AirflowLineageConfig(BaseModel):
     timeout: Optional[int] = None
     retry: Optional[int] = None
     retry_wait: Optional[int] = None
+    retry_codes: Optional[List[int]] = None
 
 
 def parse_airflow_config(
@@ -58,6 +59,12 @@ def parse_airflow_config(
         timeout=int(conf.get(LINEAGE, "timeout", fallback=0)) or None,
         retry=int(conf.get(LINEAGE, "retry", fallback=0)) or None,
         retry_wait=int(conf.get(LINEAGE, "retry_wait", fallback=0)) or None,
+        retry_codes=[
+            int(code)
+            for code in (conf.get(LINEAGE, "retry_codes", fallback="") or "").split(",")
+            if code.strip()
+        ]
+        or None,  # input e.g. 503,504
         metadata_config=OpenMetadataConnection(
             hostPort=conf.get(
                 LINEAGE,

metadata/cli/common.py

@@ -20,7 +20,6 @@ from metadata.workflow.base import BaseWorkflow
 def execute_workflow(workflow: BaseWorkflow, config_dict: Dict[str, Any]) -> None:
     """Execute the workflow and raise if needed"""
     workflow.execute()
-    workflow.print_status()
     workflow.stop()
     if config_dict.get("workflowConfig", {}).get("raiseOnError", True):
         workflow.raise_from_status()

metadata/clients/aws_client.py

@@ -36,6 +36,7 @@ class AWSServices(Enum):
     GLUE = "glue"
     SAGEMAKER = "sagemaker"
     KINESIS = "kinesis"
+    FIREHOSE = "firehose"
     QUICKSIGHT = "quicksight"
     ATHENA = "athena"
     RDS = "rds"
@@ -233,6 +234,9 @@ class AWSClient:
     def get_kinesis_client(self):
         return self.get_client(AWSServices.KINESIS.value)
 
+    def get_firehose_client(self):
+        return self.get_client(AWSServices.FIREHOSE.value)
+
     def get_quicksight_client(self):
         return self.get_client(AWSServices.QUICKSIGHT.value)
 

metadata/data_quality/interface/test_suite_interface.py

@@ -27,7 +27,7 @@ from metadata.data_quality.validations.runtime_param_setter.param_setter_factory
 )
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
-from metadata.generated.schema.tests.basic import TestCaseResult
+from metadata.generated.schema.tests.basic import TestCaseResult, TestCaseStatus
 from metadata.generated.schema.tests.testCase import TestCase
 from metadata.generated.schema.tests.testDefinition import TestDefinition
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -134,10 +134,16 @@ class TestSuiteInterface(ABC):
         try:
             return validator.run_validation()
         except Exception as err:
-            logger.error(
+            message = (
                 f"Error executing {test_case.testDefinition.fullyQualifiedName} - {err}"
             )
-            raise RuntimeError(err)
+            logger.exception(message)
+            return validator.get_test_case_result_object(
+                validator.execution_date,
+                TestCaseStatus.Aborted,
+                message,
+                [],
+            )
 
     def _get_table_config(self):
         """Get the sampling configuration for the data quality tests"""

metadata/data_quality/validations/models.py

@@ -4,8 +4,13 @@ from typing import List, Optional, Union
 
 from pydantic import BaseModel
 
-from metadata.generated.schema.entity.data.table import Column, TableProfilerConfig
+from metadata.generated.schema.entity.data.table import (
+    Column,
+    Table,
+    TableProfilerConfig,
+)
 from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseConnection,
     DatabaseServiceType,
 )
 from metadata.ingestion.models.custom_pydantic import CustomSecretStr
@@ -27,3 +32,8 @@ class TableDiffRuntimeParameters(BaseModel):
     extraColumns: List[str]
     whereClause: Optional[str]
     table_profile_config: Optional[TableProfilerConfig]
+
+
+class TableCustomSQLQueryRuntimeParameters(BaseModel):
+    conn_config: DatabaseConnection
+    entity: Table

metadata/data_quality/validations/runtime_param_setter/base_diff_params_setter.py

@@ -3,16 +3,24 @@
 from typing import List, Optional, Set, Type, Union
 from urllib.parse import urlparse
 
+from sqlalchemy.engine import make_url
+
 from metadata.data_quality.validations.models import Column, TableParameter
 from metadata.generated.schema.entity.data.table import Table
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseService,
+    DatabaseServiceType,
+)
 from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.ingestion.connections.connection import BaseConnection
 from metadata.ingestion.source.connections import get_connection
-from metadata.profiler.orm.registry import Dialects
+from metadata.profiler.orm.registry import Dialects, PythonDialects
 from metadata.utils import fqn
 from metadata.utils.collections import CaseInsensitiveList
 from metadata.utils.importer import get_module_dir, import_from_module
+from metadata.utils.logger import test_suite_logger
+
+logger = test_suite_logger()
 
 
 # TODO: Refactor to avoid the circular import that makes us unable to use the BaseSpec class and the helper methods.
@@ -68,7 +76,9 @@ class BaseTableParameter:
         """
         return TableParameter(
             database_service_type=service.serviceType,
-            path=self.get_data_diff_table_path(entity.fullyQualifiedName.root),
+            path=self.get_data_diff_table_path(
+                entity.fullyQualifiedName.root, service.serviceType
+            ),
             serviceUrl=self.get_data_diff_url(
                 service,
                 entity.fullyQualifiedName.root,
@@ -85,7 +95,9 @@ class BaseTableParameter:
         )
 
     @staticmethod
-    def get_data_diff_table_path(table_fqn: str) -> str:
+    def get_data_diff_table_path(
+        table_fqn: str, service_type: DatabaseServiceType
+    ) -> str:
         """Get the data diff table path.
 
         Args:
@@ -95,17 +107,32 @@ class BaseTableParameter:
             str
         """
         _, _, schema, table = fqn.split(table_fqn)
+        try:
+            dialect = PythonDialects[service_type.name].value
+            if dialect in (Dialects.Oracle):
+                url = make_url(f"{dialect}://")
+                dialect_instance = url.get_dialect()()
+                table = dialect_instance.denormalize_name(name=table)
+                schema = dialect_instance.denormalize_name(name=schema)
+        except Exception as e:
+            logger.debug(
+                f"[Data Diff]: Error denormalizing table and schema names. Skipping denormalization\n{e}"
+            )
         return fqn._build(  # pylint: disable=protected-access
             "___SERVICE___", "__DATABASE__", schema, table
         ).replace("___SERVICE___.__DATABASE__.", "")
 
-    @staticmethod
+    @classmethod
     def _get_service_connection_config(
+        cls,
         service_connection_config,
     ) -> Optional[Union[str, dict]]:
         """
         Get the connection dictionary for the service.
         """
+        if not service_connection_config:
+            return None
+
         service_spec_patch = ServiceSpecPatch(
             ServiceType.Database, service_connection_config.type.value.lower()
         )
@@ -127,8 +154,8 @@ class BaseTableParameter:
                 else None
             )
 
-    @staticmethod
     def get_data_diff_url(
+        self,
         db_service: DatabaseService,
         table_fqn,
         override_url: Optional[Union[str, dict]] = None,
@@ -144,9 +171,7 @@ class BaseTableParameter:
             str: The url for the data diff service
         """
         source_url = (
-            BaseTableParameter._get_service_connection_config(
-                db_service.connection.config
-            )
+            self._get_service_connection_config(db_service.connection.config)
             if not override_url
             else override_url
         )
@@ -157,12 +182,13 @@ class BaseTableParameter:
         url = urlparse(source_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
-        )
+        _, database, schema, _ = fqn.split(table_fqn)  # pylint: disable=unused-variable
         # path needs to include the database AND schema in some of the connectors
         if hasattr(db_service.connection.config, "supportsDatabase"):
-            kwargs["path"] = f"/{database}"
+            if kwargs["scheme"] in {Dialects.UnityCatalog, Dialects.Databricks}:
+                kwargs["query"] = f"catalog={database}"
+            else:
+                kwargs["path"] = f"/{database}"
         if kwargs["scheme"] in {Dialects.MSSQL, Dialects.Snowflake, Dialects.Trino}:
             kwargs["path"] = f"/{database}/{schema}"
         return url._replace(**kwargs).geturl()

metadata/data_quality/validations/runtime_param_setter/param_setter_factory.py

@@ -19,9 +19,15 @@ from typing import Dict, Set, Type
 from metadata.data_quality.validations.runtime_param_setter.param_setter import (
     RuntimeParameterSetter,
 )
+from metadata.data_quality.validations.runtime_param_setter.table_custom_sql_query_params_setter import (
+    TableCustomSQLQueryParamsSetter,
+)
 from metadata.data_quality.validations.runtime_param_setter.table_diff_params_setter import (
     TableDiffParamsSetter,
 )
+from metadata.data_quality.validations.table.sqlalchemy.tableCustomSQLQuery import (
+    TableCustomSQLQueryValidator,
+)
 from metadata.data_quality.validations.table.sqlalchemy.tableDiff import (
     TableDiffValidator,
 )
@@ -60,6 +66,9 @@ class RuntimeParameterSetterFactory:
         """Set"""
         self._setter_map: Dict[str, Set[Type[RuntimeParameterSetter]]] = {
             validator_name(TableDiffValidator): {TableDiffParamsSetter},
+            validator_name(TableCustomSQLQueryValidator): {
+                TableCustomSQLQueryParamsSetter
+            },
         }
 
     def get_runtime_param_setters(

metadata/data_quality/validations/runtime_param_setter/table_custom_sql_query_params_setter.py

@@ -0,0 +1,31 @@
+#  Copyright 2024 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  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 TableCustomSQLQueryParamsSetter class."""
+
+from metadata.data_quality.validations.models import (
+    TableCustomSQLQueryRuntimeParameters,
+)
+from metadata.data_quality.validations.runtime_param_setter.param_setter import (
+    RuntimeParameterSetter,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+
+
+class TableCustomSQLQueryParamsSetter(RuntimeParameterSetter):
+    """Set runtime parameters for a the table custom sql query test."""
+
+    def get_parameters(self, test_case) -> TableCustomSQLQueryRuntimeParameters:
+        return TableCustomSQLQueryRuntimeParameters(
+            conn_config=DatabaseConnection(
+                config=self.service_connection_config,
+            ),
+            entity=self.table_entity,
+        )

metadata/data_quality/validations/runtime_param_setter/table_diff_params_setter.py

@@ -15,7 +15,6 @@ from typing import List, Optional, Set
 from metadata.data_quality.validations import utils
 from metadata.data_quality.validations.models import Column, TableDiffRuntimeParameters
 from metadata.data_quality.validations.runtime_param_setter.base_diff_params_setter import (
-    BaseTableParameter,
     ServiceSpecPatch,
 )
 from metadata.data_quality.validations.runtime_param_setter.param_setter import (
@@ -49,32 +48,40 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
         }
 
     def get_parameters(self, test_case) -> TableDiffRuntimeParameters:
-        service_spec_patch = ServiceSpecPatch(
-            ServiceType.Database, self.service_connection_config.type.value.lower()
-        )
-        cls = service_spec_patch.get_data_diff_class()()
-
         service1: DatabaseService = self.ometa_client.get_by_id(
             DatabaseService, self.table_entity.service.id, nullable=False
         )
 
-        service1_url = BaseTableParameter._get_service_connection_config(
-            self.service_connection_config
-        )
-
         table2_fqn = self.get_parameter(test_case, "table2")
         if table2_fqn is None:
             raise ValueError("table2 not set")
+
         table2: Table = self.ometa_client.get_by_name(
             Table, fqn=table2_fqn, nullable=False
         )
         service2: DatabaseService = self.ometa_client.get_by_id(
             DatabaseService, table2.service.id, nullable=False
         )
+
+        service_spec_patch_table_1 = ServiceSpecPatch(
+            ServiceType.Database, service1.connection.config.type.value.lower()
+        )
+        data_diff_class_1 = service_spec_patch_table_1.get_data_diff_class()()
+        service_spec_patch_table_2 = ServiceSpecPatch(
+            ServiceType.Database, service2.connection.config.type.value.lower()
+        )
+        data_diff_class_2 = service_spec_patch_table_2.get_data_diff_class()()
+
+        service1_url = data_diff_class_1._get_service_connection_config(
+            service1.connection.config
+        )
         service2_url = (
             self.get_parameter(test_case, "service2Url") or service1_url
             if table2.service == self.table_entity.service
-            else None
+            else data_diff_class_2._get_service_connection_config(
+                service2.connection.config
+            )
+            or None
         )
 
         key_columns = self.get_key_columns(test_case)
@@ -93,7 +100,7 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
 
         return TableDiffRuntimeParameters(
             table_profile_config=self.table_entity.tableProfilerConfig,
-            table1=cls.get(
+            table1=data_diff_class_1.get(
                 service1,
                 self.table_entity,
                 key_columns,
@@ -101,7 +108,7 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
                 case_sensitive_columns,
                 service1_url,
             ),
-            table2=cls.get(
+            table2=data_diff_class_2.get(
                 service2,
                 table2,
                 key_columns,

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

@@ -87,24 +87,153 @@ class BaseTableCustomSQLQueryValidator(BaseTestValidator):
                 [TestResultValue(name=RESULT_ROW_COUNT, value=None)],
             )
         len_rows = rows if isinstance(rows, int) else len(rows)
-        if evaluate_threshold(
+        test_passed = evaluate_threshold(
             threshold,
             operator,
             len_rows,
-        ):
+        )
+
+        if test_passed:
             status = TestCaseStatus.Success
             result_value = len_rows
         else:
             status = TestCaseStatus.Failed
             result_value = len_rows
 
+        if self.test_case.computePassedFailedRowCount:
+            row_count = self._get_total_row_count_if_needed()
+            passed_rows, failed_rows = self._calculate_passed_failed_rows(
+                test_passed, operator, threshold, len_rows, row_count
+            )
+        else:
+            passed_rows = None
+            failed_rows = None
+            row_count = None
+
         return self.get_test_case_result_object(
             self.execution_date,
             status,
-            f"Found {result_value} row(s). Test query is expected to return {threshold} row.",
+            f"Found {result_value} row(s). Test query is expected to return {operator} {threshold} row(s).",
             [TestResultValue(name=RESULT_ROW_COUNT, value=str(result_value))],
+            row_count=row_count,
+            failed_rows=failed_rows,
+            passed_rows=passed_rows,
         )
 
     @abstractmethod
     def _run_results(self, sql_expression: str, strategy: Strategy = Strategy.ROWS):
         raise NotImplementedError
+
+    @abstractmethod
+    def compute_row_count(self):
+        """Compute row count for the given column
+
+        Raises:
+            NotImplementedError:
+        """
+        raise NotImplementedError
+
+    def get_row_count(self) -> int:
+        """Get row count
+
+        Returns:
+            Tuple[int, int]:
+        """
+        return self.compute_row_count()
+
+    def _get_total_row_count_if_needed(self) -> int:
+        """Get total row count if computePassedFailedRowCount is enabled"""
+        return self.get_row_count()
+
+    def _calculate_passed_failed_rows(
+        self,
+        test_passed: bool,
+        operator: str,
+        threshold: int,
+        len_rows: int,
+        row_count: int,
+    ) -> tuple[int, int]:
+        """Calculate passed and failed rows based on test result and operator
+
+        Args:
+            test_passed: Whether the test passed
+            operator: Comparison operator (>, >=, <, <=, ==)
+            threshold: Expected threshold value
+            len_rows: Number of rows returned by the test query
+            row_count: Total number of rows in the table (or None)
+
+        Returns:
+            Tuple of (passed_rows, failed_rows)
+        """
+        if test_passed:
+            return self._calculate_passed_rows_success(operator, len_rows, row_count)
+        return self._calculate_passed_rows_failure(
+            operator, threshold, len_rows, row_count
+        )
+
+    def _calculate_passed_rows_success(
+        self, operator: str, len_rows: int, row_count: int
+    ) -> tuple[int, int]:
+        """Calculate passed/failed rows when test passed"""
+        if operator in (">", ">="):
+            passed_rows = len_rows
+            failed_rows = (row_count - len_rows) if row_count else 0
+        elif operator in ("<", "<="):
+            passed_rows = row_count - len_rows
+            failed_rows = len_rows
+        elif operator == "==":
+            passed_rows = len_rows
+            failed_rows = row_count - len_rows
+        else:
+            passed_rows = len_rows
+            failed_rows = 0
+
+        return max(0, passed_rows), max(0, failed_rows)
+
+    def _calculate_passed_rows_failure(
+        self, operator: str, threshold: int, len_rows: int, row_count: int
+    ) -> tuple[int, int]:
+        """Calculate passed/failed rows when test failed"""
+        if operator in (">", ">="):
+            return self._calculate_greater_than_failure(len_rows, row_count)
+        if operator in ("<", "<="):
+            return self._calculate_less_than_failure(len_rows, row_count)
+        if operator == "==":
+            return self._calculate_equal_failure(threshold, len_rows, row_count)
+
+        failed_rows = row_count if row_count else len_rows
+        return 0, max(0, failed_rows)
+
+    def _calculate_greater_than_failure(
+        self, len_rows: int, row_count: int
+    ) -> tuple[int, int]:
+        """Calculate rows for > or >= operator failure (expected more rows)"""
+        passed_rows = len_rows
+        failed_rows = (row_count - len_rows) if row_count else 0
+        return max(0, passed_rows), max(0, failed_rows)
+
+    def _calculate_less_than_failure(
+        self, len_rows: int, row_count: int
+    ) -> tuple[int, int]:
+        """Calculate rows for < or <= operator failure (expected fewer rows)"""
+        failed_rows = len_rows
+        passed_rows = row_count - failed_rows
+
+        return max(0, passed_rows), max(0, failed_rows)
+
+    def _calculate_equal_failure(
+        self, threshold: int, len_rows: int, row_count: int
+    ) -> tuple[int, int]:
+        """Calculate rows for == operator failure (expected exact count)"""
+        if row_count:
+            if len_rows > threshold:
+                failed_rows = len_rows - threshold
+                passed_rows = row_count - failed_rows
+            else:
+                failed_rows = row_count - len_rows
+                passed_rows = len_rows
+        else:
+            failed_rows = abs(len_rows - threshold)
+            passed_rows = 0
+
+        return max(0, passed_rows), max(0, failed_rows)

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

@@ -13,6 +13,8 @@
 Validator for table custom SQL Query test case
 """
 
+from typing import List, Optional
+
 from metadata.data_quality.validations.mixins.pandas_validator_mixin import (
     PandasValidatorMixin,
 )
@@ -39,3 +41,39 @@ class TableCustomSQLQueryValidator(
                 if len(runner.query(sql_expression))
             ]
         )
+
+    def compute_row_count(self) -> Optional[int]:
+        """Compute row count for the given column
+
+        Returns:
+            Optional[int]: Total number of rows across all dataframes
+        """
+        runner: List["DataFrame"] = self.runner  # type: ignore
+
+        if not runner:
+            return None
+
+        total_rows = 0
+        partition_expression = next(
+            (
+                param.value
+                for param in self.test_case.parameterValues
+                if param.name == "partitionExpression"
+            ),
+            None,
+        )
+        for dataframe in runner:
+            if dataframe is not None:
+                if partition_expression:
+                    try:
+                        total_rows += len(dataframe.query(partition_expression))
+                    except Exception as e:
+                        logger.error(
+                            "Error executing partition expression, "
+                            f"expression may be invalid: {partition_expression} - {e}"
+                        )
+                        return None
+                else:
+                    total_rows += len(dataframe.index)
+
+        return total_rows if total_rows > 0 else None