openmetadata-ingestion 1.9.9.1__py3-none-any.whl → 1.9.11.0__py3-none-any.whl

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/table/base/tableCustomSQLQuery.py

@@ -87,11 +87,13 @@ 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:
@@ -99,17 +101,23 @@ class BaseTableCustomSQLQueryValidator(BaseTestValidator):
             result_value = len_rows
 
         if self.test_case.computePassedFailedRowCount:
-            row_count = self.get_row_count()
+            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=result_value,
+            failed_rows=failed_rows,
+            passed_rows=passed_rows,
         )
 
     @abstractmethod
@@ -132,3 +140,100 @@ class BaseTableCustomSQLQueryValidator(BaseTestValidator):
             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/sqlalchemy/tableCustomSQLQuery.py

@@ -13,10 +13,13 @@
 Validator for table custom SQL Query test case
 """
 
-from typing import Optional, cast
+from typing import Optional, Tuple, cast
 
+import sqlparse
 from sqlalchemy import text
 from sqlalchemy.sql import func, select
+from sqlparse.sql import Statement, Token, Where
+from sqlparse.tokens import Keyword
 
 from metadata.data_quality.validations.mixins.sqa_validator_mixin import (
     SQAValidatorMixin,
@@ -33,11 +36,244 @@ from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.orm.functions.table_metric_computer import TableMetricComputer
 from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.helpers import is_safe_sql_query
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class TableCustomSQLQueryValidator(BaseTableCustomSQLQueryValidator, SQAValidatorMixin):
     """Validator for table custom SQL Query test case"""
 
+    def _replace_where_clause(
+        self, sql_query: str, partition_expression: str
+    ) -> Optional[str]:
+        """Replace or add WHERE clause in SQL query using sqlparse.
+
+        This method properly handles:
+        - Queries with existing WHERE clause (replaces it)
+        - Queries without WHERE clause (adds it)
+        - Complex queries with joins, subqueries, CTEs
+        - Preserves GROUP BY, ORDER BY, LIMIT, etc.
+
+        Args:
+            sql_query: Original SQL query
+            partition_expression: New WHERE condition (without WHERE keyword)
+
+        Returns:
+            Modified SQL query with partition_expression as WHERE clause
+        """
+        parsed = sqlparse.parse(sql_query)
+        if not parsed or len(parsed) == 0:
+            return None
+
+        statement: Statement = parsed[0]
+        tokens = list(statement.tokens)
+
+        where_idx, where_end_idx, insert_before_idx = self._find_clause_positions(
+            tokens
+        )
+        new_tokens = self._build_new_tokens(
+            tokens, where_idx, where_end_idx, insert_before_idx, partition_expression
+        )
+
+        return "".join(str(token) for token in new_tokens)
+
+    def _find_clause_positions(
+        self, tokens: list
+    ) -> Tuple[Optional[int], Optional[int], Optional[int]]:
+        """Find positions of WHERE clause and insertion points in token list.
+
+        Args:
+            tokens: List of parsed SQL tokens
+
+        Returns:
+            Tuple of (where_idx, where_end_idx, insert_before_idx)
+        """
+        where_idx = None
+        where_end_idx = None
+        insert_before_idx = None
+        paren_depth = 0
+
+        for i, token in enumerate(tokens):
+            paren_depth = self._update_parentheses_depth(token, paren_depth)
+
+            if isinstance(token, Where) and paren_depth == 0:
+                where_idx = i
+                where_end_idx = i + 1
+                break
+
+            if self._should_insert_before_token(token, insert_before_idx, paren_depth):
+                insert_before_idx = i
+
+        return where_idx, where_end_idx, insert_before_idx
+
+    def _update_parentheses_depth(self, token: Token, current_depth: int) -> int:
+        """Update parentheses depth based on token content.
+
+        Args:
+            token: SQL token to analyze
+            current_depth: Current parentheses depth
+
+        Returns:
+            Updated parentheses depth
+        """
+        if token.ttype is None and hasattr(token, "tokens"):
+            paren_count = str(token).count("(") - str(token).count(")")
+            return current_depth + paren_count
+        elif token.value == "(":
+            return current_depth + 1
+        elif token.value == ")":
+            return current_depth - 1
+        return current_depth
+
+    def _should_insert_before_token(
+        self, token: Token, insert_before_idx: Optional[int], paren_depth: int
+    ) -> bool:
+        """Check if WHERE clause should be inserted before this token.
+
+        Args:
+            token: SQL token to check
+            insert_before_idx: Current insertion index (None if not set)
+            paren_depth: Current parentheses depth
+
+        Returns:
+            True if WHERE should be inserted before this token
+        """
+        if insert_before_idx is not None or paren_depth != 0:
+            return False
+
+        if token.ttype is not Keyword:
+            return False
+
+        clause_keywords = {
+            "GROUP BY",
+            "ORDER BY",
+            "HAVING",
+            "LIMIT",
+            "OFFSET",
+            "UNION",
+            "EXCEPT",
+            "INTERSECT",
+        }
+
+        return any(keyword in token.value.upper() for keyword in clause_keywords)
+
+    def _build_new_tokens(
+        self,
+        tokens: list,
+        where_idx: Optional[int],
+        where_end_idx: Optional[int],
+        insert_before_idx: Optional[int],
+        partition_expression: str,
+    ) -> list:
+        """Build new token list with WHERE clause inserted or replaced.
+
+        Args:
+            tokens: Original token list
+            where_idx: Index of existing WHERE clause (None if not found)
+            where_end_idx: End index of existing WHERE clause
+            insert_before_idx: Index to insert WHERE before (None if append)
+            partition_expression: WHERE condition expression
+
+        Returns:
+            New list of tokens with WHERE clause
+        """
+        if where_idx is not None:
+            return self._replace_existing_where(
+                tokens, where_idx, where_end_idx, partition_expression
+            )
+        elif insert_before_idx is not None:
+            return self._insert_where_before_clause(
+                tokens, insert_before_idx, partition_expression
+            )
+        else:
+            return self._append_where_clause(tokens, partition_expression)
+
+    def _replace_existing_where(
+        self,
+        tokens: list,
+        where_idx: int,
+        where_end_idx: int,
+        partition_expression: str,
+    ) -> list:
+        """Replace existing WHERE clause with new expression.
+
+        Args:
+            tokens: Original token list
+            where_idx: Index of WHERE clause to replace
+            where_end_idx: End index of WHERE clause
+            partition_expression: New WHERE condition
+
+        Returns:
+            Token list with replaced WHERE clause
+        """
+        original_where = str(tokens[where_idx])
+        trailing_whitespace = self._extract_trailing_whitespace(original_where)
+
+        return (
+            tokens[:where_idx]
+            + [
+                Token(Keyword, "WHERE"),
+                Token(None, f" {partition_expression}{trailing_whitespace}"),
+            ]
+            + tokens[where_end_idx:]
+        )
+
+    def _insert_where_before_clause(
+        self, tokens: list, insert_before_idx: int, partition_expression: str
+    ) -> list:
+        """Insert WHERE clause before specified token index.
+
+        Args:
+            tokens: Original token list
+            insert_before_idx: Index to insert WHERE clause before
+            partition_expression: WHERE condition expression
+
+        Returns:
+            Token list with WHERE clause inserted
+        """
+        return (
+            tokens[:insert_before_idx]
+            + [Token(Keyword, "WHERE"), Token(None, f" {partition_expression} ")]
+            + tokens[insert_before_idx:]
+        )
+
+    def _append_where_clause(self, tokens: list, partition_expression: str) -> list:
+        """Append WHERE clause to end of token list.
+
+        Args:
+            tokens: Original token list
+            partition_expression: WHERE condition expression
+
+        Returns:
+            Token list with WHERE clause appended
+        """
+        return tokens + [
+            Token(None, " "),
+            Token(Keyword, "WHERE"),
+            Token(None, f" {partition_expression}"),
+        ]
+
+    def _extract_trailing_whitespace(self, where_clause: str) -> str:
+        """Extract trailing whitespace from WHERE clause string.
+
+        Args:
+            where_clause: Original WHERE clause string
+
+        Returns:
+            Trailing whitespace string
+        """
+        if not where_clause.split():
+            return ""
+
+        last_word = where_clause.split()[-1]
+        where_content_end = where_clause.rfind(last_word) + len(last_word)
+
+        if where_content_end < len(where_clause):
+            return where_clause[where_content_end:]
+
+        return ""
+
     def run_validation(self) -> TestCaseResult:
         """Run validation for the given test case
 
@@ -85,12 +321,39 @@ class TableCustomSQLQueryValidator(BaseTableCustomSQLQueryValidator, SQAValidato
             None,
         )
         if partition_expression:
-            stmt = (
-                select(func.count())
-                .select_from(self.runner.table)
-                .filter(text(partition_expression))
+            custom_sql = self.get_test_case_param_value(
+                self.test_case.parameterValues,  # type: ignore
+                "sqlExpression",
+                str,
             )
-            return self.runner.session.execute(stmt).scalar()
+
+            if custom_sql:
+                modified_query = self._replace_where_clause(
+                    custom_sql, partition_expression
+                )
+                if modified_query is None:
+                    return None
+                count_query = f"SELECT COUNT(*) FROM ({modified_query}) AS test_results"
+
+                try:
+                    result = self.runner.session.execute(text(count_query)).scalar()
+                    return result
+                except Exception as exc:
+                    logger.error(
+                        "Failed to execute custom SQL with partition expression. "
+                        f"Query: {count_query}\n"
+                        f"Error: {exc}\n",
+                        exc_info=True,
+                    )
+                    self.runner.session.rollback()
+                    raise exc
+            else:
+                stmt = (
+                    select(func.count())
+                    .select_from(self.runner.table)
+                    .filter(text(partition_expression))
+                )
+                return self.runner.session.execute(stmt).scalar()
 
         self.runner = cast(QueryRunner, self.runner)
         dialect = self.runner._session.get_bind().dialect.name

metadata/generated/schema/analytics/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00

metadata/generated/schema/analytics/basic.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/basic.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportData.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00

metadata/generated/schema/analytics/reportDataType/aggregatedCostAnalysisReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/aggregatedCostAnalysisReportData.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/entityReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/entityReportData.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/rawCostAnalysisReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/rawCostAnalysisReportData.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/webAnalyticEntityViewReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/webAnalyticEntityViewReportData.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/webAnalyticUserActivityReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/webAnalyticUserActivityReportData.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEvent.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEventData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEventData.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEventType/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00

metadata/generated/schema/analytics/webAnalyticEventType/customEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEventType/customEvent.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEventType/pageViewEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEventType/pageViewEvent.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00

metadata/generated/schema/api/addGlossaryToAssetsRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/addGlossaryToAssetsRequest.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/addTagToAssetsRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/addTagToAssetsRequest.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/analytics/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00

metadata/generated/schema/api/analytics/createWebAnalyticEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/analytics/createWebAnalyticEvent.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/automations/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00

metadata/generated/schema/api/automations/createWorkflow.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/automations/createWorkflow.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/bulkAssets.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/bulkAssets.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/classification/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00

metadata/generated/schema/api/classification/createClassification.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/classification/createClassification.json
-#   timestamp: 2025-09-23T17:36:03+00:00
+#   timestamp: 2025-09-25T15:53:40+00:00
 
 from __future__ import annotations