Skip to content
27 changes: 27 additions & 0 deletions pyatlan/client/asset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
21 changes: 21 additions & 0 deletions pyatlan/errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
82 changes: 74 additions & 8 deletions pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2
Original file line number Diff line number Diff line change
Expand Up @@ -113,16 +113,17 @@
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(
[
"client",
"rule_type",
"asset",
"column",
"threshold_compare_operator",
"threshold_value",
"alert_priority",
],
Expand All @@ -131,25 +132,59 @@
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,
rule_name=None,
dimension=None,
custom_sql=None,
description=None,
rule_conditions=rule_conditions,
row_scope_filtering_enabled=row_scope_filtering_enabled,
)
return cls(attributes=attributes)

Expand All @@ -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

Expand All @@ -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()
Expand All @@ -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]
Expand All @@ -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,
Expand Down
Loading