Skip to content

Commit c527116

Browse files
committed
chore: update jinja templates
1 parent 650823c commit c527116

3 files changed

Lines changed: 172 additions & 16 deletions

File tree

pyatlan/generator/templates/methods/asset/alpha__d_q_rule.jinja2

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -113,16 +113,17 @@
113113
alert_priority: alpha_DQRuleAlertPriority,
114114
threshold_compare_operator: Optional[
115115
alpha_DQRuleThresholdCompareOperator
116-
] = alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL,
116+
] = None,
117117
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
118+
rule_conditions: Optional[str] = None,
119+
row_scope_filtering_enabled: Optional[bool] = False,
118120
) -> alpha_DQRule:
119121
validate_required_fields(
120122
[
121123
"client",
122124
"rule_type",
123125
"asset",
124126
"column",
125-
"threshold_compare_operator",
126127
"threshold_value",
127128
"alert_priority",
128129
],
@@ -131,11 +132,18 @@
131132
rule_type,
132133
asset,
133134
column,
134-
threshold_compare_operator,
135135
threshold_value,
136136
alert_priority,
137137
],
138138
)
139+
template_config = client.dq_template_config_cache.get_template_config(rule_type)
140+
threshold_compare_operator = alpha_DQRule.Attributes._validate_template_features(
141+
rule_type,
142+
rule_conditions,
143+
row_scope_filtering_enabled,
144+
template_config,
145+
threshold_compare_operator,
146+
)
139147

140148
attributes = alpha_DQRule.Attributes.creator(
141149
client=client,
@@ -150,6 +158,8 @@
150158
dimension=None,
151159
custom_sql=None,
152160
description=None,
161+
rule_conditions=rule_conditions,
162+
row_scope_filtering_enabled=row_scope_filtering_enabled,
153163
)
154164
return cls(attributes=attributes)
155165

@@ -169,6 +179,8 @@
169179
custom_sql: Optional[str] = None,
170180
rule_name: Optional[str] = None,
171181
description: Optional[str] = None,
182+
rule_conditions: Optional[str] = None,
183+
row_scope_filtering_enabled: Optional[bool] = False,
172184
) -> SelfAsset:
173185
from pyatlan.model.fluent_search import FluentSearch
174186

@@ -190,6 +202,7 @@
190202
.include_on_results(alpha_DQRule.USER_DESCRIPTION)
191203
.include_on_results(alpha_DQRule.ALPHADQ_RULE_DIMENSION)
192204
.include_on_results(alpha_DQRule.ALPHADQ_RULE_CONFIG_ARGUMENTS)
205+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_ROW_SCOPE_FILTERING_ENABLED)
193206
.include_on_results(alpha_DQRule.ALPHADQ_RULE_SOURCE_SYNC_STATUS)
194207
.include_on_results(alpha_DQRule.ALPHADQ_RULE_STATUS)
195208
).to_request()
@@ -208,6 +221,9 @@
208221
retrieved_dimension = search_result.alpha_dq_rule_dimension # type: ignore[attr-defined]
209222
retrieved_column = search_result.alpha_dq_rule_base_column # type: ignore[attr-defined]
210223
retrieved_alert_priority = search_result.alpha_dq_rule_alert_priority # type: ignore[attr-defined]
224+
retrieved_row_scope_filtering_enabled = (
225+
search_result.alpha_dq_rule_row_scope_filtering_enabled
226+
) # type: ignore[attr-defined]
211227
retrieved_description = search_result.user_description
212228
retrieved_asset = search_result.alpha_dq_rule_base_dataset # type: ignore[attr-defined]
213229
retrieved_template_rule_name = search_result.alpha_dq_rule_template_name # type: ignore[attr-defined]
@@ -234,35 +250,50 @@
234250
else None
235251
) # type: ignore[attr-defined]
236252

253+
retrieved_rule_type = retrieved_template_rule_name
254+
template_config = client.dq_template_config_cache.get_template_config(
255+
retrieved_rule_type
256+
)
257+
validated_threshold_operator = alpha_DQRule.Attributes._validate_template_features(
258+
retrieved_rule_type,
259+
rule_conditions,
260+
row_scope_filtering_enabled,
261+
template_config,
262+
threshold_compare_operator or retrieved_threshold_compare_operator,
263+
)
264+
237265
config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw(
238266
is_alert_enabled=True,
239267
custom_sql=custom_sql or retrieved_custom_sql,
240268
display_name=rule_name or retrieved_rule_name,
241269
dimension=dimension or retrieved_dimension,
242-
compare_operator=threshold_compare_operator
243-
or retrieved_threshold_compare_operator,
270+
compare_operator=validated_threshold_operator,
244271
threshold_value=threshold_value or retrieved_threshold_value,
245272
threshold_unit=threshold_unit or retrieved_threshold_unit,
246273
column=retrieved_column,
247274
dq_priority=alert_priority or retrieved_alert_priority,
248275
description=description or retrieved_description,
276+
rule_conditions=rule_conditions,
277+
row_scope_filtering_enabled=row_scope_filtering_enabled,
249278
)
250279

251280
attr_dq = cls.Attributes(
252281
name="",
253282
alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments(
254283
alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject(
255-
alpha_dq_rule_threshold_compare_operator=threshold_compare_operator
256-
or retrieved_threshold_compare_operator,
284+
alpha_dq_rule_threshold_compare_operator=validated_threshold_operator,
257285
alpha_dq_rule_threshold_value=threshold_value
258286
or retrieved_threshold_value,
259287
alpha_dq_rule_threshold_unit=threshold_unit
260288
or retrieved_threshold_unit,
261289
),
262290
alpha_dq_rule_config_arguments_raw=config_arguments_raw,
291+
alpha_dq_rule_config_rule_conditions=rule_conditions,
263292
),
264293
alpha_dq_rule_base_dataset_qualified_name=retrieved_asset.qualified_name,
265294
alpha_dq_rule_alert_priority=alert_priority or retrieved_alert_priority,
295+
alpha_dq_rule_row_scope_filtering_enabled=row_scope_filtering_enabled
296+
or retrieved_row_scope_filtering_enabled,
266297
alpha_dq_rule_base_dataset=retrieved_asset,
267298
qualified_name=qualified_name,
268299
alpha_dq_rule_dimension=dimension or retrieved_dimension,

pyatlan/generator/templates/methods/attribute/alpha__d_q_rule.jinja2

Lines changed: 133 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,114 @@
1+
@staticmethod
2+
def _get_template_config_value(
3+
config_value: str, property_name: str = None, value_key: str = "default"
4+
):
5+
if not config_value:
6+
return None
7+
8+
try:
9+
config_json = json.loads(config_value)
10+
11+
if property_name:
12+
properties = config_json.get("properties", {})
13+
field = properties.get(property_name, {})
14+
return field.get(value_key)
15+
else:
16+
return config_json.get(value_key)
17+
except (json.JSONDecodeError, KeyError):
18+
return None
19+
20+
@staticmethod
21+
def _validate_template_features(
22+
rule_type: str,
23+
rule_conditions: Optional[str],
24+
row_scope_filtering_enabled: Optional[bool],
25+
template_config: Optional[dict],
26+
threshold_compare_operator: Optional[
27+
alpha_DQRuleThresholdCompareOperator
28+
] = None,
29+
) -> alpha_DQRuleThresholdCompareOperator:
30+
if not template_config or not template_config.get("config"):
31+
return
32+
33+
config = template_config["config"]
34+
35+
if (
36+
rule_conditions
37+
and config.alpha_dq_rule_template_config_rule_conditions is None
38+
):
39+
raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters(
40+
rule_type, "rule conditions"
41+
)
42+
43+
if row_scope_filtering_enabled:
44+
advanced_settings = config.alpha_dq_rule_template_advanced_settings or ""
45+
if "alpha_dqRuleRowScopeFilteringEnabled" not in str(advanced_settings):
46+
raise ErrorCode.DQ_RULE_TYPE_NOT_SUPPORTED.exception_with_parameters(
47+
rule_type, "row scope filtering"
48+
)
49+
50+
if rule_conditions:
51+
allowed_rule_conditions = alpha_DQRule.Attributes._get_template_config_value(
52+
config.alpha_dq_rule_template_config_rule_conditions, None, "enum"
53+
)
54+
if allowed_rule_conditions:
55+
try:
56+
rule_conditions_json = json.loads(rule_conditions)
57+
conditions = rule_conditions_json.get("conditions", [])
58+
if len(conditions) != 1:
59+
raise ErrorCode.DQ_RULE_CONDITIONS_INVALID.exception_with_parameters(
60+
f"exactly one condition required, found {len(conditions)}"
61+
)
62+
condition_type = conditions[0].get("type")
63+
except json.JSONDecodeError:
64+
condition_type = rule_conditions
65+
66+
if condition_type not in allowed_rule_conditions:
67+
raise ErrorCode.DQ_RULE_CONDITIONS_INVALID.exception_with_parameters(
68+
f"condition type '{condition_type}' not supported, allowed: {allowed_rule_conditions}"
69+
)
70+
71+
if threshold_compare_operator is None:
72+
return alpha_DQRuleThresholdCompareOperator.EQUAL
73+
elif (
74+
threshold_compare_operator != alpha_DQRuleThresholdCompareOperator.EQUAL
75+
):
76+
raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters(
77+
f"threshold_compare_operator={threshold_compare_operator.value}",
78+
"threshold_compare_operator",
79+
"EQUAL when rule_conditions are provided",
80+
)
81+
82+
if threshold_compare_operator is not None:
83+
allowed_operators = alpha_DQRule.Attributes._get_template_config_value(
84+
config.alpha_dq_rule_template_config_threshold_object,
85+
"alpha_dqRuleTemplateConfigThresholdCompareOperator",
86+
"enum",
87+
)
88+
if (
89+
allowed_operators
90+
and threshold_compare_operator.value not in allowed_operators
91+
):
92+
raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters(
93+
f"threshold_compare_operator={threshold_compare_operator.value}",
94+
"threshold_compare_operator",
95+
f"must be one of {allowed_operators}",
96+
)
97+
elif threshold_compare_operator is None:
98+
default_value = alpha_DQRule.Attributes._get_template_config_value(
99+
config.alpha_dq_rule_template_config_threshold_object,
100+
"alpha_dqRuleTemplateConfigThresholdCompareOperator",
101+
"default",
102+
)
103+
if default_value:
104+
threshold_compare_operator = alpha_DQRuleThresholdCompareOperator(
105+
default_value
106+
)
107+
108+
return (
109+
threshold_compare_operator
110+
or alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL
111+
)
1112

2113
@staticmethod
3114
def _generate_config_arguments_raw(
@@ -12,6 +123,8 @@
12123
column: Optional[Asset] = None,
13124
dq_priority: alpha_DQRuleAlertPriority,
14125
description: Optional[str] = None,
126+
rule_conditions: Optional[str] = None,
127+
row_scope_filtering_enabled: Optional[bool] = None,
15128
) -> str:
16129
config = {
17130
"isAlertEnabled": is_alert_enabled,
@@ -40,6 +153,16 @@
40153
if dimension is not None:
41154
config["alpha_dqRuleTemplateConfigDimension"] = dimension
42155

156+
if rule_conditions is not None:
157+
config["alpha_dqRuleTemplateConfigRuleConditions"] = json.loads(
158+
rule_conditions
159+
)
160+
161+
if row_scope_filtering_enabled is not None:
162+
config[
163+
"alpha_dqRuleTemplateAdvancedSettings.alpha_dqRuleRowScopeFilteringEnabled"
164+
] = row_scope_filtering_enabled
165+
43166
return json.dumps(config)
44167

45168
@staticmethod
@@ -81,6 +204,8 @@
81204
dimension: Optional[alpha_DQDimension] = None,
82205
custom_sql: Optional[str] = None,
83206
description: Optional[str] = None,
207+
rule_conditions: Optional[str] = None,
208+
row_scope_filtering_enabled: Optional[bool] = False,
84209
) -> alpha_DQRule.Attributes:
85210
template_config = client.dq_template_config_cache.get_template_config(
86211
rule_type
@@ -98,16 +223,11 @@
98223
if threshold_unit is None:
99224
config = template_config.get("config")
100225
if config is not None:
101-
threashold_object = (
102-
config.alpha_dq_rule_template_config_threshold_object
103-
)
104-
threashold_object_json = json.loads(threashold_object)
105-
properties = threashold_object_json.get("properties", {})
106-
threshold_unit_field = properties.get(
107-
"alpha_dqRuleTemplateConfigThresholdUnit", {}
226+
threshold_unit = alpha_DQRule.Attributes._get_template_config_value(
227+
config.alpha_dq_rule_template_config_threshold_object,
228+
"alpha_dqRuleTemplateConfigThresholdUnit",
229+
"default",
108230
)
109-
default_value = threshold_unit_field.get("default")
110-
threshold_unit = default_value
111231

112232
config_arguments_raw = (
113233
alpha_DQRule.Attributes._generate_config_arguments_raw(
@@ -121,6 +241,8 @@
121241
column=column,
122242
dq_priority=alert_priority,
123243
description=description,
244+
rule_conditions=rule_conditions,
245+
row_scope_filtering_enabled=row_scope_filtering_enabled,
124246
)
125247
)
126248

@@ -133,9 +255,11 @@
133255
alpha_dq_rule_threshold_unit=threshold_unit,
134256
),
135257
alpha_dq_rule_config_arguments_raw=config_arguments_raw,
258+
alpha_dq_rule_config_rule_conditions=rule_conditions,
136259
),
137260
alpha_dq_rule_base_dataset_qualified_name=asset.qualified_name,
138261
alpha_dq_rule_alert_priority=alert_priority,
262+
alpha_dq_rule_row_scope_filtering_enabled=row_scope_filtering_enabled,
139263
alpha_dq_rule_source_sync_status=alpha_DQSourceSyncStatus.IN_PROGRESS,
140264
alpha_dq_rule_status=alpha_DQRuleStatus.ACTIVE,
141265
alpha_dq_rule_base_dataset=asset,

pyatlan/generator/templates/methods/imports/alpha__d_q_rule.jinja2

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ from pyatlan.model.enums import (
88
alpha_DQRuleThresholdCompareOperator,
99
alpha_DQRuleThresholdUnit,
1010
alpha_DQSourceSyncStatus,
11+
alpha_dqRuleTemplateConfigRuleConditions,
1112
)
1213
from pyatlan.model.structs import (
1314
alpha_DQRuleConfigArguments,

0 commit comments

Comments
 (0)