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 5918a27e0..390889ea8 100644 --- a/pyatlan/errors.py +++ b/pyatlan/errors.py @@ -647,6 +647,27 @@ 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, + ) + DQ_ROW_SCOPE_FILTER_COLUMN_MISSING = ( + 400, + "ATLAN-PYTHON-400-075", + "Row scope filter column not configured for asset '{0}'.", + "Use client.asset.set_dq_row_scope_filter_column() to configure the row scope filter column first.", + InvalidRequestError, + ) AUTHENTICATION_PASSTHROUGH = ( 401, "ATLAN-PYTHON-401-000", 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..36e41bdcf 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,18 +132,50 @@ rule_type, asset, column, - threshold_compare_operator, threshold_value, alert_priority, ], ) + 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] + + validated_threshold_operator = ( + alpha_DQRule.Attributes._validate_template_features( + rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator, + asset_for_validation, + ) + ) + + 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, @@ -150,6 +183,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 +204,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 +227,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 +246,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 +275,60 @@ 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, + retrieved_asset, + ) + ) + + final_compare_operator = ( + 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( 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=final_compare_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=final_compare_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..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,3 +1,133 @@ + @staticmethod + def _get_template_config_value( + config_value: str, + property_name: Optional[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, + asset: Optional[Asset] = None, + ) -> Optional[alpha_DQRuleThresholdCompareOperator]: + if not template_config or not template_config.get("config"): + return None + + 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 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 or "", + 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 +142,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 +172,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 +223,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 +242,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 +260,8 @@ column=column, dq_priority=alert_priority, description=description, + rule_conditions=rule_conditions, + row_scope_filtering_enabled=row_scope_filtering_enabled, ) ) @@ -133,9 +274,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, diff --git a/pyatlan/model/assets/core/alpha__d_q_rule.py b/pyatlan/model/assets/core/alpha__d_q_rule.py index c5a4cd95f..df25a2c7d 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, @@ -40,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): @@ -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,18 +180,50 @@ 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) + + 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] + + validated_threshold_operator = ( + alpha_DQRule.Attributes._validate_template_features( + rule_type, + rule_conditions, + row_scope_filtering_enabled, + template_config, + threshold_compare_operator, + asset_for_validation, + ) + ) + + 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, @@ -197,6 +231,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 +252,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 +275,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 +294,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 +323,60 @@ 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, + retrieved_asset, + ) + ) + + final_compare_operator = ( + 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( 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=final_compare_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=final_compare_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 +454,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 +590,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 +685,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 +1036,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 +1096,137 @@ class Attributes(DataQuality.Attributes): default=None, description="" ) # relationship + @staticmethod + def _get_template_config_value( + config_value: str, + property_name: Optional[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, + asset: Optional[Asset] = None, + ) -> Optional[alpha_DQRuleThresholdCompareOperator]: + if not template_config or not template_config.get("config"): + return None + + 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 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 or "", + 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 +1240,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 +1270,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 +1321,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 +1340,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 +1358,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 +1372,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/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="" ) diff --git a/pyatlan/model/dq_rule_conditions.py b/pyatlan/model/dq_rule_conditions.py new file mode 100644 index 000000000..6daddfad2 --- /dev/null +++ b/pyatlan/model/dq_rule_conditions.py @@ -0,0 +1,124 @@ +# 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.""" + + 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 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 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", + "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]: + 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"] = {"value": self.value} + + return result + + +class DQRuleConditionsBuilder: + """Builder for data quality rule conditions.""" + + def __init__(self) -> None: + 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: + 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): diff --git a/tests/unit/model/alpha__d_q_rule_test.py b/tests/unit/model/alpha__d_q_rule_test.py index c9898f159..ef312f295 100644 --- a/tests/unit/model/alpha__d_q_rule_test.py +++ b/tests/unit/model/alpha__d_q_rule_test.py @@ -3,11 +3,14 @@ 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 ( alpha_DQDimension, alpha_DQRuleAlertPriority, alpha_DQRuleStatus, + alpha_dqRuleTemplateConfigRuleConditions, alpha_DQRuleThresholdCompareOperator, alpha_DQRuleThresholdUnit, ) @@ -39,6 +42,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 +392,136 @@ 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): + 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): + 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", [ @@ -397,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 1a58858d5..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, @@ -2767,3 +2769,19 @@ 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): + 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