Skip to content

Commit 9c4ea2f

Browse files
vaibhavatlanvaibhavatlan
committed
Added updater and QA fixes
1 parent c1e5117 commit 9c4ea2f

File tree

1 file changed

+120
-8
lines changed

1 file changed

+120
-8
lines changed

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

Lines changed: 120 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def custom_sql_creator(
5454
asset: Asset,
5555
custom_sql: str,
5656
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
57-
threshold_value: float,
57+
threshold_value: int,
5858
alert_priority: alpha_DQRuleAlertPriority,
5959
dimension: alpha_DQDimension,
6060
description: Optional[str] = None,
@@ -105,7 +105,7 @@ def table_level_rule_creator(
105105
rule_type: str,
106106
asset: Asset,
107107
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
108-
threshold_value: float,
108+
threshold_value: int,
109109
alert_priority: alpha_DQRuleAlertPriority,
110110
) -> alpha_DQRule:
111111
validate_required_fields(
@@ -150,9 +150,11 @@ def column_level_rule_creator(
150150
rule_type: str,
151151
asset: Asset,
152152
column: Asset,
153-
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
154-
threshold_value: float,
153+
threshold_value: int,
155154
alert_priority: alpha_DQRuleAlertPriority,
155+
threshold_compare_operator: Optional[
156+
alpha_DQRuleThresholdCompareOperator
157+
] = alpha_DQRuleThresholdCompareOperator.LESS_THAN_EQUAL,
156158
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
157159
) -> alpha_DQRule:
158160
validate_required_fields(
@@ -190,6 +192,116 @@ def column_level_rule_creator(
190192
)
191193
return cls(attributes=attributes)
192194

195+
@classmethod
196+
@init_guid
197+
def updater(
198+
cls: type[SelfAsset],
199+
client: AtlanClient,
200+
qualified_name: str,
201+
threshold_compare_operator: Optional[
202+
alpha_DQRuleThresholdCompareOperator
203+
] = None,
204+
threshold_value: Optional[int] = None,
205+
alert_priority: Optional[alpha_DQRuleAlertPriority] = None,
206+
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
207+
dimension: Optional[alpha_DQDimension] = None,
208+
custom_sql: Optional[str] = None,
209+
rule_name: Optional[str] = None,
210+
description: Optional[str] = None,
211+
) -> SelfAsset:
212+
from pyatlan.model.fluent_search import FluentSearch
213+
214+
validate_required_fields(
215+
["qualified_name"],
216+
[qualified_name],
217+
)
218+
request = (
219+
FluentSearch()
220+
.where(alpha_DQRule.QUALIFIED_NAME.eq(qualified_name))
221+
.include_on_results(alpha_DQRule.NAME)
222+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_TEMPLATE_NAME)
223+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_TEMPLATE)
224+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_BASE_DATASET)
225+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_BASE_COLUMN)
226+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_ALERT_PRIORITY)
227+
.include_on_results(alpha_DQRule.DISPLAY_NAME)
228+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_CUSTOM_SQL)
229+
.include_on_results(alpha_DQRule.USER_DESCRIPTION)
230+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_DIMENSION)
231+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_CONFIG_ARGUMENTS)
232+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_SOURCE_SYNC_STATUS)
233+
.include_on_results(alpha_DQRule.ALPHADQ_RULE_STATUS)
234+
).to_request()
235+
236+
for result in client.asset.search(request):
237+
retrieved_custom_sql = result.alpha_dq_rule_custom_s_q_l # type: ignore[attr-defined]
238+
retrieved_rule_name = result.display_name
239+
retrieved_dimension = result.alpha_dq_rule_dimension # type: ignore[attr-defined]
240+
retrieved_column = result.alpha_dq_rule_base_column # type: ignore[attr-defined]
241+
retrieved_alert_priority = result.alpha_dq_rule_alert_priority # type: ignore[attr-defined]
242+
retrieved_description = result.user_description
243+
retrieved_asset = result.alpha_dq_rule_base_dataset # type: ignore[attr-defined]
244+
retrieved_template_rule_name = result.alpha_dq_rule_template_name # type: ignore[attr-defined]
245+
retrieved_template = result.alpha_dq_rule_template # type: ignore[attr-defined]
246+
retrieved_threshold_compare_operator = result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_compare_operator # type: ignore[attr-defined]
247+
retrieved_threshold_value = result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_value # type: ignore[attr-defined]
248+
retrieved_threshold_unit = result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_unit # type: ignore[attr-defined]
249+
250+
config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw(
251+
is_alert_enabled=True,
252+
custom_sql=custom_sql or retrieved_custom_sql,
253+
display_name=rule_name or retrieved_rule_name,
254+
dimension=dimension or retrieved_dimension,
255+
compare_operator=threshold_compare_operator
256+
or retrieved_threshold_compare_operator,
257+
threshold_value=threshold_value or retrieved_threshold_value,
258+
threshold_unit=threshold_unit or retrieved_threshold_unit,
259+
column=retrieved_column,
260+
dq_priority=alert_priority or retrieved_alert_priority,
261+
description=description or retrieved_description,
262+
)
263+
264+
attr_dq = cls.Attributes(
265+
name="",
266+
alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments(
267+
alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject(
268+
alpha_dq_rule_threshold_compare_operator=threshold_compare_operator
269+
or retrieved_threshold_compare_operator,
270+
alpha_dq_rule_threshold_value=threshold_value
271+
or retrieved_threshold_value,
272+
alpha_dq_rule_threshold_unit=threshold_unit
273+
or retrieved_threshold_unit,
274+
),
275+
alpha_dq_rule_config_arguments_raw=config_arguments_raw,
276+
),
277+
alpha_dq_rule_base_dataset_qualified_name=retrieved_asset.qualified_name,
278+
alpha_dq_rule_alert_priority=alert_priority or retrieved_alert_priority,
279+
alpha_dq_rule_source_sync_status=alpha_DQSourceSyncStatus.IN_PROGRESS,
280+
alpha_dq_rule_status=alpha_DQRuleStatus.ACTIVE,
281+
alpha_dq_rule_base_dataset=retrieved_asset,
282+
qualified_name=retrieved_asset.qualified_name,
283+
alpha_dq_rule_dimension=dimension or retrieved_dimension,
284+
alpha_dq_rule_template_name=retrieved_template_rule_name,
285+
alpha_dq_rule_template=alpha_DQRuleTemplate.ref_by_qualified_name(
286+
qualified_name=retrieved_template.qualified_name
287+
),
288+
)
289+
290+
if retrieved_column is not None:
291+
attr_dq.alpha_dq_rule_base_column_qualified_name = (
292+
retrieved_column.qualified_name
293+
)
294+
attr_dq.alpha_dq_rule_base_column = retrieved_column # type: ignore
295+
296+
custom_sql = custom_sql or retrieved_custom_sql
297+
if custom_sql is not None:
298+
attr_dq.alpha_dq_rule_custom_s_q_l = custom_sql
299+
attr_dq.display_name = rule_name or retrieved_rule_name
300+
if description is not None:
301+
attr_dq.user_description = description or retrieved_description
302+
303+
return cls(attributes=attr_dq)
304+
193305
type_name: str = Field(default="alpha_DQRule", allow_mutation=False)
194306

195307
@validator("type_name")
@@ -865,7 +977,7 @@ def _generate_config_arguments_raw(
865977
display_name: Optional[str] = None,
866978
dimension: Optional[alpha_DQDimension] = None,
867979
compare_operator: alpha_DQRuleThresholdCompareOperator,
868-
threshold_value: float,
980+
threshold_value: int,
869981
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
870982
column: Optional[Asset] = None,
871983
dq_priority: alpha_DQRuleAlertPriority,
@@ -932,7 +1044,7 @@ def creator(
9321044
rule_type: str,
9331045
asset: Asset,
9341046
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
935-
threshold_value: float,
1047+
threshold_value: int,
9361048
alert_priority: alpha_DQRuleAlertPriority,
9371049
column: Optional[Asset] = None,
9381050
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
@@ -1007,7 +1119,7 @@ def creator(
10071119
attr_dq.alpha_dq_rule_base_column_qualified_name = column.qualified_name
10081120
attr_dq.alpha_dq_rule_base_column = column # type: ignore
10091121

1010-
if rule_type == "Custom SQL":
1122+
if custom_sql is not None:
10111123
attr_dq.alpha_dq_rule_custom_s_q_l = custom_sql
10121124
attr_dq.display_name = rule_name
10131125
if description is not None:
@@ -1026,5 +1138,5 @@ def creator(
10261138

10271139

10281140
from .alpha__d_q_rule_template import alpha_DQRuleTemplate # noqa: E402, F401
1029-
from .asset import Asset # noqa: E402, F401
1141+
from .asset import Asset, SelfAsset # noqa: E402, F401
10301142
# from .column import co

0 commit comments

Comments
 (0)