openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/data_quality/validations/column/sqlalchemy/columnValueMinToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValueMinToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValueStdDevToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValueStdDevToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesMissingCount.py

@@ -42,7 +42,7 @@ class ColumnValuesMissingCountValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesSumToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValuesSumToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeAtExpectedLocation.py

@@ -37,7 +37,7 @@ class ColumnValuesToBeAtExpectedLocationValidator(
     def _fetch_data(self, columns: List[str]) -> Iterator:
         """Fetch data from the runner object"""
         self.runner = cast(QueryRunner, self.runner)
-        inspection = inspect(self.runner.table)
+        inspection = inspect(self.runner.dataset)
         table_columns: List[Column] = inspection.c if inspection is not None else []
         cols = [col for col in table_columns if col.name in columns]
         for col in cols:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeBetween.py

@@ -39,7 +39,7 @@ class ColumnValuesToBeBetweenValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeInSet.py

@@ -39,7 +39,7 @@ class ColumnValuesToBeInSetValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeNotInSet.py

@@ -39,7 +39,7 @@ class ColumnValuesToBeNotInSetValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeNotNull.py

@@ -42,7 +42,7 @@ class ColumnValuesToBeNotNullValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:

metadata/data_quality/validations/column/sqlalchemy/columnValuesToBeUnique.py

@@ -17,7 +17,6 @@ from typing import Optional
 
 from sqlalchemy import Column, inspect
 from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.orm.util import AliasedClass
 
 from metadata.data_quality.validations.column.base.columnValuesToBeUnique import (
     BaseColumnValuesToBeUniqueValidator,
@@ -41,7 +40,7 @@ class ColumnValuesToBeUniqueValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column) -> Optional[int]:
@@ -53,12 +52,7 @@ class ColumnValuesToBeUniqueValidator(
         """
         count = Metrics.COUNT.value(column).fn()
         unique_count = Metrics.UNIQUE_COUNT.value(column).query(
-            sample=self.runner._sample  # pylint: disable=protected-access
-            if isinstance(
-                self.runner._sample,  # pylint: disable=protected-access
-                AliasedClass,
-            )
-            else self.runner.table,
+            sample=self.runner.dataset,
             session=self.runner._session,  # pylint: disable=protected-access
         )  # type: ignore
 

metadata/data_quality/validations/column/sqlalchemy/columnValuesToMatchRegex.py

@@ -43,7 +43,7 @@ class ColumnValuesToMatchRegexValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(

metadata/data_quality/validations/column/sqlalchemy/columnValuesToNotMatchRegex.py

@@ -43,7 +43,7 @@ class ColumnValuesToNotMatchRegexValidator(
         """
         return self.get_column_name(
             self.test_case.entityLink.root,
-            inspect(self.runner.table).c,
+            inspect(self.runner.dataset).c,
         )
 
     def _run_results(self, metric: Metrics, column: Column, **kwargs) -> Optional[int]:

metadata/data_quality/validations/models.py

@@ -4,13 +4,17 @@ from typing import List, Optional
 
 from pydantic import BaseModel
 
-from metadata.generated.schema.entity.data.table import Column
+from metadata.generated.schema.entity.data.table import Column, TableProfilerConfig
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseServiceType,
+)
 
 
 class TableParameter(BaseModel):
     serviceUrl: str
     path: str
     columns: List[Column]
+    database_service_type: DatabaseServiceType
 
 
 class TableDiffRuntimeParameters(BaseModel):
@@ -19,3 +23,4 @@ class TableDiffRuntimeParameters(BaseModel):
     keyColumns: List[str]
     extraColumns: List[str]
     whereClause: Optional[str]
+    table_profile_config: Optional[TableProfilerConfig]

metadata/data_quality/validations/runtime_param_setter/param_setter.py

@@ -15,14 +15,13 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.entity.data.table import Table
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.profiler.processor.sampler.sqlalchemy.sampler import SQASampler
 
 
 class RuntimeParameterSetter(ABC):
     """
     Set runtime parameters for a test case. This can be useful for resolving parameters based on the
     state of OpenMetadata entities.
-    You can see an example implemnentation in `TableDiffParamsSetter`.
+    You can see an example implementation in `TableDiffParamsSetter`.
     """
 
     def __init__(
@@ -30,7 +29,7 @@ class RuntimeParameterSetter(ABC):
         ometa_client: OpenMetadata,
         service_connection_config,
         table_entity: Table,
-        sampler: SQASampler,
+        sampler,
     ):
         self.ometa_client = ometa_client
         self.service_connection_config = service_connection_config

metadata/data_quality/validations/runtime_param_setter/param_setter_factory.py

@@ -13,8 +13,8 @@ Module that defines the RuntimeParameterFactory class.
 This class is responsible for creating instances of the RuntimeParameterSetter 
 based on the test case.
 """
-
-from typing import Optional
+import sys
+from typing import Dict, Set, Type
 
 from metadata.data_quality.validations.runtime_param_setter.param_setter import (
     RuntimeParameterSetter,
@@ -22,7 +22,35 @@ from metadata.data_quality.validations.runtime_param_setter.param_setter import
 from metadata.data_quality.validations.runtime_param_setter.table_diff_params_setter import (
     TableDiffParamsSetter,
 )
+from metadata.data_quality.validations.table.sqlalchemy.tableDiff import (
+    TableDiffValidator,
+)
+from metadata.generated.schema.entity.data.table import Table
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.sampler.sampler_interface import SamplerInterface
+
+
+def removesuffix(s: str, suffix: str) -> str:
+    """A custom implementation of removesuffix for python versions < 3.9
+
+    Args:
+        s (str): The string to remove the suffix from
+        suffix (str): The suffix to remove
+
+    Returns:
+        str: The string with the suffix removed
+    """
+    if sys.version_info >= (3, 9):
+        return s.removesuffix(suffix)
+    if s.endswith(suffix):
+        return s[: -len(suffix)]
+    return s
+
+
+def validator_name(test_case_class: Type) -> str:
+    return removesuffix(
+        test_case_class.__name__[0].lower() + test_case_class.__name__[1:], "Validator"
+    )
 
 
 class RuntimeParameterSetterFactory:
@@ -30,25 +58,25 @@ class RuntimeParameterSetterFactory:
 
     def __init__(self) -> None:
         """Set"""
-        self._setter_map = {
-            TableDiffParamsSetter: {"tableDiff"},
+        self._setter_map: Dict[str, Set[Type[RuntimeParameterSetter]]] = {
+            validator_name(TableDiffValidator): {TableDiffParamsSetter},
         }
 
-    def get_runtime_param_setter(
+    def get_runtime_param_setters(
         self,
         name: str,
         ometa: OpenMetadata,
         service_connection_config,
-        table_entity,
-        sampler,
-    ) -> Optional[RuntimeParameterSetter]:
+        table_entity: Table,
+        sampler: SamplerInterface,
+    ) -> Set[RuntimeParameterSetter]:
         """Get the runtime parameter setter"""
-        for setter_cls, validator_names in self._setter_map.items():
-            if name in validator_names:
-                return setter_cls(
-                    ometa,
-                    service_connection_config,
-                    table_entity,
-                    sampler,
-                )
-        return None
+        return {
+            setter(
+                ometa,
+                service_connection_config,
+                table_entity,
+                sampler,
+            )
+            for setter in self._setter_map.get(name, set())
+        }

metadata/data_quality/validations/runtime_param_setter/table_diff_params_setter.py

@@ -10,7 +10,7 @@
 #  limitations under the License.
 """Module that defines the TableDiffParamsSetter class."""
 from ast import literal_eval
-from typing import List, Optional
+from typing import List, Optional, Set
 from urllib.parse import urlparse
 
 from metadata.data_quality.validations import utils
@@ -75,9 +75,13 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
             DatabaseService, table2.service.id, nullable=False
         )
         key_columns = self.get_key_columns(test_case)
-        extra_columns = self.get_extra_columns(key_columns, test_case)
+        extra_columns = self.get_extra_columns(
+            key_columns, test_case, self.table_entity.columns, table2.columns
+        )
         return TableDiffRuntimeParameters(
+            table_profile_config=self.table_entity.tableProfilerConfig,
             table1=TableParameter(
+                database_service_type=service1.serviceType,
                 path=self.get_data_diff_table_path(
                     self.table_entity.fullyQualifiedName.root
                 ),
@@ -94,6 +98,7 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
                 ),
             ),
             table2=TableParameter(
+                database_service_type=service2.serviceType,
                 path=self.get_data_diff_table_path(table2_fqn),
                 serviceUrl=self.get_data_diff_url(
                     service2,
@@ -108,18 +113,19 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
                     case_sensitive=case_sensitive_columns,
                 ),
             ),
-            keyColumns=key_columns,
-            extraColumns=extra_columns,
+            keyColumns=list(key_columns),
+            extraColumns=list(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
+            if not (
+                self.sampler.partition_details
+                and self.sampler.partition_details.enablePartitioning
+            )
             else self.sampler.get_partitioned_query().whereclause.compile(
                 compile_kwargs={"literal_binds": True}
             )
@@ -130,21 +136,25 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
         return " AND ".join(where_clauses)
 
     def get_extra_columns(
-        self, key_columns: List[str], test_case
-    ) -> Optional[List[str]]:
+        self,
+        key_columns: Set[str],
+        test_case,
+        left_columns: List[Column],
+        right_columns: List[Column],
+    ) -> Optional[Set[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
+            return set(extra_columns)
         if extra_columns_param is None:
             extra_columns_param = []
-            for column in self.table_entity.columns:
+            for column in left_columns + right_columns:
                 if column.name.root not in key_columns:
                     extra_columns_param.insert(0, column.name.root)
-        return extra_columns_param
+        return set(extra_columns_param)
 
-    def get_key_columns(self, test_case) -> List[str]:
+    def get_key_columns(self, test_case) -> Set[str]:
         key_columns_param = self.get_parameter(test_case, "keyColumns", "[]")
         key_columns: List[str] = literal_eval(key_columns_param)
         if key_columns:
@@ -163,13 +173,13 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
                 "Could not find primary key or unique constraint columns.\n",
                 "Specify 'keyColumns' to explicitly set the columns to use as keys.",
             )
-        return key_columns
+        return set(key_columns)
 
     @staticmethod
     def filter_relevant_columns(
         columns: List[Column],
-        key_columns: List[str],
-        extra_columns: List[str],
+        key_columns: Set[str],
+        extra_columns: Set[str],
         case_sensitive: bool,
     ) -> List[Column]:
         validated_columns = (

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

@@ -37,7 +37,13 @@ class TableCustomSQLQueryValidator(BaseTableCustomSQLQueryValidator, SQAValidato
                 text(sql_expression)
             )
             if strategy == Strategy.COUNT:
-                return cursor.scalar()
+                result = cursor.scalar()
+                if not isinstance(result, int):
+                    raise ValueError(
+                        f"When using COUNT strategy, the result must be an integer. Received: {type(result)}\n"
+                        "Example: SELECT COUNT(*) FROM table_name WHERE my_value IS NOT NULL"
+                    )
+                return result
             return cursor.fetchall()
         except Exception as exc:
             self.runner._session.rollback()  # pylint: disable=protected-access