Skip to content

Commit 4d61f37

Browse files
UdaySagar-GitUdaySagar-Git
committed
chore: update jinja templates
1 parent 330da12 commit 4d61f37

3 files changed

Lines changed: 114 additions & 18 deletions

File tree

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

Lines changed: 43 additions & 18 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(
@@ -47,6 +48,7 @@
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,
@@ -64,6 +66,7 @@
6466
threshold_compare_operator: DataQualityRuleThresholdCompareOperator,
6567
threshold_value: int,
6668
alert_priority: DataQualityRuleAlertPriority,
69+
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
6770
) -> DataQualityRule:
6871
validate_required_fields(
6972
[
@@ -93,7 +96,7 @@
9396
alert_priority=alert_priority,
9497
rule_name=None,
9598
column=None,
96-
threshold_unit=None,
99+
threshold_unit=threshold_unit,
97100
dimension=None,
98101
custom_sql=None,
99102
description=None,
@@ -138,20 +141,11 @@
138141
)
139142
template_config = client.dq_template_config_cache.get_template_config(rule_type)
140143

141-
asset_for_validation = asset
142-
if row_scope_filtering_enabled and asset.qualified_name:
143-
from pyatlan.model.fluent_search import FluentSearch
144-
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]
144+
asset_for_validation, target_table_asset = (
145+
DataQualityRule.Attributes._fetch_assets_for_row_scope_validation(
146+
client, asset, rule_conditions, row_scope_filtering_enabled
147+
)
148+
)
155149

156150
validated_threshold_operator = (
157151
DataQualityRule.Attributes._validate_template_features(
@@ -161,6 +155,7 @@
161155
template_config,
162156
threshold_compare_operator,
163157
asset_for_validation,
158+
target_table_asset,
164159
)
165160
)
166161

@@ -202,6 +197,7 @@
202197
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
203198
dimension: Optional[DataQualityDimension] = None,
204199
custom_sql: Optional[str] = None,
200+
custom_sql_return_type: Optional[DataQualityRuleCustomSQLReturnType] = None,
205201
rule_name: Optional[str] = None,
206202
description: Optional[str] = None,
207203
rule_conditions: Optional[str] = None,
@@ -224,6 +220,7 @@
224220
.include_on_results(DataQualityRule.DQ_RULE_ALERT_PRIORITY)
225221
.include_on_results(DataQualityRule.DISPLAY_NAME)
226222
.include_on_results(DataQualityRule.DQ_RULE_CUSTOM_SQL)
223+
.include_on_results(DataQualityRule.DQ_RULE_CUSTOM_SQL_RETURN_TYPE)
227224
.include_on_results(DataQualityRule.USER_DESCRIPTION)
228225
.include_on_results(DataQualityRule.DQ_RULE_DIMENSION)
229226
.include_on_results(DataQualityRule.DQ_RULE_CONFIG_ARGUMENTS)
@@ -242,6 +239,9 @@
242239
search_result = results.current_page()[0]
243240

244241
retrieved_custom_sql = search_result.dq_rule_custom_s_q_l # type: ignore[attr-defined]
242+
retrieved_custom_sql_return_type = (
243+
search_result.dq_rule_custom_s_q_l_return_type # type: ignore[attr-defined]
244+
)
245245
retrieved_rule_name = search_result.display_name
246246
retrieved_dimension = search_result.dq_rule_dimension # type: ignore[attr-defined]
247247
retrieved_column = search_result.dq_rule_base_column # type: ignore[attr-defined]
@@ -279,14 +279,37 @@
279279
template_config = client.dq_template_config_cache.get_template_config(
280280
retrieved_rule_type
281281
)
282+
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+
)
288+
289+
final_row_scope_filtering_enabled = (
290+
row_scope_filtering_enabled or retrieved_row_scope_filtering_enabled
291+
)
292+
if retrieved_asset:
293+
retrieved_asset, target_table_asset = (
294+
DataQualityRule.Attributes._fetch_assets_for_row_scope_validation(
295+
client,
296+
retrieved_asset,
297+
final_rule_conditions,
298+
final_row_scope_filtering_enabled,
299+
)
300+
)
301+
else:
302+
target_table_asset = None
303+
282304
validated_threshold_operator = (
283305
DataQualityRule.Attributes._validate_template_features(
284306
retrieved_rule_type,
285307
rule_conditions,
286-
row_scope_filtering_enabled,
308+
final_row_scope_filtering_enabled,
287309
template_config,
288310
threshold_compare_operator or retrieved_threshold_compare_operator,
289311
retrieved_asset,
312+
target_table_asset,
290313
)
291314
)
292315

@@ -310,8 +333,7 @@
310333
),
311334
dq_rule_base_dataset_qualified_name=retrieved_asset.qualified_name,
312335
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,
336+
dq_rule_row_scope_filtering_enabled=final_row_scope_filtering_enabled,
315337
dq_rule_base_dataset=retrieved_asset,
316338
qualified_name=qualified_name,
317339
dq_rule_dimension=dimension or retrieved_dimension,
@@ -329,6 +351,9 @@
329351
if custom_sql is not None:
330352
attr_dq.dq_rule_custom_s_q_l = custom_sql
331353
attr_dq.display_name = rule_name or retrieved_rule_name
354+
attr_dq.dq_rule_custom_s_q_l_return_type = (
355+
custom_sql_return_type or retrieved_custom_sql_return_type
356+
)
332357
if description is not None:
333358
attr_dq.user_description = description or retrieved_description
334359

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,60 @@
11

2+
@staticmethod
3+
def _fetch_assets_for_row_scope_validation(
4+
client: AtlanClient,
5+
base_asset: Asset,
6+
rule_conditions: Optional[str],
7+
row_scope_filtering_enabled: bool,
8+
) -> tuple[Asset, Optional[Asset]]:
9+
asset_for_validation = base_asset
10+
target_table_asset = None
11+
12+
if not row_scope_filtering_enabled:
13+
return asset_for_validation, target_table_asset
14+
15+
# Extract target_table from rule_conditions
16+
target_table_qualified_name = None
17+
if rule_conditions:
18+
try:
19+
rule_conditions_json = json.loads(rule_conditions)
20+
conditions = rule_conditions_json.get("conditions", [])
21+
if conditions:
22+
condition_value = conditions[0].get("value", {})
23+
target_table_qualified_name = condition_value.get(
24+
"target_table"
25+
)
26+
except (json.JSONDecodeError, (KeyError, TypeError)):
27+
pass
28+
29+
qualified_names_to_search = []
30+
if base_asset.qualified_name:
31+
qualified_names_to_search.append(base_asset.qualified_name)
32+
if target_table_qualified_name:
33+
qualified_names_to_search.append(target_table_qualified_name)
34+
35+
if qualified_names_to_search:
36+
from pyatlan.model.fluent_search import FluentSearch
37+
38+
search_request = (
39+
FluentSearch()
40+
.where(Asset.QUALIFIED_NAME.within(qualified_names_to_search))
41+
.include_on_results(
42+
Asset.ASSET_DQ_ROW_SCOPE_FILTER_COLUMN_QUALIFIED_NAME
43+
)
44+
).to_request()
45+
results = client.asset.search(search_request)
46+
47+
for result in results.current_page():
48+
if result.qualified_name == base_asset.qualified_name:
49+
asset_for_validation = result
50+
elif (
51+
target_table_qualified_name
52+
and result.qualified_name == target_table_qualified_name
53+
):
54+
target_table_asset = result
55+
56+
return asset_for_validation, target_table_asset
57+
258
@staticmethod
359
def _get_template_config_value(
460
config_value: str,
@@ -30,6 +86,7 @@
3086
DataQualityRuleThresholdCompareOperator
3187
] = None,
3288
asset: Optional[Asset] = None,
89+
target_table_asset: Optional[Asset] = None,
3390
) -> Optional[DataQualityRuleThresholdCompareOperator]:
3491
if not template_config or not template_config.get("config"):
3592
return None
@@ -62,6 +119,16 @@
62119
getattr(asset, "qualified_name", "unknown")
63120
)
64121

122+
if target_table_asset:
123+
if not getattr(
124+
target_table_asset,
125+
"asset_d_q_row_scope_filter_column_qualified_name",
126+
None,
127+
):
128+
raise ErrorCode.DQ_ROW_SCOPE_FILTER_COLUMN_MISSING.exception_with_parameters(
129+
getattr(target_table_asset, "qualified_name", "unknown")
130+
)
131+
65132
if rule_conditions:
66133
allowed_rule_conditions = (
67134
DataQualityRule.Attributes._get_template_config_value(
@@ -170,6 +237,7 @@
170237
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
171238
dimension: Optional[DataQualityDimension] = None,
172239
custom_sql: Optional[str] = None,
240+
custom_sql_return_type: Optional[DataQualityRuleCustomSQLReturnType] = None,
173241
description: Optional[str] = None,
174242
rule_conditions: Optional[str] = None,
175243
row_scope_filtering_enabled: Optional[bool] = False,
@@ -229,6 +297,8 @@
229297
if custom_sql is not None:
230298
attr_dq.dq_rule_custom_s_q_l = custom_sql
231299
attr_dq.display_name = rule_name
300+
if custom_sql_return_type is not None:
301+
attr_dq.dq_rule_custom_s_q_l_return_type = custom_sql_return_type
232302
if description is not None:
233303
attr_dq.user_description = description
234304

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ from pyatlan.model.enums import (
44
DataQualityDimension,
55
DataQualityResult,
66
DataQualityRuleAlertPriority,
7+
DataQualityRuleCustomSQLReturnType,
78
DataQualityRuleStatus,
89
DataQualityRuleThresholdCompareOperator,
910
DataQualityRuleThresholdUnit,

0 commit comments

Comments
 (0)