Skip to content

Commit fed0f11

Browse files
Aryamanz29Aryamanz29
authored
Merge pull request #789 from atlanhq/DQ-421
DQ-421 | feat: add support for new data quality rule templates
2 parents 32abd94 + 44381eb commit fed0f11

10 files changed

Lines changed: 686 additions & 145 deletions

File tree

‎pyatlan/cache/aio/dq_template_config_cache.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ async def get_template_config(self, rule_type: str) -> Optional[Dict]:
3232
"""
3333
Get template configuration for a specific rule type.
3434
35-
:param rule_type: The display name of the rule type
35+
:param rule_type: The name of the rule template
3636
:returns: Template configuration dict or None if not found
3737
"""
3838
if not self._initialized:

‎pyatlan/cache/common/dq_template_config_cache.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ def process_search_results(
6060
"dimension": result.dq_rule_template_dimension, # type: ignore
6161
"config": result.dq_rule_template_config, # type: ignore
6262
}
63-
cache[result.display_name] = template_config # type: ignore
63+
cache[result.name] = template_config # type: ignore
6464
return True, None
6565
except Exception as e:
6666
return False, e

‎pyatlan/cache/dq_template_config_cache.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def get_template_config(self, rule_type: str) -> Optional[Dict]:
3232
"""
3333
Get template configuration for a specific rule type.
3434
35-
:param rule_type: The display name of the rule type
35+
:param rule_type: The name of the rule template
3636
:returns: Template configuration dict or None if not found
3737
"""
3838
if not self._initialized:

‎pyatlan/generator/templates/methods/asset/data_quality_rule.jinja2‎

Lines changed: 109 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
threshold_value: int,
1313
alert_priority: DataQualityRuleAlertPriority,
1414
dimension: DataQualityDimension,
15+
custom_sql_return_type: Optional[DataQualityRuleCustomSQLReturnType] = None,
1516
description: Optional[str] = None,
1617
) -> DataQualityRule:
1718
validate_required_fields(
@@ -40,13 +41,14 @@
4041
attributes = DataQualityRule.Attributes.creator(
4142
client=client,
4243
rule_name=rule_name,
43-
rule_type="Custom SQL",
44+
rule_type=DataQualityRuleTemplateType.CUSTOM_SQL,
4445
asset=asset,
4546
threshold_compare_operator=threshold_compare_operator,
4647
threshold_value=threshold_value,
4748
alert_priority=alert_priority,
4849
dimension=dimension,
4950
custom_sql=custom_sql,
51+
custom_sql_return_type=custom_sql_return_type,
5052
description=description,
5153
column=None,
5254
threshold_unit=None,
@@ -59,44 +61,76 @@
5961
cls,
6062
*,
6163
client: AtlanClient,
62-
rule_type: str,
64+
rule_type: DataQualityRuleTemplateType,
6365
asset: Asset,
64-
threshold_compare_operator: DataQualityRuleThresholdCompareOperator,
6566
threshold_value: int,
6667
alert_priority: DataQualityRuleAlertPriority,
68+
threshold_compare_operator: Optional[
69+
DataQualityRuleThresholdCompareOperator
70+
] = None,
71+
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
72+
rule_conditions: Optional[str] = None,
73+
row_scope_filtering_enabled: Optional[bool] = False,
6774
) -> DataQualityRule:
6875
validate_required_fields(
6976
[
7077
"client",
7178
"rule_type",
7279
"asset",
73-
"threshold_compare_operator",
7480
"threshold_value",
7581
"alert_priority",
7682
],
7783
[
7884
client,
7985
rule_type,
8086
asset,
81-
threshold_compare_operator,
8287
threshold_value,
8388
alert_priority,
8489
],
8590
)
91+
template_config = client.dq_template_config_cache.get_template_config(
92+
rule_type.value
93+
)
94+
95+
asset_for_validation, target_table_asset = (
96+
DataQualityRule.Attributes._fetch_assets_for_row_scope_validation(
97+
client, asset, rule_conditions, row_scope_filtering_enabled or False
98+
)
99+
)
100+
101+
validated_threshold_operator = (
102+
DataQualityRule.Attributes._validate_template_features(
103+
rule_type,
104+
rule_conditions,
105+
row_scope_filtering_enabled,
106+
template_config,
107+
threshold_compare_operator,
108+
asset_for_validation,
109+
target_table_asset,
110+
)
111+
)
112+
113+
final_threshold_compare_operator = (
114+
validated_threshold_operator
115+
or threshold_compare_operator
116+
or DataQualityRuleThresholdCompareOperator.LESS_THAN_EQUAL
117+
)
86118

87119
attributes = DataQualityRule.Attributes.creator(
88120
client=client,
89121
rule_type=rule_type,
90122
asset=asset,
91-
threshold_compare_operator=threshold_compare_operator,
123+
threshold_compare_operator=final_threshold_compare_operator,
92124
threshold_value=threshold_value,
93125
alert_priority=alert_priority,
94126
rule_name=None,
95127
column=None,
96-
threshold_unit=None,
128+
threshold_unit=threshold_unit,
97129
dimension=None,
98130
custom_sql=None,
99131
description=None,
132+
rule_conditions=rule_conditions,
133+
row_scope_filtering_enabled=row_scope_filtering_enabled,
100134
)
101135
return cls(attributes=attributes)
102136

@@ -106,7 +140,7 @@
106140
cls,
107141
*,
108142
client: AtlanClient,
109-
rule_type: str,
143+
rule_type: DataQualityRuleTemplateType,
110144
asset: Asset,
111145
column: Asset,
112146
threshold_value: int,
@@ -136,22 +170,15 @@
136170
alert_priority,
137171
],
138172
)
139-
template_config = client.dq_template_config_cache.get_template_config(rule_type)
140-
141-
asset_for_validation = asset
142-
if row_scope_filtering_enabled and asset.qualified_name:
143-
from pyatlan.model.fluent_search import FluentSearch
173+
template_config = client.dq_template_config_cache.get_template_config(
174+
rule_type.value
175+
)
144176

145-
search_request = (
146-
FluentSearch()
147-
.where(Asset.QUALIFIED_NAME.eq(asset.qualified_name))
148-
.include_on_results(
149-
Asset.ASSET_DQ_ROW_SCOPE_FILTER_COLUMN_QUALIFIED_NAME
150-
)
151-
).to_request()
152-
results = client.asset.search(search_request)
153-
if results.count == 1:
154-
asset_for_validation = results.current_page()[0]
177+
asset_for_validation, target_table_asset = (
178+
DataQualityRule.Attributes._fetch_assets_for_row_scope_validation(
179+
client, asset, rule_conditions, row_scope_filtering_enabled or False
180+
)
181+
)
155182

156183
validated_threshold_operator = (
157184
DataQualityRule.Attributes._validate_template_features(
@@ -161,6 +188,7 @@
161188
template_config,
162189
threshold_compare_operator,
163190
asset_for_validation,
191+
target_table_asset,
164192
)
165193
)
166194

@@ -202,6 +230,7 @@
202230
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
203231
dimension: Optional[DataQualityDimension] = None,
204232
custom_sql: Optional[str] = None,
233+
custom_sql_return_type: Optional[DataQualityRuleCustomSQLReturnType] = None,
205234
rule_name: Optional[str] = None,
206235
description: Optional[str] = None,
207236
rule_conditions: Optional[str] = None,
@@ -224,6 +253,7 @@
224253
.include_on_results(DataQualityRule.DQ_RULE_ALERT_PRIORITY)
225254
.include_on_results(DataQualityRule.DISPLAY_NAME)
226255
.include_on_results(DataQualityRule.DQ_RULE_CUSTOM_SQL)
256+
.include_on_results(DataQualityRule.DQ_RULE_CUSTOM_SQL_RETURN_TYPE)
227257
.include_on_results(DataQualityRule.USER_DESCRIPTION)
228258
.include_on_results(DataQualityRule.DQ_RULE_DIMENSION)
229259
.include_on_results(DataQualityRule.DQ_RULE_CONFIG_ARGUMENTS)
@@ -242,6 +272,9 @@
242272
search_result = results.current_page()[0]
243273

244274
retrieved_custom_sql = search_result.dq_rule_custom_s_q_l # type: ignore[attr-defined]
275+
retrieved_custom_sql_return_type = (
276+
search_result.dq_rule_custom_s_q_l_return_type # type: ignore[attr-defined]
277+
)
245278
retrieved_rule_name = search_result.display_name
246279
retrieved_dimension = search_result.dq_rule_dimension # type: ignore[attr-defined]
247280
retrieved_column = search_result.dq_rule_base_column # type: ignore[attr-defined]
@@ -275,20 +308,56 @@
275308
else None
276309
) # type: ignore[attr-defined]
277310

278-
retrieved_rule_type = retrieved_template_rule_name
279-
template_config = client.dq_template_config_cache.get_template_config(
280-
retrieved_rule_type
281-
)
282-
validated_threshold_operator = (
283-
DataQualityRule.Attributes._validate_template_features(
284-
retrieved_rule_type,
285-
rule_conditions,
286-
row_scope_filtering_enabled,
287-
template_config,
288-
threshold_compare_operator or retrieved_threshold_compare_operator,
289-
retrieved_asset,
311+
template_config = None
312+
if retrieved_template_rule_name:
313+
template_config = client.dq_template_config_cache.get_template_config(
314+
retrieved_template_rule_name
315+
)
316+
317+
if rule_conditions:
318+
final_rule_conditions = rule_conditions
319+
elif search_result.dq_rule_config_arguments is not None: # type: ignore[attr-defined]
320+
final_rule_conditions = (
321+
search_result.dq_rule_config_arguments.dq_rule_config_rule_conditions # type: ignore[attr-defined]
290322
)
323+
else:
324+
final_rule_conditions = None
325+
326+
final_row_scope_filtering_enabled = (
327+
row_scope_filtering_enabled or retrieved_row_scope_filtering_enabled
291328
)
329+
if retrieved_asset:
330+
retrieved_asset, target_table_asset = (
331+
DataQualityRule.Attributes._fetch_assets_for_row_scope_validation(
332+
client,
333+
retrieved_asset,
334+
final_rule_conditions,
335+
final_row_scope_filtering_enabled,
336+
)
337+
)
338+
else:
339+
target_table_asset = None
340+
341+
validated_threshold_operator = None
342+
if retrieved_template_rule_name and template_config:
343+
try:
344+
retrieved_rule_type = DataQualityRuleTemplateType(
345+
retrieved_template_rule_name
346+
)
347+
validated_threshold_operator = (
348+
DataQualityRule.Attributes._validate_template_features(
349+
retrieved_rule_type,
350+
final_rule_conditions,
351+
final_row_scope_filtering_enabled,
352+
template_config,
353+
threshold_compare_operator
354+
or retrieved_threshold_compare_operator,
355+
retrieved_asset,
356+
target_table_asset,
357+
)
358+
)
359+
except ValueError:
360+
pass
292361

293362
final_compare_operator = (
294363
validated_threshold_operator
@@ -306,12 +375,11 @@
306375
or retrieved_threshold_value,
307376
dq_rule_threshold_unit=threshold_unit or retrieved_threshold_unit,
308377
),
309-
dq_rule_config_rule_conditions=rule_conditions,
378+
dq_rule_config_rule_conditions=final_rule_conditions,
310379
),
311380
dq_rule_base_dataset_qualified_name=retrieved_asset.qualified_name,
312381
dq_rule_alert_priority=alert_priority or retrieved_alert_priority,
313-
dq_rule_row_scope_filtering_enabled=row_scope_filtering_enabled
314-
or retrieved_row_scope_filtering_enabled,
382+
dq_rule_row_scope_filtering_enabled=final_row_scope_filtering_enabled,
315383
dq_rule_base_dataset=retrieved_asset,
316384
qualified_name=qualified_name,
317385
dq_rule_dimension=dimension or retrieved_dimension,
@@ -329,6 +397,9 @@
329397
if custom_sql is not None:
330398
attr_dq.dq_rule_custom_s_q_l = custom_sql
331399
attr_dq.display_name = rule_name or retrieved_rule_name
400+
attr_dq.dq_rule_custom_s_q_l_return_type = (
401+
custom_sql_return_type or retrieved_custom_sql_return_type
402+
)
332403
if description is not None:
333404
attr_dq.user_description = description or retrieved_description
334405

0 commit comments

Comments
 (0)