From 650823c77e43b69455c9781aae42aa64190f4060 Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Sun, 21 Sep 2025 17:06:06 +0530 Subject: [PATCH 01/10] feat: add data quality rule conditions and incremental run support --- pyatlan/errors.py | 14 ++ pyatlan/model/assets/core/alpha__d_q_rule.py | 228 +++++++++++++++++-- pyatlan/model/dq_rule_conditions.py | 127 +++++++++++ pyatlan/model/enums.py | 18 ++ pyatlan/model/structs.py | 6 + 5 files changed, 377 insertions(+), 16 deletions(-) create mode 100644 pyatlan/model/dq_rule_conditions.py diff --git a/pyatlan/errors.py b/pyatlan/errors.py index 5918a27e0..064f5db35 100644 --- a/pyatlan/errors.py +++ b/pyatlan/errors.py @@ -647,6 +647,20 @@ class ErrorCode(Enum): "Ensure your product instance has a valid `data_product_assets_d_s_l` value before making the request.", InvalidRequestError, ) + DQ_RULE_TYPE_NOT_SUPPORTED = ( + 400, + "ATLAN-PYTHON-400-073", + "Rule type '{0}' does not support {1}.", + "Choose a rule type that supports the specified template configuration.", + InvalidRequestError, + ) + DQ_RULE_CONDITIONS_INVALID = ( + 400, + "ATLAN-PYTHON-400-074", + "Invalid rule conditions: {0}.", + "Ensure your rule conditions are valid and match the expected format.", + InvalidRequestError, + ) AUTHENTICATION_PASSTHROUGH = ( 401, "ATLAN-PYTHON-401-000", diff --git a/pyatlan/model/assets/core/alpha__d_q_rule.py b/pyatlan/model/assets/core/alpha__d_q_rule.py index c5a4cd95f..1df2674e4 100644 --- a/pyatlan/model/assets/core/alpha__d_q_rule.py +++ b/pyatlan/model/assets/core/alpha__d_q_rule.py @@ -23,6 +23,7 @@ alpha_DQSourceSyncStatus, ) from pyatlan.model.fields.atlan_fields import ( + BooleanField, KeywordField, KeywordTextField, NumericField, @@ -160,8 +161,10 @@ def column_level_rule_creator( alert_priority: alpha_DQRuleAlertPriority, threshold_compare_operator: Optional[ alpha_DQRuleThresholdCompareOperator - ] = alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL, + ] = None, threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = False, ) -> alpha_DQRule: validate_required_fields( [ @@ -169,7 +172,6 @@ def column_level_rule_creator( "rule_type", "asset", "column", - "threshold_compare_operator", "threshold_value", "alert_priority", ], @@ -178,11 +180,20 @@ def column_level_rule_creator( rule_type, asset, column, - threshold_compare_operator, threshold_value, alert_priority, ], ) + template_config = client.dq_template_config_cache.get_template_config(rule_type) + threshold_compare_operator = ( + alpha_DQRule.Attributes._validate_template_features( + rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator, + ) + ) attributes = alpha_DQRule.Attributes.creator( client=client, @@ -197,6 +208,8 @@ def column_level_rule_creator( dimension=None, custom_sql=None, description=None, + rule_conditions=rule_conditions, + row_scope_filtering_enabled=row_scope_filtering_enabled, ) return cls(attributes=attributes) @@ -216,6 +229,8 @@ def updater( custom_sql: Optional[str] = None, rule_name: Optional[str] = None, description: Optional[str] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = False, ) -> SelfAsset: from pyatlan.model.fluent_search import FluentSearch @@ -237,6 +252,7 @@ def updater( .include_on_results(alpha_DQRule.USER_DESCRIPTION) .include_on_results(alpha_DQRule.ALPHADQ_RULE_DIMENSION) .include_on_results(alpha_DQRule.ALPHADQ_RULE_CONFIG_ARGUMENTS) + .include_on_results(alpha_DQRule.ALPHADQ_RULE_ROW_SCOPE_FILTERING_ENABLED) .include_on_results(alpha_DQRule.ALPHADQ_RULE_SOURCE_SYNC_STATUS) .include_on_results(alpha_DQRule.ALPHADQ_RULE_STATUS) ).to_request() @@ -255,6 +271,9 @@ def updater( retrieved_dimension = search_result.alpha_dq_rule_dimension # type: ignore[attr-defined] retrieved_column = search_result.alpha_dq_rule_base_column # type: ignore[attr-defined] retrieved_alert_priority = search_result.alpha_dq_rule_alert_priority # type: ignore[attr-defined] + retrieved_row_scope_filtering_enabled = ( + search_result.alpha_dq_rule_row_scope_filtering_enabled + ) # type: ignore[attr-defined] retrieved_description = search_result.user_description retrieved_asset = search_result.alpha_dq_rule_base_dataset # type: ignore[attr-defined] retrieved_template_rule_name = search_result.alpha_dq_rule_template_name # type: ignore[attr-defined] @@ -281,35 +300,52 @@ def updater( else None ) # type: ignore[attr-defined] + retrieved_rule_type = retrieved_template_rule_name + template_config = client.dq_template_config_cache.get_template_config( + retrieved_rule_type + ) + validated_threshold_operator = ( + alpha_DQRule.Attributes._validate_template_features( + retrieved_rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator or retrieved_threshold_compare_operator, + ) + ) + config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw( is_alert_enabled=True, custom_sql=custom_sql or retrieved_custom_sql, display_name=rule_name or retrieved_rule_name, dimension=dimension or retrieved_dimension, - compare_operator=threshold_compare_operator - or retrieved_threshold_compare_operator, + compare_operator=validated_threshold_operator, threshold_value=threshold_value or retrieved_threshold_value, threshold_unit=threshold_unit or retrieved_threshold_unit, column=retrieved_column, dq_priority=alert_priority or retrieved_alert_priority, description=description or retrieved_description, + rule_conditions=rule_conditions, + row_scope_filtering_enabled=row_scope_filtering_enabled, ) attr_dq = cls.Attributes( name="", alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments( alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject( - alpha_dq_rule_threshold_compare_operator=threshold_compare_operator - or retrieved_threshold_compare_operator, + alpha_dq_rule_threshold_compare_operator=validated_threshold_operator, alpha_dq_rule_threshold_value=threshold_value or retrieved_threshold_value, alpha_dq_rule_threshold_unit=threshold_unit or retrieved_threshold_unit, ), alpha_dq_rule_config_arguments_raw=config_arguments_raw, + alpha_dq_rule_config_rule_conditions=rule_conditions, ), alpha_dq_rule_base_dataset_qualified_name=retrieved_asset.qualified_name, alpha_dq_rule_alert_priority=alert_priority or retrieved_alert_priority, + alpha_dq_rule_row_scope_filtering_enabled=row_scope_filtering_enabled + or retrieved_row_scope_filtering_enabled, alpha_dq_rule_base_dataset=retrieved_asset, qualified_name=qualified_name, alpha_dq_rule_dimension=dimension or retrieved_dimension, @@ -387,6 +423,12 @@ def __setattr__(self, name, value): """ List of unique reference dataset's qualified names related to this rule. """ + ALPHADQ_RULE_ROW_SCOPE_FILTERING_ENABLED: ClassVar[BooleanField] = BooleanField( + "alpha_dqRuleRowScopeFilteringEnabled", "alpha_dqRuleRowScopeFilteringEnabled" + ) + """ + Flag to enable row scope filtering for the rule + """ ALPHADQ_RULE_SOURCE_SYNC_STATUS: ClassVar[KeywordField] = KeywordField( "alpha_dqRuleSourceSyncStatus", "alpha_dqRuleSourceSyncStatus" ) @@ -517,6 +559,7 @@ def __setattr__(self, name, value): "alpha_dq_rule_base_column_qualified_name", "alpha_dq_rule_reference_dataset_qualified_names", "alpha_dq_rule_reference_column_qualified_names", + "alpha_dq_rule_row_scope_filtering_enabled", "alpha_dq_rule_source_sync_status", "alpha_dq_rule_source_sync_error_code", "alpha_dq_rule_source_sync_error_message", @@ -611,6 +654,24 @@ def alpha_dq_rule_reference_column_qualified_names( alpha_dq_rule_reference_column_qualified_names ) + @property + def alpha_dq_rule_row_scope_filtering_enabled(self) -> Optional[bool]: + return ( + None + if self.attributes is None + else self.attributes.alpha_dq_rule_row_scope_filtering_enabled + ) + + @alpha_dq_rule_row_scope_filtering_enabled.setter + def alpha_dq_rule_row_scope_filtering_enabled( + self, alpha_dq_rule_row_scope_filtering_enabled: Optional[bool] + ): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.alpha_dq_rule_row_scope_filtering_enabled = ( + alpha_dq_rule_row_scope_filtering_enabled + ) + @property def alpha_dq_rule_source_sync_status(self) -> Optional[alpha_DQSourceSyncStatus]: return ( @@ -944,6 +1005,9 @@ class Attributes(DataQuality.Attributes): alpha_dq_rule_reference_column_qualified_names: Optional[Set[str]] = Field( default=None, description="" ) + alpha_dq_rule_row_scope_filtering_enabled: Optional[bool] = Field( + default=None, description="" + ) alpha_dq_rule_source_sync_status: Optional[alpha_DQSourceSyncStatus] = Field( default=None, description="" ) @@ -1001,6 +1065,125 @@ class Attributes(DataQuality.Attributes): default=None, description="" ) # relationship + @staticmethod + def _get_template_config_value( + config_value: str, property_name: str = None, value_key: str = "default" + ): + if not config_value: + return None + + try: + config_json = json.loads(config_value) + + if property_name: + properties = config_json.get("properties", {}) + field = properties.get(property_name, {}) + return field.get(value_key) + else: + return config_json.get(value_key) + except (json.JSONDecodeError, KeyError): + return None + + @staticmethod + def _validate_template_features( + rule_type: str, + rule_conditions: Optional[str], + row_scope_filtering_enabled: Optional[bool], + template_config: Optional[dict], + threshold_compare_operator: Optional[ + alpha_DQRuleThresholdCompareOperator + ] = None, + ) -> alpha_DQRuleThresholdCompareOperator: + if not template_config or not template_config.get("config"): + return + + config = template_config["config"] + + if ( + rule_conditions + and config.alpha_dq_rule_template_config_rule_conditions is None + ): + raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters( + rule_type, "rule conditions" + ) + + if row_scope_filtering_enabled: + advanced_settings = ( + config.alpha_dq_rule_template_advanced_settings or "" + ) + if "alpha_dqRuleRowScopeFilteringEnabled" not in str(advanced_settings): + raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters( + rule_type, "row scope filtering" + ) + + if rule_conditions: + allowed_rule_conditions = ( + alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_rule_conditions, + None, + "enum", + ) + ) + if allowed_rule_conditions: + try: + rule_conditions_json = json.loads(rule_conditions) + conditions = rule_conditions_json.get("conditions", []) + if len(conditions) != 1: + raise ErrorCode.DQ_RULE_CONDITIONS_INVALID.exception_with_parameters( + f"exactly one condition required, found {len(conditions)}" + ) + condition_type = conditions[0].get("type") + except json.JSONDecodeError: + condition_type = rule_conditions + + if condition_type not in allowed_rule_conditions: + raise ErrorCode.DQ_RULE_CONDITIONS_INVALID.exception_with_parameters( + f"condition type '{condition_type}' not supported, allowed: {allowed_rule_conditions}" + ) + + if threshold_compare_operator is None: + return alpha_DQRuleThresholdCompareOperator.EQUAL + elif ( + threshold_compare_operator + != alpha_DQRuleThresholdCompareOperator.EQUAL + ): + raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( + f"threshold_compare_operator={threshold_compare_operator.value}", + "threshold_compare_operator", + "EQUAL when rule_conditions are provided", + ) + + if threshold_compare_operator is not None: + allowed_operators = alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_threshold_object, + "alpha_dqRuleTemplateConfigThresholdCompareOperator", + "enum", + ) + if ( + allowed_operators + and threshold_compare_operator.value not in allowed_operators + ): + raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( + f"threshold_compare_operator={threshold_compare_operator.value}", + "threshold_compare_operator", + f"must be one of {allowed_operators}", + ) + elif threshold_compare_operator is None: + default_value = alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_threshold_object, + "alpha_dqRuleTemplateConfigThresholdCompareOperator", + "default", + ) + if default_value: + threshold_compare_operator = alpha_DQRuleThresholdCompareOperator( + default_value + ) + + return ( + threshold_compare_operator + or alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL + ) + @staticmethod def _generate_config_arguments_raw( *, @@ -1014,6 +1197,8 @@ def _generate_config_arguments_raw( column: Optional[Asset] = None, dq_priority: alpha_DQRuleAlertPriority, description: Optional[str] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = None, ) -> str: config = { "isAlertEnabled": is_alert_enabled, @@ -1042,6 +1227,16 @@ def _generate_config_arguments_raw( if dimension is not None: config["alpha_dqRuleTemplateConfigDimension"] = dimension + if rule_conditions is not None: + config["alpha_dqRuleTemplateConfigRuleConditions"] = json.loads( + rule_conditions + ) + + if row_scope_filtering_enabled is not None: + config[ + "alpha_dqRuleTemplateAdvancedSettings.alpha_dqRuleRowScopeFilteringEnabled" + ] = row_scope_filtering_enabled + return json.dumps(config) @staticmethod @@ -1083,6 +1278,8 @@ def creator( dimension: Optional[alpha_DQDimension] = None, custom_sql: Optional[str] = None, description: Optional[str] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = False, ) -> alpha_DQRule.Attributes: template_config = client.dq_template_config_cache.get_template_config( rule_type @@ -1100,16 +1297,11 @@ def creator( if threshold_unit is None: config = template_config.get("config") if config is not None: - threashold_object = ( - config.alpha_dq_rule_template_config_threshold_object - ) - threashold_object_json = json.loads(threashold_object) - properties = threashold_object_json.get("properties", {}) - threshold_unit_field = properties.get( - "alpha_dqRuleTemplateConfigThresholdUnit", {} + threshold_unit = alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_threshold_object, + "alpha_dqRuleTemplateConfigThresholdUnit", + "default", ) - default_value = threshold_unit_field.get("default") - threshold_unit = default_value config_arguments_raw = ( alpha_DQRule.Attributes._generate_config_arguments_raw( @@ -1123,6 +1315,8 @@ def creator( column=column, dq_priority=alert_priority, description=description, + rule_conditions=rule_conditions, + row_scope_filtering_enabled=row_scope_filtering_enabled, ) ) @@ -1135,9 +1329,11 @@ def creator( alpha_dq_rule_threshold_unit=threshold_unit, ), alpha_dq_rule_config_arguments_raw=config_arguments_raw, + alpha_dq_rule_config_rule_conditions=rule_conditions, ), alpha_dq_rule_base_dataset_qualified_name=asset.qualified_name, alpha_dq_rule_alert_priority=alert_priority, + alpha_dq_rule_row_scope_filtering_enabled=row_scope_filtering_enabled, alpha_dq_rule_source_sync_status=alpha_DQSourceSyncStatus.IN_PROGRESS, alpha_dq_rule_status=alpha_DQRuleStatus.ACTIVE, alpha_dq_rule_base_dataset=asset, diff --git a/pyatlan/model/dq_rule_conditions.py b/pyatlan/model/dq_rule_conditions.py new file mode 100644 index 000000000..f86b26b7b --- /dev/null +++ b/pyatlan/model/dq_rule_conditions.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 Atlan Pte. Ltd. +from __future__ import annotations + +import json +from typing import Any, Dict, List, Optional, Union + +from pydantic.v1 import Field + +from pyatlan.errors import ErrorCode +from pyatlan.model.enums import alpha_dqRuleTemplateConfigRuleConditions +from pyatlan.model.structs import AtlanObject +from pyatlan.utils import validate_required_fields, validate_type + + +class DQCondition(AtlanObject): + """ + Data quality rule condition. + + :param type: the condition type enum value + :param value: value of type str, int, or list depending on condition type + :param min_value: minimum value for range-based conditions + :param max_value: maximum value for range-based conditions + """ + + type: alpha_dqRuleTemplateConfigRuleConditions = Field(description="") + value: Optional[Union[str, int, List[str], Dict[str, Any]]] = Field( + default=None, description="" + ) + min_value: Optional[int] = Field(default=None, description="") + max_value: Optional[int] = Field(default=None, description="") + + def __init__( + self, + type: alpha_dqRuleTemplateConfigRuleConditions, + value: Optional[Union[str, int, List[str], Dict[str, Any]]] = None, + min_value: Optional[int] = None, + max_value: Optional[int] = None, + **kwargs, + ): + super().__init__( + type=type, value=value, min_value=min_value, max_value=max_value, **kwargs + ) + + if self.type == alpha_dqRuleTemplateConfigRuleConditions.STRING_LENGTH_BETWEEN: + validate_required_fields( + ["min_value", "max_value"], [self.min_value, self.max_value] + ) + if self.min_value < 0 or self.max_value < 0: + raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( + f"min_value={self.min_value}, max_value={self.max_value}", + "min_value, max_value", + "non-negative integers", + ) + if self.min_value > self.max_value: + raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( + f"min_value={self.min_value}, max_value={self.max_value}", + "min_value, max_value", + "min_value <= max_value", + ) + else: + validate_required_fields(["value"], [self.value]) + if self.type in [ + alpha_dqRuleTemplateConfigRuleConditions.IN_LIST, + alpha_dqRuleTemplateConfigRuleConditions.NOT_IN_LIST, + ]: + validate_type("value", list, self.value) + elif self.type in [ + alpha_dqRuleTemplateConfigRuleConditions.REGEX_MATCH, + alpha_dqRuleTemplateConfigRuleConditions.REGEX_NOT_MATCH, + ]: + validate_type("value", str, self.value) + + def to_dict(self) -> Dict[str, Any]: + """Convert the condition to a dictionary representation.""" + result = {"type": self.type.value} + + if self.type == alpha_dqRuleTemplateConfigRuleConditions.STRING_LENGTH_BETWEEN: + result["value"] = {"minValue": self.min_value, "maxValue": self.max_value} + else: + result["value"] = self.value + + return result + + +class DQRuleConditionsBuilder: + """Builder for data quality rule conditions.""" + + def __init__(self): + self._conditions: List[DQCondition] = [] + + def add_condition( + self, + type: alpha_dqRuleTemplateConfigRuleConditions, + value: Optional[Union[str, int, List[str]]] = None, + min_value: Optional[int] = None, + max_value: Optional[int] = None, + ) -> "DQRuleConditionsBuilder": + """ + Add a condition to the builder. + + :param type: the condition type enum value + :param value: value of type str, int, or list depending on condition type + :param min_value: minimum value for range-based conditions + :param max_value: maximum value for range-based conditions + :returns: the builder for method chaining + """ + self._conditions.append( + DQCondition( + type=type, value=value, min_value=min_value, max_value=max_value + ) + ) + return self + + def build(self) -> str: + """Build the conditions as a JSON string.""" + if not self._conditions: + raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( + "empty conditions list", "conditions", "at least one condition" + ) + + return json.dumps( + {"conditions": [condition.to_dict() for condition in self._conditions]} + ) + + +DQCondition.update_forward_refs() diff --git a/pyatlan/model/enums.py b/pyatlan/model/enums.py index 48a2c7d9b..0fbc8249c 100644 --- a/pyatlan/model/enums.py +++ b/pyatlan/model/enums.py @@ -2505,6 +2505,24 @@ class alpha_DQSourceSyncStatus(str, Enum): WAITING_FOR_SCHEDULE = "WAITING_FOR_SCHEDULE" +class alpha_dqRuleTemplateConfigRuleConditions(str, Enum): + # String Length conditions + STRING_LENGTH_BETWEEN = "STRING_LENGTH_BETWEEN" + STRING_LENGTH_EQUALS = "STRING_LENGTH_EQUALS" + STRING_LENGTH_GREATER_THAN = "STRING_LENGTH_GREATER_THAN" + STRING_LENGTH_LESS_THAN = "STRING_LENGTH_LESS_THAN" + STRING_LENGTH_GREATER_THAN_EQUALS = "STRING_LENGTH_GREATER_THAN_EQUALS" + STRING_LENGTH_LESS_THAN_EQUALS = "STRING_LENGTH_LESS_THAN_EQUALS" + + # Regex conditions + REGEX_MATCH = "REGEX_MATCH" + REGEX_NOT_MATCH = "REGEX_NOT_MATCH" + + # List conditions + IN_LIST = "IN_LIST" + NOT_IN_LIST = "NOT_IN_LIST" + + # ************************************** # CODE BELOW IS GENERATED NOT MODIFY ** # ************************************** diff --git a/pyatlan/model/structs.py b/pyatlan/model/structs.py index 61ee3e781..6d03bd9d1 100644 --- a/pyatlan/model/structs.py +++ b/pyatlan/model/structs.py @@ -177,6 +177,9 @@ class alpha_DQRuleTemplateConfig(AtlanObject): alpha_dq_rule_template_config_user_description: Optional[str] = Field( default=None, description="" ) + alpha_dq_rule_template_config_rule_conditions: Optional[str] = Field( + default=None, description="" + ) alpha_dq_rule_template_advanced_settings: Optional[str] = Field( default=None, description="" ) @@ -508,6 +511,9 @@ class alpha_DQRuleConfigArguments(AtlanObject): alpha_dq_rule_config_arguments_raw: Optional[str] = Field( default=None, description="" ) + alpha_dq_rule_config_rule_conditions: Optional[str] = Field( + default=None, description="" + ) class alpha_DQRuleRangeForTesting(AtlanObject): From c527116f734785cca39e499b2dd350bd05e89b55 Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Sun, 21 Sep 2025 17:06:41 +0530 Subject: [PATCH 02/10] chore: update jinja templates --- .../methods/asset/alpha__d_q_rule.jinja2 | 45 +++++- .../methods/attribute/alpha__d_q_rule.jinja2 | 142 ++++++++++++++++-- .../methods/imports/alpha__d_q_rule.jinja2 | 1 + 3 files changed, 172 insertions(+), 16 deletions(-) diff --git a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 index 66587297e..106023b44 100644 --- a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 +++ b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 @@ -113,8 +113,10 @@ alert_priority: alpha_DQRuleAlertPriority, threshold_compare_operator: Optional[ alpha_DQRuleThresholdCompareOperator - ] = alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL, + ] = None, threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = False, ) -> alpha_DQRule: validate_required_fields( [ @@ -122,7 +124,6 @@ "rule_type", "asset", "column", - "threshold_compare_operator", "threshold_value", "alert_priority", ], @@ -131,11 +132,18 @@ rule_type, asset, column, - threshold_compare_operator, threshold_value, alert_priority, ], ) + template_config = client.dq_template_config_cache.get_template_config(rule_type) + threshold_compare_operator = alpha_DQRule.Attributes._validate_template_features( + rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator, + ) attributes = alpha_DQRule.Attributes.creator( client=client, @@ -150,6 +158,8 @@ dimension=None, custom_sql=None, description=None, + rule_conditions=rule_conditions, + row_scope_filtering_enabled=row_scope_filtering_enabled, ) return cls(attributes=attributes) @@ -169,6 +179,8 @@ custom_sql: Optional[str] = None, rule_name: Optional[str] = None, description: Optional[str] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = False, ) -> SelfAsset: from pyatlan.model.fluent_search import FluentSearch @@ -190,6 +202,7 @@ .include_on_results(alpha_DQRule.USER_DESCRIPTION) .include_on_results(alpha_DQRule.ALPHADQ_RULE_DIMENSION) .include_on_results(alpha_DQRule.ALPHADQ_RULE_CONFIG_ARGUMENTS) + .include_on_results(alpha_DQRule.ALPHADQ_RULE_ROW_SCOPE_FILTERING_ENABLED) .include_on_results(alpha_DQRule.ALPHADQ_RULE_SOURCE_SYNC_STATUS) .include_on_results(alpha_DQRule.ALPHADQ_RULE_STATUS) ).to_request() @@ -208,6 +221,9 @@ retrieved_dimension = search_result.alpha_dq_rule_dimension # type: ignore[attr-defined] retrieved_column = search_result.alpha_dq_rule_base_column # type: ignore[attr-defined] retrieved_alert_priority = search_result.alpha_dq_rule_alert_priority # type: ignore[attr-defined] + retrieved_row_scope_filtering_enabled = ( + search_result.alpha_dq_rule_row_scope_filtering_enabled + ) # type: ignore[attr-defined] retrieved_description = search_result.user_description retrieved_asset = search_result.alpha_dq_rule_base_dataset # type: ignore[attr-defined] retrieved_template_rule_name = search_result.alpha_dq_rule_template_name # type: ignore[attr-defined] @@ -234,35 +250,50 @@ else None ) # type: ignore[attr-defined] + retrieved_rule_type = retrieved_template_rule_name + template_config = client.dq_template_config_cache.get_template_config( + retrieved_rule_type + ) + validated_threshold_operator = alpha_DQRule.Attributes._validate_template_features( + retrieved_rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator or retrieved_threshold_compare_operator, + ) + config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw( is_alert_enabled=True, custom_sql=custom_sql or retrieved_custom_sql, display_name=rule_name or retrieved_rule_name, dimension=dimension or retrieved_dimension, - compare_operator=threshold_compare_operator - or retrieved_threshold_compare_operator, + compare_operator=validated_threshold_operator, threshold_value=threshold_value or retrieved_threshold_value, threshold_unit=threshold_unit or retrieved_threshold_unit, column=retrieved_column, dq_priority=alert_priority or retrieved_alert_priority, description=description or retrieved_description, + rule_conditions=rule_conditions, + row_scope_filtering_enabled=row_scope_filtering_enabled, ) attr_dq = cls.Attributes( name="", alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments( alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject( - alpha_dq_rule_threshold_compare_operator=threshold_compare_operator - or retrieved_threshold_compare_operator, + alpha_dq_rule_threshold_compare_operator=validated_threshold_operator, alpha_dq_rule_threshold_value=threshold_value or retrieved_threshold_value, alpha_dq_rule_threshold_unit=threshold_unit or retrieved_threshold_unit, ), alpha_dq_rule_config_arguments_raw=config_arguments_raw, + alpha_dq_rule_config_rule_conditions=rule_conditions, ), alpha_dq_rule_base_dataset_qualified_name=retrieved_asset.qualified_name, alpha_dq_rule_alert_priority=alert_priority or retrieved_alert_priority, + alpha_dq_rule_row_scope_filtering_enabled=row_scope_filtering_enabled + or retrieved_row_scope_filtering_enabled, alpha_dq_rule_base_dataset=retrieved_asset, qualified_name=qualified_name, alpha_dq_rule_dimension=dimension or retrieved_dimension, diff --git a/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 b/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 index 96b4ba00f..1a5c2d9d1 100644 --- a/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 +++ b/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 @@ -1,3 +1,114 @@ + @staticmethod + def _get_template_config_value( + config_value: str, property_name: str = None, value_key: str = "default" + ): + if not config_value: + return None + + try: + config_json = json.loads(config_value) + + if property_name: + properties = config_json.get("properties", {}) + field = properties.get(property_name, {}) + return field.get(value_key) + else: + return config_json.get(value_key) + except (json.JSONDecodeError, KeyError): + return None + + @staticmethod + def _validate_template_features( + rule_type: str, + rule_conditions: Optional[str], + row_scope_filtering_enabled: Optional[bool], + template_config: Optional[dict], + threshold_compare_operator: Optional[ + alpha_DQRuleThresholdCompareOperator + ] = None, + ) -> alpha_DQRuleThresholdCompareOperator: + if not template_config or not template_config.get("config"): + return + + config = template_config["config"] + + if ( + rule_conditions + and config.alpha_dq_rule_template_config_rule_conditions is None + ): + raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters( + rule_type, "rule conditions" + ) + + if row_scope_filtering_enabled: + advanced_settings = config.alpha_dq_rule_template_advanced_settings or "" + if "alpha_dqRuleRowScopeFilteringEnabled" not in str(advanced_settings): + raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters( + rule_type, "row scope filtering" + ) + + if rule_conditions: + allowed_rule_conditions = alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_rule_conditions, None, "enum" + ) + if allowed_rule_conditions: + try: + rule_conditions_json = json.loads(rule_conditions) + conditions = rule_conditions_json.get("conditions", []) + if len(conditions) != 1: + raise ErrorCode.DQ_RULE_CONDITIONS_INVALID.exception_with_parameters( + f"exactly one condition required, found {len(conditions)}" + ) + condition_type = conditions[0].get("type") + except json.JSONDecodeError: + condition_type = rule_conditions + + if condition_type not in allowed_rule_conditions: + raise ErrorCode.DQ_RULE_CONDITIONS_INVALID.exception_with_parameters( + f"condition type '{condition_type}' not supported, allowed: {allowed_rule_conditions}" + ) + + if threshold_compare_operator is None: + return alpha_DQRuleThresholdCompareOperator.EQUAL + elif ( + threshold_compare_operator != alpha_DQRuleThresholdCompareOperator.EQUAL + ): + raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( + f"threshold_compare_operator={threshold_compare_operator.value}", + "threshold_compare_operator", + "EQUAL when rule_conditions are provided", + ) + + if threshold_compare_operator is not None: + allowed_operators = alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_threshold_object, + "alpha_dqRuleTemplateConfigThresholdCompareOperator", + "enum", + ) + if ( + allowed_operators + and threshold_compare_operator.value not in allowed_operators + ): + raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( + f"threshold_compare_operator={threshold_compare_operator.value}", + "threshold_compare_operator", + f"must be one of {allowed_operators}", + ) + elif threshold_compare_operator is None: + default_value = alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_threshold_object, + "alpha_dqRuleTemplateConfigThresholdCompareOperator", + "default", + ) + if default_value: + threshold_compare_operator = alpha_DQRuleThresholdCompareOperator( + default_value + ) + + return ( + threshold_compare_operator + or alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL + ) @staticmethod def _generate_config_arguments_raw( @@ -12,6 +123,8 @@ column: Optional[Asset] = None, dq_priority: alpha_DQRuleAlertPriority, description: Optional[str] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = None, ) -> str: config = { "isAlertEnabled": is_alert_enabled, @@ -40,6 +153,16 @@ if dimension is not None: config["alpha_dqRuleTemplateConfigDimension"] = dimension + if rule_conditions is not None: + config["alpha_dqRuleTemplateConfigRuleConditions"] = json.loads( + rule_conditions + ) + + if row_scope_filtering_enabled is not None: + config[ + "alpha_dqRuleTemplateAdvancedSettings.alpha_dqRuleRowScopeFilteringEnabled" + ] = row_scope_filtering_enabled + return json.dumps(config) @staticmethod @@ -81,6 +204,8 @@ dimension: Optional[alpha_DQDimension] = None, custom_sql: Optional[str] = None, description: Optional[str] = None, + rule_conditions: Optional[str] = None, + row_scope_filtering_enabled: Optional[bool] = False, ) -> alpha_DQRule.Attributes: template_config = client.dq_template_config_cache.get_template_config( rule_type @@ -98,16 +223,11 @@ if threshold_unit is None: config = template_config.get("config") if config is not None: - threashold_object = ( - config.alpha_dq_rule_template_config_threshold_object - ) - threashold_object_json = json.loads(threashold_object) - properties = threashold_object_json.get("properties", {}) - threshold_unit_field = properties.get( - "alpha_dqRuleTemplateConfigThresholdUnit", {} + threshold_unit = alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_threshold_object, + "alpha_dqRuleTemplateConfigThresholdUnit", + "default", ) - default_value = threshold_unit_field.get("default") - threshold_unit = default_value config_arguments_raw = ( alpha_DQRule.Attributes._generate_config_arguments_raw( @@ -121,6 +241,8 @@ column=column, dq_priority=alert_priority, description=description, + rule_conditions=rule_conditions, + row_scope_filtering_enabled=row_scope_filtering_enabled, ) ) @@ -133,9 +255,11 @@ alpha_dq_rule_threshold_unit=threshold_unit, ), alpha_dq_rule_config_arguments_raw=config_arguments_raw, + alpha_dq_rule_config_rule_conditions=rule_conditions, ), alpha_dq_rule_base_dataset_qualified_name=asset.qualified_name, alpha_dq_rule_alert_priority=alert_priority, + alpha_dq_rule_row_scope_filtering_enabled=row_scope_filtering_enabled, alpha_dq_rule_source_sync_status=alpha_DQSourceSyncStatus.IN_PROGRESS, alpha_dq_rule_status=alpha_DQRuleStatus.ACTIVE, alpha_dq_rule_base_dataset=asset, diff --git a/pyatlan/generator/templates/methods/imports/alpha__d_q_rule.jinja2 b/pyatlan/generator/templates/methods/imports/alpha__d_q_rule.jinja2 index 41fa507fb..ccb86d258 100644 --- a/pyatlan/generator/templates/methods/imports/alpha__d_q_rule.jinja2 +++ b/pyatlan/generator/templates/methods/imports/alpha__d_q_rule.jinja2 @@ -8,6 +8,7 @@ from pyatlan.model.enums import ( alpha_DQRuleThresholdCompareOperator, alpha_DQRuleThresholdUnit, alpha_DQSourceSyncStatus, + alpha_dqRuleTemplateConfigRuleConditions, ) from pyatlan.model.structs import ( alpha_DQRuleConfigArguments, From 44fce934c5c613726e2cfbd60c66cac0aed08e00 Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Sun, 21 Sep 2025 20:27:31 +0530 Subject: [PATCH 03/10] feat: add row scope filter column check for incremental rules --- pyatlan/errors.py | 7 +++++ pyatlan/model/assets/core/alpha__d_q_rule.py | 30 ++++++++++++++++++- pyatlan/model/assets/core/asset.py | 31 ++++++++++++++++++++ 3 files changed, 67 insertions(+), 1 deletion(-) diff --git a/pyatlan/errors.py b/pyatlan/errors.py index 064f5db35..a963df373 100644 --- a/pyatlan/errors.py +++ b/pyatlan/errors.py @@ -661,6 +661,13 @@ class ErrorCode(Enum): "Ensure your rule conditions are valid and match the expected format.", InvalidRequestError, ) + DQ_ROW_SCOPE_FILTER_COLUMN_MISSING = ( + 400, + "ATLAN-PYTHON-400-075", + "Row scope filter column not configured for asset '{0}'.", + "please configure the row scope filter column first.", + InvalidRequestError, + ) AUTHENTICATION_PASSTHROUGH = ( 401, "ATLAN-PYTHON-401-000", diff --git a/pyatlan/model/assets/core/alpha__d_q_rule.py b/pyatlan/model/assets/core/alpha__d_q_rule.py index 1df2674e4..5c157e651 100644 --- a/pyatlan/model/assets/core/alpha__d_q_rule.py +++ b/pyatlan/model/assets/core/alpha__d_q_rule.py @@ -41,7 +41,7 @@ if TYPE_CHECKING: from pyatlan.client.atlan import AtlanClient - from pyatlan.model.assets import Column + from pyatlan.model.assets import Asset, Column class alpha_DQRule(DataQuality): @@ -185,6 +185,22 @@ def column_level_rule_creator( ], ) template_config = client.dq_template_config_cache.get_template_config(rule_type) + + asset_for_validation = asset + if row_scope_filtering_enabled and asset.qualified_name: + from pyatlan.model.fluent_search import FluentSearch + + search_request = ( + FluentSearch() + .where(Asset.QUALIFIED_NAME.eq(asset.qualified_name)) + .include_on_results( + Asset.ALPHAASSET_DQ_ROW_SCOPE_FILTER_COLUMN_QUALIFIED_NAME + ) + ).to_request() + results = client.asset.search(search_request) + if results.count == 1: + asset_for_validation = results.current_page()[0] + threshold_compare_operator = ( alpha_DQRule.Attributes._validate_template_features( rule_type, @@ -192,6 +208,7 @@ def column_level_rule_creator( row_scope_filtering_enabled, template_config, threshold_compare_operator, + asset_for_validation, ) ) @@ -311,6 +328,7 @@ def updater( row_scope_filtering_enabled, template_config, threshold_compare_operator or retrieved_threshold_compare_operator, + retrieved_asset, ) ) @@ -1093,6 +1111,7 @@ def _validate_template_features( threshold_compare_operator: Optional[ alpha_DQRuleThresholdCompareOperator ] = None, + asset: Optional[Asset] = None, ) -> alpha_DQRuleThresholdCompareOperator: if not template_config or not template_config.get("config"): return @@ -1116,6 +1135,15 @@ def _validate_template_features( rule_type, "row scope filtering" ) + if asset and not getattr( + asset, + "alpha_asset_d_q_row_scope_filter_column_qualified_name", + None, + ): + raise ErrorCode.DQ_ROW_SCOPE_FILTER_COLUMN_MISSING.exception_with_parameters( + getattr(asset, "qualified_name", "unknown") + ) + if rule_conditions: allowed_rule_conditions = ( alpha_DQRule.Attributes._get_template_config_value( diff --git a/pyatlan/model/assets/core/asset.py b/pyatlan/model/assets/core/asset.py index 694d71b02..4e84d766c 100644 --- a/pyatlan/model/assets/core/asset.py +++ b/pyatlan/model/assets/core/asset.py @@ -1143,6 +1143,15 @@ def __setattr__(self, name, value): """ Internal Popularity score for this asset. """ + ALPHAASSET_DQ_ROW_SCOPE_FILTER_COLUMN_QUALIFIED_NAME: ClassVar[KeywordField] = ( + KeywordField( + "alpha_assetDQRowScopeFilterColumnQualifiedName", + "alpha_assetDQRowScopeFilterColumnQualifiedName", + ) + ) + """ + Qualified name of the column to be used for row scope filter + """ ALPHAASSET_DQ_SCHEDULE_TYPE: ClassVar[KeywordField] = KeywordField( "alpha_assetDQScheduleType", "alpha_assetDQScheduleType" ) @@ -1510,6 +1519,7 @@ def __setattr__(self, name, value): "application_field_qualified_name", "asset_user_defined_type", "asset_internal_popularity_score", + "alpha_asset_d_q_row_scope_filter_column_qualified_name", "alpha_asset_d_q_schedule_type", "alpha_asset_d_q_schedule_crontab", "alpha_asset_d_q_schedule_time_zone", @@ -3561,6 +3571,24 @@ def asset_internal_popularity_score( asset_internal_popularity_score ) + @property + def alpha_asset_d_q_row_scope_filter_column_qualified_name(self) -> Optional[str]: + return ( + None + if self.attributes is None + else self.attributes.alpha_asset_d_q_row_scope_filter_column_qualified_name + ) + + @alpha_asset_d_q_row_scope_filter_column_qualified_name.setter + def alpha_asset_d_q_row_scope_filter_column_qualified_name( + self, alpha_asset_d_q_row_scope_filter_column_qualified_name: Optional[str] + ): + if self.attributes is None: + self.attributes = self.Attributes() + self.attributes.alpha_asset_d_q_row_scope_filter_column_qualified_name = ( + alpha_asset_d_q_row_scope_filter_column_qualified_name + ) + @property def alpha_asset_d_q_schedule_type(self) -> Optional[alpha_DQScheduleType]: return ( @@ -4451,6 +4479,9 @@ class Attributes(Referenceable.Attributes): asset_internal_popularity_score: Optional[float] = Field( default=None, description="" ) + alpha_asset_d_q_row_scope_filter_column_qualified_name: Optional[str] = Field( + default=None, description="" + ) alpha_asset_d_q_schedule_type: Optional[alpha_DQScheduleType] = Field( default=None, description="" ) From b0d3b3530739c25919c20eff2c048390b3414635 Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Sun, 21 Sep 2025 20:28:19 +0530 Subject: [PATCH 04/10] feat: implement set_dq_row_scope_filter_column method --- pyatlan/client/asset.py | 27 +++++++++++++++++++++++++++ pyatlan/errors.py | 2 +- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/pyatlan/client/asset.py b/pyatlan/client/asset.py index 164090a11..ad62bb393 100644 --- a/pyatlan/client/asset.py +++ b/pyatlan/client/asset.py @@ -1782,6 +1782,33 @@ def add_dq_rule_schedule( response = self.save(updated_asset) return response + @validate_arguments + def set_dq_row_scope_filter_column( + self, + asset_type: Type[A], + asset_name: str, + asset_qualified_name: str, + row_scope_filter_column_qualified_name: str, + ) -> AssetMutationResponse: + """ + Set the row scope filter column for data quality rules on an asset. + + :param asset_type: the type of asset to update (e.g., Table) + :param asset_name: the name of the asset to update + :param asset_qualified_name: the qualified name of the asset to update + :param row_scope_filter_column_qualified_name: the qualified name of the column to use for row scope filtering + :returns: the result of the save + :raises AtlanError: on any API communication issue + """ + updated_asset = asset_type.updater( + qualified_name=asset_qualified_name, name=asset_name + ) + updated_asset.alpha_asset_d_q_row_scope_filter_column_qualified_name = ( + row_scope_filter_column_qualified_name + ) + response = self.save(updated_asset) + return response + class SearchResults(ABC, Iterable): """ diff --git a/pyatlan/errors.py b/pyatlan/errors.py index a963df373..390889ea8 100644 --- a/pyatlan/errors.py +++ b/pyatlan/errors.py @@ -665,7 +665,7 @@ class ErrorCode(Enum): 400, "ATLAN-PYTHON-400-075", "Row scope filter column not configured for asset '{0}'.", - "please configure the row scope filter column first.", + "Use client.asset.set_dq_row_scope_filter_column() to configure the row scope filter column first.", InvalidRequestError, ) AUTHENTICATION_PASSTHROUGH = ( From 250ea41bd2f9087bca112f5ece91c8ae9e2a761c Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Sun, 21 Sep 2025 20:53:45 +0530 Subject: [PATCH 05/10] feat: add tests for row scope filtering and rule conditions --- tests/unit/model/alpha__d_q_rule_test.py | 142 +++++++++++++++++++++++ tests/unit/test_client.py | 22 ++++ 2 files changed, 164 insertions(+) diff --git a/tests/unit/model/alpha__d_q_rule_test.py b/tests/unit/model/alpha__d_q_rule_test.py index c9898f159..d097aa957 100644 --- a/tests/unit/model/alpha__d_q_rule_test.py +++ b/tests/unit/model/alpha__d_q_rule_test.py @@ -4,10 +4,12 @@ import pytest from pyatlan.model.assets import Column, Table, alpha_DQRule +from pyatlan.model.dq_rule_conditions import DQRuleConditionsBuilder from pyatlan.model.enums import ( alpha_DQDimension, alpha_DQRuleAlertPriority, alpha_DQRuleStatus, + alpha_dqRuleTemplateConfigRuleConditions, alpha_DQRuleThresholdCompareOperator, alpha_DQRuleThresholdUnit, ) @@ -39,6 +41,12 @@ def mock_client(): } } ) + config.alpha_dq_rule_template_config_rule_conditions = json.dumps( + {"enum": ["STRING_LENGTH_BETWEEN", "STRING_LENGTH_EQUAL"]} + ) + config.alpha_dq_rule_template_advanced_settings = json.dumps( + {"alpha_dqRuleRowScopeFilteringEnabled": True} + ) client.dq_template_config_cache.get_template_config.return_value = { "name": "Test Template", @@ -383,6 +391,140 @@ def test_column_level_rule_creator_with_optional_parameters(mock_client): ) +def test_column_level_rule_creator_with_row_scope_filtering(mock_client): + asset = Table.ref_by_qualified_name(qualified_name=ALPHA_DQ_TABLE_QUALIFIED_NAME) + asset.alpha_asset_d_q_row_scope_filter_column_qualified_name = ( + ALPHA_DQ_COLUMN_QUALIFIED_NAME + ) + column = Column.ref_by_qualified_name(qualified_name=ALPHA_DQ_COLUMN_QUALIFIED_NAME) + + dq_rule = alpha_DQRule.column_level_rule_creator( + client=mock_client, + rule_type=ALPHA_DQ_RULE_TYPE_COLUMN, + asset=asset, + column=column, + threshold_value=ALPHA_DQ_RULE_THRESHOLD_VALUE, + alert_priority=alpha_DQRuleAlertPriority.NORMAL, + row_scope_filtering_enabled=True, + ) + + assert dq_rule.alpha_dq_rule_alert_priority == alpha_DQRuleAlertPriority.NORMAL + assert dq_rule.alpha_dq_rule_status == alpha_DQRuleStatus.ACTIVE + assert dq_rule.alpha_dq_rule_row_scope_filtering_enabled is True + assert ( + dq_rule.alpha_dq_rule_base_column_qualified_name + == ALPHA_DQ_COLUMN_QUALIFIED_NAME + ) + + +def test_column_level_rule_creator_with_rule_conditions(mock_client): + asset = Table.ref_by_qualified_name(qualified_name=ALPHA_DQ_TABLE_QUALIFIED_NAME) + column = Column.ref_by_qualified_name(qualified_name=ALPHA_DQ_COLUMN_QUALIFIED_NAME) + + rule_conditions = ( + DQRuleConditionsBuilder() + .add_condition( + type=alpha_dqRuleTemplateConfigRuleConditions.STRING_LENGTH_BETWEEN, + min_value=5, + max_value=50, + ) + .build() + ) + + dq_rule = alpha_DQRule.column_level_rule_creator( + client=mock_client, + rule_type=ALPHA_DQ_RULE_TYPE_COLUMN, + asset=asset, + column=column, + threshold_value=ALPHA_DQ_RULE_THRESHOLD_VALUE, + alert_priority=alpha_DQRuleAlertPriority.NORMAL, + rule_conditions=rule_conditions, + ) + + assert dq_rule.alpha_dq_rule_alert_priority == alpha_DQRuleAlertPriority.NORMAL + assert dq_rule.alpha_dq_rule_status == alpha_DQRuleStatus.ACTIVE + assert ( + dq_rule.alpha_dq_rule_config_arguments.alpha_dq_rule_config_rule_conditions + == rule_conditions + ) + assert ( + dq_rule.alpha_dq_rule_base_column_qualified_name + == ALPHA_DQ_COLUMN_QUALIFIED_NAME + ) + + +def test_validate_template_features_rule_conditions_not_supported(mock_client): + from pyatlan.errors import InvalidRequestError + + config = Mock() + config.alpha_dq_rule_template_config_rule_conditions = None + config.alpha_dq_rule_template_advanced_settings = json.dumps({}) + + template_config = { + "name": "Test Template", + "qualified_name": "test/template/123", + "config": config, + } + + mock_client.dq_template_config_cache.get_template_config.return_value = ( + template_config + ) + + rule_conditions = ( + DQRuleConditionsBuilder() + .add_condition( + type=alpha_dqRuleTemplateConfigRuleConditions.STRING_LENGTH_BETWEEN, + min_value=5, + max_value=50, + ) + .build() + ) + + with pytest.raises( + InvalidRequestError, + match="Rule type 'Blank Count' does not support rule conditions", + ): + alpha_DQRule.Attributes._validate_template_features( + rule_type="Blank Count", + rule_conditions=rule_conditions, + row_scope_filtering_enabled=False, + template_config=template_config, + threshold_compare_operator=alpha_DQRuleThresholdCompareOperator.EQUAL, + ) + + +def test_validate_template_features_row_scope_filtering_not_supported(mock_client): + from pyatlan.errors import InvalidRequestError + + config = Mock() + config.alpha_dq_rule_template_config_rule_conditions = json.dumps( + {"enum": ["STRING_LENGTH_BETWEEN"]} + ) + config.alpha_dq_rule_template_advanced_settings = json.dumps({}) + + template_config = { + "name": "Test Template", + "qualified_name": "test/template/123", + "config": config, + } + + mock_client.dq_template_config_cache.get_template_config.return_value = ( + template_config + ) + + with pytest.raises( + InvalidRequestError, + match="Rule type 'Blank Count' does not support row scope filtering", + ): + alpha_DQRule.Attributes._validate_template_features( + rule_type="Blank Count", + rule_conditions=None, + row_scope_filtering_enabled=True, + template_config=template_config, + threshold_compare_operator=alpha_DQRuleThresholdCompareOperator.EQUAL, + ) + + @pytest.mark.parametrize( "qualified_name, message", [ diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 1a58858d5..167920e61 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2767,3 +2767,25 @@ def test_add_dq_rule_schedule(mock_api_caller): ) mock_save.assert_called_once_with(updated_table) assert result == mock_response + + +def test_set_dq_row_scope_filter_column(mock_api_caller): + from pyatlan.model.response import AssetMutationResponse + from tests.unit.model.constants import ( + ALPHA_DQ_COLUMN_QUALIFIED_NAME, + ALPHA_DQ_TABLE_QUALIFIED_NAME, + ) + + asset_client = AssetClient(mock_api_caller) + mock_response = Mock(spec=AssetMutationResponse) + + with patch.object(asset_client, "save", return_value=mock_response) as mock_save: + result = asset_client.set_dq_row_scope_filter_column( + asset_type=Table, + asset_name="TestTable", + asset_qualified_name=ALPHA_DQ_TABLE_QUALIFIED_NAME, + row_scope_filter_column_qualified_name=ALPHA_DQ_COLUMN_QUALIFIED_NAME, + ) + + mock_save.assert_called_once() + assert result == mock_response From adec96e5cd5b52f65aab29bc79ee10dd0c56fdd5 Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Sun, 21 Sep 2025 21:59:21 +0530 Subject: [PATCH 06/10] chore: qa fixes --- pyatlan/model/assets/core/alpha__d_q_rule.py | 24 +++++++++++++------- pyatlan/model/dq_rule_conditions.py | 14 ++++++++---- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/pyatlan/model/assets/core/alpha__d_q_rule.py b/pyatlan/model/assets/core/alpha__d_q_rule.py index 5c157e651..789ef6015 100644 --- a/pyatlan/model/assets/core/alpha__d_q_rule.py +++ b/pyatlan/model/assets/core/alpha__d_q_rule.py @@ -289,8 +289,8 @@ def updater( retrieved_column = search_result.alpha_dq_rule_base_column # type: ignore[attr-defined] retrieved_alert_priority = search_result.alpha_dq_rule_alert_priority # type: ignore[attr-defined] retrieved_row_scope_filtering_enabled = ( - search_result.alpha_dq_rule_row_scope_filtering_enabled - ) # type: ignore[attr-defined] + search_result.alpha_dq_rule_row_scope_filtering_enabled # type: ignore[attr-defined] + ) retrieved_description = search_result.user_description retrieved_asset = search_result.alpha_dq_rule_base_dataset # type: ignore[attr-defined] retrieved_template_rule_name = search_result.alpha_dq_rule_template_name # type: ignore[attr-defined] @@ -332,12 +332,18 @@ def updater( ) ) + final_compare_operator = ( + validated_threshold_operator + or threshold_compare_operator + or retrieved_threshold_compare_operator + ) + config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw( is_alert_enabled=True, custom_sql=custom_sql or retrieved_custom_sql, display_name=rule_name or retrieved_rule_name, dimension=dimension or retrieved_dimension, - compare_operator=validated_threshold_operator, + compare_operator=final_compare_operator, threshold_value=threshold_value or retrieved_threshold_value, threshold_unit=threshold_unit or retrieved_threshold_unit, column=retrieved_column, @@ -351,7 +357,7 @@ def updater( name="", alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments( alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject( - alpha_dq_rule_threshold_compare_operator=validated_threshold_operator, + alpha_dq_rule_threshold_compare_operator=final_compare_operator, alpha_dq_rule_threshold_value=threshold_value or retrieved_threshold_value, alpha_dq_rule_threshold_unit=threshold_unit @@ -1085,7 +1091,9 @@ class Attributes(DataQuality.Attributes): @staticmethod def _get_template_config_value( - config_value: str, property_name: str = None, value_key: str = "default" + config_value: str, + property_name: Optional[str] = None, + value_key: str = "default", ): if not config_value: return None @@ -1112,9 +1120,9 @@ def _validate_template_features( alpha_DQRuleThresholdCompareOperator ] = None, asset: Optional[Asset] = None, - ) -> alpha_DQRuleThresholdCompareOperator: + ) -> Optional[alpha_DQRuleThresholdCompareOperator]: if not template_config or not template_config.get("config"): - return + return None config = template_config["config"] @@ -1147,7 +1155,7 @@ def _validate_template_features( if rule_conditions: allowed_rule_conditions = ( alpha_DQRule.Attributes._get_template_config_value( - config.alpha_dq_rule_template_config_rule_conditions, + config.alpha_dq_rule_template_config_rule_conditions or "", None, "enum", ) diff --git a/pyatlan/model/dq_rule_conditions.py b/pyatlan/model/dq_rule_conditions.py index f86b26b7b..2e9b2a0b0 100644 --- a/pyatlan/model/dq_rule_conditions.py +++ b/pyatlan/model/dq_rule_conditions.py @@ -46,13 +46,19 @@ def __init__( validate_required_fields( ["min_value", "max_value"], [self.min_value, self.max_value] ) - if self.min_value < 0 or self.max_value < 0: + if (self.min_value is not None and self.min_value < 0) or ( + self.max_value is not None and self.max_value < 0 + ): raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( f"min_value={self.min_value}, max_value={self.max_value}", "min_value, max_value", "non-negative integers", ) - if self.min_value > self.max_value: + if ( + self.min_value is not None + and self.max_value is not None + and self.min_value > self.max_value + ): raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( f"min_value={self.min_value}, max_value={self.max_value}", "min_value, max_value", @@ -73,7 +79,7 @@ def __init__( def to_dict(self) -> Dict[str, Any]: """Convert the condition to a dictionary representation.""" - result = {"type": self.type.value} + result: Dict[str, Any] = {"type": self.type.value} if self.type == alpha_dqRuleTemplateConfigRuleConditions.STRING_LENGTH_BETWEEN: result["value"] = {"minValue": self.min_value, "maxValue": self.max_value} @@ -86,7 +92,7 @@ def to_dict(self) -> Dict[str, Any]: class DQRuleConditionsBuilder: """Builder for data quality rule conditions.""" - def __init__(self): + def __init__(self) -> None: self._conditions: List[DQCondition] = [] def add_condition( From a0b2493dbf2db089f0532514ee16b04355be8db8 Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Sun, 21 Sep 2025 23:36:41 +0530 Subject: [PATCH 07/10] fix: dq rule condition builder --- pyatlan/model/dq_rule_conditions.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/pyatlan/model/dq_rule_conditions.py b/pyatlan/model/dq_rule_conditions.py index 2e9b2a0b0..bf3a3c927 100644 --- a/pyatlan/model/dq_rule_conditions.py +++ b/pyatlan/model/dq_rule_conditions.py @@ -14,14 +14,7 @@ class DQCondition(AtlanObject): - """ - Data quality rule condition. - - :param type: the condition type enum value - :param value: value of type str, int, or list depending on condition type - :param min_value: minimum value for range-based conditions - :param max_value: maximum value for range-based conditions - """ + """Data quality rule condition.""" type: alpha_dqRuleTemplateConfigRuleConditions = Field(description="") value: Optional[Union[str, int, List[str], Dict[str, Any]]] = Field( @@ -78,13 +71,12 @@ def __init__( validate_type("value", str, self.value) def to_dict(self) -> Dict[str, Any]: - """Convert the condition to a dictionary representation.""" result: Dict[str, Any] = {"type": self.type.value} if self.type == alpha_dqRuleTemplateConfigRuleConditions.STRING_LENGTH_BETWEEN: result["value"] = {"minValue": self.min_value, "maxValue": self.max_value} else: - result["value"] = self.value + result["value"] = {"value": self.value} return result @@ -119,7 +111,6 @@ def add_condition( return self def build(self) -> str: - """Build the conditions as a JSON string.""" if not self._conditions: raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( "empty conditions list", "conditions", "at least one condition" From ffba9c53ad7cf135768fd561c72caf5e3825bdd4 Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Mon, 22 Sep 2025 10:36:49 +0530 Subject: [PATCH 08/10] chore: update jinja templates --- .../methods/asset/alpha__d_q_rule.jinja2 | 62 ++++++++++++++----- .../methods/attribute/alpha__d_q_rule.jinja2 | 33 +++++++--- 2 files changed, 71 insertions(+), 24 deletions(-) diff --git a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 index 106023b44..a3e7b590f 100644 --- a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 +++ b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 @@ -1,5 +1,5 @@ - @classmethod +@classmethod @init_guid def custom_sql_creator( cls, @@ -137,12 +137,31 @@ ], ) template_config = client.dq_template_config_cache.get_template_config(rule_type) - threshold_compare_operator = alpha_DQRule.Attributes._validate_template_features( - rule_type, - rule_conditions, - row_scope_filtering_enabled, - template_config, - threshold_compare_operator, + + asset_for_validation = asset + if row_scope_filtering_enabled and asset.qualified_name: + from pyatlan.model.fluent_search import FluentSearch + + search_request = ( + FluentSearch() + .where(Asset.QUALIFIED_NAME.eq(asset.qualified_name)) + .include_on_results( + Asset.ALPHAASSET_DQ_ROW_SCOPE_FILTER_COLUMN_QUALIFIED_NAME + ) + ).to_request() + results = client.asset.search(search_request) + if results.count == 1: + asset_for_validation = results.current_page()[0] + + threshold_compare_operator = ( + alpha_DQRule.Attributes._validate_template_features( + rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator, + asset_for_validation, + ) ) attributes = alpha_DQRule.Attributes.creator( @@ -222,8 +241,8 @@ retrieved_column = search_result.alpha_dq_rule_base_column # type: ignore[attr-defined] retrieved_alert_priority = search_result.alpha_dq_rule_alert_priority # type: ignore[attr-defined] retrieved_row_scope_filtering_enabled = ( - search_result.alpha_dq_rule_row_scope_filtering_enabled - ) # type: ignore[attr-defined] + search_result.alpha_dq_rule_row_scope_filtering_enabled # type: ignore[attr-defined] + ) retrieved_description = search_result.user_description retrieved_asset = search_result.alpha_dq_rule_base_dataset # type: ignore[attr-defined] retrieved_template_rule_name = search_result.alpha_dq_rule_template_name # type: ignore[attr-defined] @@ -254,12 +273,21 @@ template_config = client.dq_template_config_cache.get_template_config( retrieved_rule_type ) - validated_threshold_operator = alpha_DQRule.Attributes._validate_template_features( - retrieved_rule_type, - rule_conditions, - row_scope_filtering_enabled, - template_config, - threshold_compare_operator or retrieved_threshold_compare_operator, + validated_threshold_operator = ( + alpha_DQRule.Attributes._validate_template_features( + retrieved_rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator or retrieved_threshold_compare_operator, + retrieved_asset, + ) + ) + + final_compare_operator = ( + validated_threshold_operator + or threshold_compare_operator + or retrieved_threshold_compare_operator ) config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw( @@ -267,7 +295,7 @@ custom_sql=custom_sql or retrieved_custom_sql, display_name=rule_name or retrieved_rule_name, dimension=dimension or retrieved_dimension, - compare_operator=validated_threshold_operator, + compare_operator=final_compare_operator, threshold_value=threshold_value or retrieved_threshold_value, threshold_unit=threshold_unit or retrieved_threshold_unit, column=retrieved_column, @@ -281,7 +309,7 @@ name="", alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments( alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject( - alpha_dq_rule_threshold_compare_operator=validated_threshold_operator, + alpha_dq_rule_threshold_compare_operator=final_compare_operator, alpha_dq_rule_threshold_value=threshold_value or retrieved_threshold_value, alpha_dq_rule_threshold_unit=threshold_unit diff --git a/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 b/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 index 1a5c2d9d1..5d3d911af 100644 --- a/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 +++ b/pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2 @@ -1,6 +1,8 @@ @staticmethod def _get_template_config_value( - config_value: str, property_name: str = None, value_key: str = "default" + config_value: str, + property_name: Optional[str] = None, + value_key: str = "default", ): if not config_value: return None @@ -26,9 +28,10 @@ threshold_compare_operator: Optional[ alpha_DQRuleThresholdCompareOperator ] = None, - ) -> alpha_DQRuleThresholdCompareOperator: + asset: Optional[Asset] = None, + ) -> Optional[alpha_DQRuleThresholdCompareOperator]: if not template_config or not template_config.get("config"): - return + return None config = template_config["config"] @@ -41,15 +44,30 @@ ) if row_scope_filtering_enabled: - advanced_settings = config.alpha_dq_rule_template_advanced_settings or "" + advanced_settings = ( + config.alpha_dq_rule_template_advanced_settings or "" + ) if "alpha_dqRuleRowScopeFilteringEnabled" not in str(advanced_settings): raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters( rule_type, "row scope filtering" ) + if asset and not getattr( + asset, + "alpha_asset_d_q_row_scope_filter_column_qualified_name", + None, + ): + raise ErrorCode.DQ_ROW_SCOPE_FILTER_COLUMN_MISSING.exception_with_parameters( + getattr(asset, "qualified_name", "unknown") + ) + if rule_conditions: - allowed_rule_conditions = alpha_DQRule.Attributes._get_template_config_value( - config.alpha_dq_rule_template_config_rule_conditions, None, "enum" + allowed_rule_conditions = ( + alpha_DQRule.Attributes._get_template_config_value( + config.alpha_dq_rule_template_config_rule_conditions or "", + None, + "enum", + ) ) if allowed_rule_conditions: try: @@ -71,7 +89,8 @@ if threshold_compare_operator is None: return alpha_DQRuleThresholdCompareOperator.EQUAL elif ( - threshold_compare_operator != alpha_DQRuleThresholdCompareOperator.EQUAL + threshold_compare_operator + != alpha_DQRuleThresholdCompareOperator.EQUAL ): raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters( f"threshold_compare_operator={threshold_compare_operator.value}", From 67c171c2aeb8f2077e3b1fde6d54d0b5d8693b9c Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Mon, 22 Sep 2025 10:48:04 +0530 Subject: [PATCH 09/10] chore: add fallback --- .../templates/methods/asset/alpha__d_q_rule.jinja2 | 13 ++++++++++--- pyatlan/model/assets/core/alpha__d_q_rule.py | 11 +++++++++-- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 index a3e7b590f..9aab02d8c 100644 --- a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 +++ b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 @@ -1,5 +1,5 @@ -@classmethod + @classmethod @init_guid def custom_sql_creator( cls, @@ -153,7 +153,7 @@ if results.count == 1: asset_for_validation = results.current_page()[0] - threshold_compare_operator = ( + validated_threshold_operator = ( alpha_DQRule.Attributes._validate_template_features( rule_type, rule_conditions, @@ -164,12 +164,18 @@ ) ) + final_threshold_compare_operator = ( + validated_threshold_operator + or threshold_compare_operator + or alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL + ) + attributes = alpha_DQRule.Attributes.creator( client=client, rule_type=rule_type, asset=asset, column=column, - threshold_compare_operator=threshold_compare_operator, + threshold_compare_operator=final_threshold_compare_operator, threshold_value=threshold_value, alert_priority=alert_priority, threshold_unit=threshold_unit, @@ -288,6 +294,7 @@ validated_threshold_operator or threshold_compare_operator or retrieved_threshold_compare_operator + or alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL ) config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw( diff --git a/pyatlan/model/assets/core/alpha__d_q_rule.py b/pyatlan/model/assets/core/alpha__d_q_rule.py index 789ef6015..df25a2c7d 100644 --- a/pyatlan/model/assets/core/alpha__d_q_rule.py +++ b/pyatlan/model/assets/core/alpha__d_q_rule.py @@ -201,7 +201,7 @@ def column_level_rule_creator( if results.count == 1: asset_for_validation = results.current_page()[0] - threshold_compare_operator = ( + validated_threshold_operator = ( alpha_DQRule.Attributes._validate_template_features( rule_type, rule_conditions, @@ -212,12 +212,18 @@ def column_level_rule_creator( ) ) + final_threshold_compare_operator = ( + validated_threshold_operator + or threshold_compare_operator + or alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL + ) + attributes = alpha_DQRule.Attributes.creator( client=client, rule_type=rule_type, asset=asset, column=column, - threshold_compare_operator=threshold_compare_operator, + threshold_compare_operator=final_threshold_compare_operator, threshold_value=threshold_value, alert_priority=alert_priority, threshold_unit=threshold_unit, @@ -336,6 +342,7 @@ def updater( validated_threshold_operator or threshold_compare_operator or retrieved_threshold_compare_operator + or alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL ) config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw( From ecb563c2649eb7a96cbc46d29c4e520b2bea88af Mon Sep 17 00:00:00 2001 From: Uday Sagar Date: Mon, 22 Sep 2025 20:12:09 +0530 Subject: [PATCH 10/10] fix: qa fixes --- .../methods/asset/alpha__d_q_rule.jinja2 | 2 +- pyatlan/model/dq_rule_conditions.py | 2 +- tests/unit/model/alpha__d_q_rule_test.py | 78 ++++++++++++++++++- tests/unit/test_client.py | 8 +- 4 files changed, 78 insertions(+), 12 deletions(-) diff --git a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 index 9aab02d8c..36e41bdcf 100644 --- a/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 +++ b/pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2 @@ -1,5 +1,5 @@ - @classmethod + @classmethod @init_guid def custom_sql_creator( cls, diff --git a/pyatlan/model/dq_rule_conditions.py b/pyatlan/model/dq_rule_conditions.py index bf3a3c927..6daddfad2 100644 --- a/pyatlan/model/dq_rule_conditions.py +++ b/pyatlan/model/dq_rule_conditions.py @@ -93,7 +93,7 @@ def add_condition( value: Optional[Union[str, int, List[str]]] = None, min_value: Optional[int] = None, max_value: Optional[int] = None, - ) -> "DQRuleConditionsBuilder": + ) -> DQRuleConditionsBuilder: """ Add a condition to the builder. diff --git a/tests/unit/model/alpha__d_q_rule_test.py b/tests/unit/model/alpha__d_q_rule_test.py index d097aa957..ef312f295 100644 --- a/tests/unit/model/alpha__d_q_rule_test.py +++ b/tests/unit/model/alpha__d_q_rule_test.py @@ -3,6 +3,7 @@ import pytest +from pyatlan.errors import ErrorCode, InvalidRequestError from pyatlan.model.assets import Column, Table, alpha_DQRule from pyatlan.model.dq_rule_conditions import DQRuleConditionsBuilder from pyatlan.model.enums import ( @@ -454,8 +455,6 @@ def test_column_level_rule_creator_with_rule_conditions(mock_client): def test_validate_template_features_rule_conditions_not_supported(mock_client): - from pyatlan.errors import InvalidRequestError - config = Mock() config.alpha_dq_rule_template_config_rule_conditions = None config.alpha_dq_rule_template_advanced_settings = json.dumps({}) @@ -494,8 +493,6 @@ def test_validate_template_features_rule_conditions_not_supported(mock_client): def test_validate_template_features_row_scope_filtering_not_supported(mock_client): - from pyatlan.errors import InvalidRequestError - config = Mock() config.alpha_dq_rule_template_config_rule_conditions = json.dumps( {"enum": ["STRING_LENGTH_BETWEEN"]} @@ -539,3 +536,76 @@ def test_updater_with_invalid_parameter_raises_value_error( client=mock_client, qualified_name=qualified_name, ) + + +def test_validate_template_features_invalid_rule_conditions(mock_client): + config = Mock() + config.alpha_dq_rule_template_config_rule_conditions = json.dumps( + {"enum": ["STRING_LENGTH_BETWEEN"]} + ) + config.alpha_dq_rule_template_advanced_settings = json.dumps({}) + + template_config = { + "name": "Test Template", + "qualified_name": "test/template/123", + "config": config, + } + + mock_client.dq_template_config_cache.get_template_config.return_value = ( + template_config + ) + + unsupported_condition = json.dumps( + {"conditions": [{"type": "UNSUPPORTED_CONDITION", "value": "test"}]} + ) + + with pytest.raises( + InvalidRequestError, + match="Invalid rule conditions: condition type 'UNSUPPORTED_CONDITION' not supported, allowed: \\['STRING_LENGTH_BETWEEN'\\]", + ): + alpha_DQRule.Attributes._validate_template_features( + rule_type="Test Rule", + rule_conditions=unsupported_condition, + row_scope_filtering_enabled=False, + template_config=template_config, + threshold_compare_operator=alpha_DQRuleThresholdCompareOperator.EQUAL, + ) + + +def test_validate_template_features_row_scope_filter_column_missing(mock_client): + config = Mock() + config.alpha_dq_rule_template_config_rule_conditions = json.dumps( + {"enum": ["STRING_LENGTH_BETWEEN"]} + ) + config.alpha_dq_rule_template_advanced_settings = json.dumps( + {"alpha_dqRuleRowScopeFilteringEnabled": True} + ) + + template_config = { + "name": "Test Template", + "qualified_name": "test/template/123", + "config": config, + } + + mock_client.dq_template_config_cache.get_template_config.return_value = ( + template_config + ) + + table_asset = Table.ref_by_qualified_name( + qualified_name=ALPHA_DQ_TABLE_QUALIFIED_NAME + ) + + with pytest.raises( + InvalidRequestError, + match=ErrorCode.DQ_ROW_SCOPE_FILTER_COLUMN_MISSING.error_message.format( + ALPHA_DQ_TABLE_QUALIFIED_NAME + ), + ): + alpha_DQRule.Attributes._validate_template_features( + rule_type="Test Rule", + rule_conditions=None, + row_scope_filtering_enabled=True, + template_config=template_config, + threshold_compare_operator=alpha_DQRuleThresholdCompareOperator.EQUAL, + asset=table_asset, + ) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 167920e61..0ad804fe9 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -71,6 +71,8 @@ TEST_USER_CLIENT_METHODS, ) from tests.unit.model.constants import ( + ALPHA_DQ_COLUMN_QUALIFIED_NAME, + ALPHA_DQ_TABLE_QUALIFIED_NAME, CONNECTION_NAME, CONNECTOR_TYPE, DATA_DOMAIN_NAME, @@ -2770,12 +2772,6 @@ def test_add_dq_rule_schedule(mock_api_caller): def test_set_dq_row_scope_filter_column(mock_api_caller): - from pyatlan.model.response import AssetMutationResponse - from tests.unit.model.constants import ( - ALPHA_DQ_COLUMN_QUALIFIED_NAME, - ALPHA_DQ_TABLE_QUALIFIED_NAME, - ) - asset_client = AssetClient(mock_api_caller) mock_response = Mock(spec=AssetMutationResponse)