Skip to content

Commit 2764ec9

Browse files
UdaySagar-GitUdaySagar-Git
committed
refactor: update rule_type parameter to use DataQualityRuleTemplateType enum (name) instead of display name
1 parent b658113 commit 2764ec9

9 files changed

Lines changed: 174 additions & 107 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: 39 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
attributes = DataQualityRule.Attributes.creator(
4242
client=client,
4343
rule_name=rule_name,
44-
rule_type="Custom SQL",
44+
rule_type=DataQualityRuleTemplateType.CUSTOM_SQL,
4545
asset=asset,
4646
threshold_compare_operator=threshold_compare_operator,
4747
threshold_value=threshold_value,
@@ -61,7 +61,7 @@
6161
cls,
6262
*,
6363
client: AtlanClient,
64-
rule_type: str,
64+
rule_type: DataQualityRuleTemplateType,
6565
asset: Asset,
6666
threshold_compare_operator: DataQualityRuleThresholdCompareOperator,
6767
threshold_value: int,
@@ -109,7 +109,7 @@
109109
cls,
110110
*,
111111
client: AtlanClient,
112-
rule_type: str,
112+
rule_type: DataQualityRuleTemplateType,
113113
asset: Asset,
114114
column: Asset,
115115
threshold_value: int,
@@ -139,7 +139,9 @@
139139
alert_priority,
140140
],
141141
)
142-
template_config = client.dq_template_config_cache.get_template_config(rule_type)
142+
template_config = client.dq_template_config_cache.get_template_config(
143+
rule_type.value
144+
)
143145

144146
asset_for_validation, target_table_asset = (
145147
DataQualityRule.Attributes._fetch_assets_for_row_scope_validation(
@@ -275,16 +277,20 @@
275277
else None
276278
) # type: ignore[attr-defined]
277279

278-
retrieved_rule_type = retrieved_template_rule_name
279-
template_config = client.dq_template_config_cache.get_template_config(
280-
retrieved_rule_type
281-
)
280+
template_config = None
281+
if retrieved_template_rule_name:
282+
template_config = client.dq_template_config_cache.get_template_config(
283+
retrieved_template_rule_name
284+
)
282285

283-
final_rule_conditions = rule_conditions or (
284-
search_result.dq_rule_config_arguments.dq_rule_config_rule_conditions # type: ignore[attr-defined]
285-
if search_result.dq_rule_config_arguments is not None # type: ignore[attr-defined]
286-
else None
287-
)
286+
if rule_conditions:
287+
final_rule_conditions = rule_conditions
288+
elif search_result.dq_rule_config_arguments is not None: # type: ignore[attr-defined]
289+
final_rule_conditions = (
290+
search_result.dq_rule_config_arguments.dq_rule_config_rule_conditions # type: ignore[attr-defined]
291+
)
292+
else:
293+
final_rule_conditions = None
288294

289295
final_row_scope_filtering_enabled = (
290296
row_scope_filtering_enabled or retrieved_row_scope_filtering_enabled
@@ -301,17 +307,26 @@
301307
else:
302308
target_table_asset = None
303309

304-
validated_threshold_operator = (
305-
DataQualityRule.Attributes._validate_template_features(
306-
retrieved_rule_type,
307-
final_rule_conditions,
308-
final_row_scope_filtering_enabled,
309-
template_config,
310-
threshold_compare_operator or retrieved_threshold_compare_operator,
311-
retrieved_asset,
312-
target_table_asset,
313-
)
314-
)
310+
validated_threshold_operator = None
311+
if retrieved_template_rule_name and template_config:
312+
try:
313+
retrieved_rule_type = DataQualityRuleTemplateType(
314+
retrieved_template_rule_name
315+
)
316+
validated_threshold_operator = (
317+
DataQualityRule.Attributes._validate_template_features(
318+
retrieved_rule_type,
319+
final_rule_conditions,
320+
final_row_scope_filtering_enabled,
321+
template_config,
322+
threshold_compare_operator
323+
or retrieved_threshold_compare_operator,
324+
retrieved_asset,
325+
target_table_asset,
326+
)
327+
)
328+
except ValueError:
329+
pass
315330

316331
final_compare_operator = (
317332
validated_threshold_operator

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

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@
7878

7979
@staticmethod
8080
def _validate_template_features(
81-
rule_type: str,
81+
rule_type: DataQualityRuleTemplateType,
8282
rule_conditions: Optional[str],
8383
row_scope_filtering_enabled: Optional[bool],
8484
template_config: Optional[dict],
@@ -98,7 +98,7 @@
9898
and config.dq_rule_template_config_rule_conditions is None
9999
):
100100
raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters(
101-
rule_type, "rule conditions"
101+
rule_type.value, "rule conditions"
102102
)
103103

104104
if row_scope_filtering_enabled:
@@ -107,7 +107,7 @@
107107
)
108108
if "dqRuleRowScopeFilteringEnabled" not in str(advanced_settings):
109109
raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters(
110-
rule_type, "row scope filtering"
110+
rule_type.value, "row scope filtering"
111111
)
112112

113113
if asset and not getattr(
@@ -228,7 +228,7 @@
228228
*,
229229
client: AtlanClient,
230230
rule_name: str,
231-
rule_type: str,
231+
rule_type: DataQualityRuleTemplateType,
232232
asset: Asset,
233233
threshold_compare_operator: DataQualityRuleThresholdCompareOperator,
234234
threshold_value: int,
@@ -243,11 +243,13 @@
243243
row_scope_filtering_enabled: Optional[bool] = False,
244244
) -> DataQualityRule.Attributes:
245245
template_config = client.dq_template_config_cache.get_template_config(
246-
rule_type
246+
rule_type.value
247247
)
248248

249249
if template_config is None:
250-
raise ErrorCode.DQ_RULE_NOT_FOUND.exception_with_parameters(rule_type)
250+
raise ErrorCode.DQ_RULE_NOT_FOUND.exception_with_parameters(
251+
rule_type.value
252+
)
251253

252254
template_rule_name = template_config.get("name")
253255
template_qualified_name = template_config.get("qualified_name")

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ from pyatlan.model.enums import (
66
DataQualityRuleAlertPriority,
77
DataQualityRuleCustomSQLReturnType,
88
DataQualityRuleStatus,
9+
DataQualityRuleTemplateType,
910
DataQualityRuleThresholdCompareOperator,
1011
DataQualityRuleThresholdUnit,
1112
DataQualitySourceSyncStatus,

‎pyatlan/model/assets/core/data_quality_rule.py‎

Lines changed: 48 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
DataQualityRuleAlertPriority,
2020
DataQualityRuleCustomSQLReturnType,
2121
DataQualityRuleStatus,
22+
DataQualityRuleTemplateType,
2223
DataQualityRuleThresholdCompareOperator,
2324
DataQualityRuleThresholdUnit,
2425
DataQualitySourceSyncStatus,
@@ -88,7 +89,7 @@ def custom_sql_creator(
8889
attributes = DataQualityRule.Attributes.creator(
8990
client=client,
9091
rule_name=rule_name,
91-
rule_type="Custom SQL",
92+
rule_type=DataQualityRuleTemplateType.CUSTOM_SQL,
9293
asset=asset,
9394
threshold_compare_operator=threshold_compare_operator,
9495
threshold_value=threshold_value,
@@ -108,7 +109,7 @@ def table_level_rule_creator(
108109
cls,
109110
*,
110111
client: AtlanClient,
111-
rule_type: str,
112+
rule_type: DataQualityRuleTemplateType,
112113
asset: Asset,
113114
threshold_compare_operator: DataQualityRuleThresholdCompareOperator,
114115
threshold_value: int,
@@ -156,7 +157,7 @@ def column_level_rule_creator(
156157
cls,
157158
*,
158159
client: AtlanClient,
159-
rule_type: str,
160+
rule_type: DataQualityRuleTemplateType,
160161
asset: Asset,
161162
column: Asset,
162163
threshold_value: int,
@@ -186,7 +187,9 @@ def column_level_rule_creator(
186187
alert_priority,
187188
],
188189
)
189-
template_config = client.dq_template_config_cache.get_template_config(rule_type)
190+
template_config = client.dq_template_config_cache.get_template_config(
191+
rule_type.value
192+
)
190193

191194
asset_for_validation, target_table_asset = (
192195
DataQualityRule.Attributes._fetch_assets_for_row_scope_validation(
@@ -322,16 +325,20 @@ def updater(
322325
else None
323326
) # type: ignore[attr-defined]
324327

325-
retrieved_rule_type = retrieved_template_rule_name
326-
template_config = client.dq_template_config_cache.get_template_config(
327-
retrieved_rule_type
328-
)
328+
template_config = None
329+
if retrieved_template_rule_name:
330+
template_config = client.dq_template_config_cache.get_template_config(
331+
retrieved_template_rule_name
332+
)
329333

330-
final_rule_conditions = rule_conditions or (
331-
search_result.dq_rule_config_arguments.dq_rule_config_rule_conditions # type: ignore[attr-defined]
332-
if search_result.dq_rule_config_arguments is not None # type: ignore[attr-defined]
333-
else None
334-
)
334+
if rule_conditions:
335+
final_rule_conditions = rule_conditions
336+
elif search_result.dq_rule_config_arguments is not None: # type: ignore[attr-defined]
337+
final_rule_conditions = (
338+
search_result.dq_rule_config_arguments.dq_rule_config_rule_conditions # type: ignore[attr-defined]
339+
)
340+
else:
341+
final_rule_conditions = None
335342

336343
final_row_scope_filtering_enabled = (
337344
row_scope_filtering_enabled or retrieved_row_scope_filtering_enabled
@@ -348,17 +355,26 @@ def updater(
348355
else:
349356
target_table_asset = None
350357

351-
validated_threshold_operator = (
352-
DataQualityRule.Attributes._validate_template_features(
353-
retrieved_rule_type,
354-
final_rule_conditions,
355-
final_row_scope_filtering_enabled,
356-
template_config,
357-
threshold_compare_operator or retrieved_threshold_compare_operator,
358-
retrieved_asset,
359-
target_table_asset,
360-
)
361-
)
358+
validated_threshold_operator = None
359+
if retrieved_template_rule_name and template_config:
360+
try:
361+
retrieved_rule_type = DataQualityRuleTemplateType(
362+
retrieved_template_rule_name
363+
)
364+
validated_threshold_operator = (
365+
DataQualityRule.Attributes._validate_template_features(
366+
retrieved_rule_type,
367+
final_rule_conditions,
368+
final_row_scope_filtering_enabled,
369+
template_config,
370+
threshold_compare_operator
371+
or retrieved_threshold_compare_operator,
372+
retrieved_asset,
373+
target_table_asset,
374+
)
375+
)
376+
except ValueError:
377+
pass
362378

363379
final_compare_operator = (
364380
validated_threshold_operator
@@ -1190,7 +1206,7 @@ def _get_template_config_value(
11901206

11911207
@staticmethod
11921208
def _validate_template_features(
1193-
rule_type: str,
1209+
rule_type: DataQualityRuleTemplateType,
11941210
rule_conditions: Optional[str],
11951211
row_scope_filtering_enabled: Optional[bool],
11961212
template_config: Optional[dict],
@@ -1210,7 +1226,7 @@ def _validate_template_features(
12101226
and config.dq_rule_template_config_rule_conditions is None
12111227
):
12121228
raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters(
1213-
rule_type, "rule conditions"
1229+
rule_type.value, "rule conditions"
12141230
)
12151231

12161232
if row_scope_filtering_enabled:
@@ -1219,7 +1235,7 @@ def _validate_template_features(
12191235
)
12201236
if "dqRuleRowScopeFilteringEnabled" not in str(advanced_settings):
12211237
raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters(
1222-
rule_type, "row scope filtering"
1238+
rule_type.value, "row scope filtering"
12231239
)
12241240

12251241
if asset and not getattr(
@@ -1340,7 +1356,7 @@ def creator(
13401356
*,
13411357
client: AtlanClient,
13421358
rule_name: str,
1343-
rule_type: str,
1359+
rule_type: DataQualityRuleTemplateType,
13441360
asset: Asset,
13451361
threshold_compare_operator: DataQualityRuleThresholdCompareOperator,
13461362
threshold_value: int,
@@ -1355,11 +1371,13 @@ def creator(
13551371
row_scope_filtering_enabled: Optional[bool] = False,
13561372
) -> DataQualityRule.Attributes:
13571373
template_config = client.dq_template_config_cache.get_template_config(
1358-
rule_type
1374+
rule_type.value
13591375
)
13601376

13611377
if template_config is None:
1362-
raise ErrorCode.DQ_RULE_NOT_FOUND.exception_with_parameters(rule_type)
1378+
raise ErrorCode.DQ_RULE_NOT_FOUND.exception_with_parameters(
1379+
rule_type.value
1380+
)
13631381

13641382
template_rule_name = template_config.get("name")
13651383
template_qualified_name = template_config.get("qualified_name")

‎pyatlan/model/enums.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2451,6 +2451,31 @@ class DataQualityRuleThresholdCompareOperator(str, Enum):
24512451
LESS_THAN = "LT"
24522452

24532453

2454+
class DataQualityRuleTemplateType(str, Enum):
2455+
BLANK_COUNT = "BLANK_COUNT"
2456+
BLANK_PERCENTAGE = "BLANK_PERCENTAGE"
2457+
NULL_COUNT = "NULL_COUNT"
2458+
NULL_PERCENTAGE = "NULL_PERCENTAGE"
2459+
FRESHNESS = "FRESHNESS"
2460+
AVERAGE = "AVERAGE"
2461+
MAX_VALUE = "MAX_VALUE"
2462+
MIN_VALUE = "MIN_VALUE"
2463+
STANDARD_DEVIATION = "STANDARD_DEVIATION"
2464+
DUPLICATE_COUNT = "DUPLICATE_COUNT"
2465+
UNIQUE_COUNT = "UNIQUE_COUNT"
2466+
ROW_COUNT = "ROW_COUNT"
2467+
CUSTOM_SQL = "Custom SQL"
2468+
REGEX_MATCH = "REGEX_MATCH"
2469+
VALID_STRING_VALUES = "VALID_STRING_VALUES"
2470+
VALID_STRING_VALUES_REFERENCE = "VALID_STRING_VALUES_REFERENCE"
2471+
STRING_LENGTH = "STRING_LENGTH"
2472+
RECON_ROW_COUNT = "RECON_ROW_COUNT"
2473+
RECON_AVERAGE = "RECON_AVERAGE"
2474+
RECON_SUM = "RECON_SUM"
2475+
RECON_DUPLICATE_COUNT = "RECON_DUPLICATE_COUNT"
2476+
RECON_UNIQUE_COUNT = "RECON_UNIQUE_COUNT"
2477+
2478+
24542479
class DataQualityRuleTemplateConfigRuleConditions(str, Enum):
24552480
# String Length conditions
24562481
STRING_LENGTH_BETWEEN = "STRING_LENGTH_BETWEEN"

0 commit comments

Comments
 (0)