Skip to content

Commit 9789c07

Browse files
vaibhavatlanvaibhavatlan
committed
Generator code added
1 parent 4e6560a commit 9789c07

File tree

3 files changed

+448
-1
lines changed

3 files changed

+448
-1
lines changed

‎pyatlan/generator/templates/imports.jinja2

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,6 @@ from pyatlan.model.structs import (
113113
CustomRatings,
114114
alpha_DQRuleConfigArguments,
115115
alpha_DQRuleTemplateConfig,
116-
alpha_DQRuleThresholdUnit,
117116
ResponseValue,
118117
FormField,
119118
DatabricksAIModelVersionMetric

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

Lines changed: 288 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,288 @@
1+
2+
@classmethod
3+
@init_guid
4+
def custom_sql_creator(
5+
cls,
6+
*,
7+
client: AtlanClient,
8+
rule_name: str,
9+
asset: Asset,
10+
custom_sql: str,
11+
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
12+
threshold_value: int,
13+
alert_priority: alpha_DQRuleAlertPriority,
14+
dimension: alpha_DQDimension,
15+
description: Optional[str] = None,
16+
) -> alpha_DQRule:
17+
validate_required_fields(
18+
[
19+
"client",
20+
"rule_name",
21+
"asset",
22+
"threshold_compare_operator",
23+
"threshold_value",
24+
"alert_priority",
25+
"dimension",
26+
"custom_sql",
27+
],
28+
[
29+
client,
30+
rule_name,
31+
asset,
32+
threshold_compare_operator,
33+
threshold_value,
34+
alert_priority,
35+
dimension,
36+
custom_sql,
37+
],
38+
)
39+
40+
attributes = alpha_DQRule.Attributes.creator(
41+
client=client,
42+
rule_name=rule_name,
43+
rule_type="Custom SQL",
44+
asset=asset,
45+
threshold_compare_operator=threshold_compare_operator,
46+
threshold_value=threshold_value,
47+
alert_priority=alert_priority,
48+
dimension=dimension,
49+
custom_sql=custom_sql,
50+
description=description,
51+
column=None,
52+
threshold_unit=None,
53+
)
54+
return cls(attributes=attributes)
55+
56+
@classmethod
57+
@init_guid
58+
def table_level_rule_creator(
59+
cls,
60+
*,
61+
client: AtlanClient,
62+
rule_type: str,
63+
asset: Asset,
64+
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
65+
threshold_value: int,
66+
alert_priority: alpha_DQRuleAlertPriority,
67+
) -> alpha_DQRule:
68+
validate_required_fields(
69+
[
70+
"client",
71+
"rule_type",
72+
"asset",
73+
"threshold_compare_operator",
74+
"threshold_value",
75+
"alert_priority",
76+
],
77+
[
78+
client,
79+
rule_type,
80+
asset,
81+
threshold_compare_operator,
82+
threshold_value,
83+
alert_priority,
84+
],
85+
)
86+
87+
attributes = alpha_DQRule.Attributes.creator(
88+
client=client,
89+
rule_type=rule_type,
90+
asset=asset,
91+
threshold_compare_operator=threshold_compare_operator,
92+
threshold_value=threshold_value,
93+
alert_priority=alert_priority,
94+
rule_name=None,
95+
column=None,
96+
threshold_unit=None,
97+
dimension=None,
98+
custom_sql=None,
99+
description=None,
100+
)
101+
return cls(attributes=attributes)
102+
103+
@classmethod
104+
@init_guid
105+
def column_level_rule_creator(
106+
cls,
107+
*,
108+
client: AtlanClient,
109+
rule_type: str,
110+
asset: Asset,
111+
column: Asset,
112+
threshold_value: int,
113+
alert_priority: alpha_DQRuleAlertPriority,
114+
threshold_compare_operator: Optional[
115+
alpha_DQRuleThresholdCompareOperator
116+
] = alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL,
117+
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
118+
) -> alpha_DQRule:
119+
validate_required_fields(
120+
[
121+
"client",
122+
"rule_type",
123+
"asset",
124+
"column",
125+
"threshold_compare_operator",
126+
"threshold_value",
127+
"alert_priority",
128+
],
129+
[
130+
client,
131+
rule_type,
132+
asset,
133+
column,
134+
threshold_compare_operator,
135+
threshold_value,
136+
alert_priority,
137+
],
138+
)
139+
140+
attributes = alpha_DQRule.Attributes.creator(
141+
client=client,
142+
rule_type=rule_type,
143+
asset=asset,
144+
column=column,
145+
threshold_compare_operator=threshold_compare_operator,
146+
threshold_value=threshold_value,
147+
alert_priority=alert_priority,
148+
threshold_unit=threshold_unit,
149+
rule_name=None,
150+
dimension=None,
151+
custom_sql=None,
152+
description=None,
153+
)
154+
return cls(attributes=attributes)
155+
156+
@classmethod
157+
@init_guid
158+
def updater(
159+
cls: type[SelfAsset],
160+
client: AtlanClient,
161+
qualified_name: str,
162+
threshold_compare_operator: Optional[
163+
alpha_DQRuleThresholdCompareOperator
164+
] = None,
165+
threshold_value: Optional[int] = None,
166+
alert_priority: Optional[alpha_DQRuleAlertPriority] = None,
167+
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
168+
dimension: Optional[alpha_DQDimension] = None,
169+
custom_sql: Optional[str] = None,
170+
rule_name: Optional[str] = None,
171+
description: Optional[str] = None,
172+
) -> SelfAsset:
173+
from pyatlan.model.fluent_search import FluentSearch
174+
175+
validate_required_fields(
176+
["client", "qualified_name"],
177+
[client, qualified_name],
178+
)
179+
request = (
180+
FluentSearch()
181+
.where(alpha_DQRule.QUALIFIED_NAME.eq(qualified_name))
182+
.include_on_results(alpha_DQRule.NAME)
183+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_TEMPLATE_NAME)
184+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_TEMPLATE)
185+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_BASE_DATASET)
186+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_BASE_COLUMN)
187+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_ALERT_PRIORITY)
188+
.include_on_results(alpha_DQRule.DISPLAY_NAME)
189+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_CUSTOM_SQL)
190+
.include_on_results(alpha_DQRule.USER_DESCRIPTION)
191+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_DIMENSION)
192+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_CONFIG_ARGUMENTS)
193+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_SOURCE_SYNC_STATUS)
194+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_STATUS)
195+
).to_request()
196+
197+
results = client.asset.search(request)
198+
199+
if results.count != 1:
200+
raise ValueError(
201+
f"Expected exactly 1 asset for qualified_name: {qualified_name}, "
202+
f"but found: {results.count}"
203+
)
204+
search_result = results.current_page()[0]
205+
206+
retrieved_custom_sql = search_result.alpha_dq_rule_custom_s_q_l # type: ignore[attr-defined]
207+
retrieved_rule_name = search_result.display_name
208+
retrieved_dimension = search_result.alpha_dq_rule_dimension # type: ignore[attr-defined]
209+
retrieved_column = search_result.alpha_dq_rule_base_column # type: ignore[attr-defined]
210+
retrieved_alert_priority = search_result.alpha_dq_rule_alert_priority # type: ignore[attr-defined]
211+
retrieved_description = search_result.user_description
212+
retrieved_asset = search_result.alpha_dq_rule_base_dataset # type: ignore[attr-defined]
213+
retrieved_template_rule_name = search_result.alpha_dq_rule_template_name # type: ignore[attr-defined]
214+
retrieved_template = search_result.alpha_dq_rule_template # type: ignore[attr-defined]
215+
retrieved_threshold_compare_operator = (
216+
search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_compare_operator # type: ignore[attr-defined]
217+
if search_result.alpha_dq_rule_config_arguments is not None # type: ignore[attr-defined]
218+
and search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object # type: ignore[attr-defined]
219+
is not None
220+
else None
221+
)
222+
retrieved_threshold_value = (
223+
search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_value # type: ignore[attr-defined]
224+
if search_result.alpha_dq_rule_config_arguments is not None # type: ignore[attr-defined]
225+
and search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object # type: ignore[attr-defined]
226+
is not None
227+
else None
228+
) # type: ignore[attr-defined]
229+
retrieved_threshold_unit = (
230+
search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_unit # type: ignore[attr-defined]
231+
if search_result.alpha_dq_rule_config_arguments is not None # type: ignore[attr-defined]
232+
and search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object # type: ignore[attr-defined]
233+
is not None
234+
else None
235+
) # type: ignore[attr-defined]
236+
237+
config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw(
238+
is_alert_enabled=True,
239+
custom_sql=custom_sql or retrieved_custom_sql,
240+
display_name=rule_name or retrieved_rule_name,
241+
dimension=dimension or retrieved_dimension,
242+
compare_operator=threshold_compare_operator
243+
or retrieved_threshold_compare_operator,
244+
threshold_value=threshold_value or retrieved_threshold_value,
245+
threshold_unit=threshold_unit or retrieved_threshold_unit,
246+
column=retrieved_column,
247+
dq_priority=alert_priority or retrieved_alert_priority,
248+
description=description or retrieved_description,
249+
)
250+
251+
attr_dq = cls.Attributes(
252+
name="",
253+
alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments(
254+
alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject(
255+
alpha_dq_rule_threshold_compare_operator=threshold_compare_operator
256+
or retrieved_threshold_compare_operator,
257+
alpha_dq_rule_threshold_value=threshold_value
258+
or retrieved_threshold_value,
259+
alpha_dq_rule_threshold_unit=threshold_unit
260+
or retrieved_threshold_unit,
261+
),
262+
alpha_dq_rule_config_arguments_raw=config_arguments_raw,
263+
),
264+
alpha_dq_rule_base_dataset_qualified_name=retrieved_asset.qualified_name,
265+
alpha_dq_rule_alert_priority=alert_priority or retrieved_alert_priority,
266+
alpha_dq_rule_base_dataset=retrieved_asset,
267+
qualified_name=qualified_name,
268+
alpha_dq_rule_dimension=dimension or retrieved_dimension,
269+
alpha_dq_rule_template_name=retrieved_template_rule_name,
270+
alpha_dq_rule_template=alpha_DQRuleTemplate.ref_by_qualified_name(
271+
qualified_name=retrieved_template.qualified_name
272+
),
273+
)
274+
275+
if retrieved_column is not None:
276+
attr_dq.alpha_dq_rule_base_column_qualified_name = (
277+
retrieved_column.qualified_name
278+
)
279+
attr_dq.alpha_dq_rule_base_column = retrieved_column # type: ignore
280+
281+
custom_sql = custom_sql or retrieved_custom_sql
282+
if custom_sql is not None:
283+
attr_dq.alpha_dq_rule_custom_s_q_l = custom_sql
284+
attr_dq.display_name = rule_name or retrieved_rule_name
285+
if description is not None:
286+
attr_dq.user_description = description or retrieved_description
287+
288+
return cls(attributes=attr_dq)

0 commit comments

Comments
 (0)