Skip to content

Commit 362da79

Browse files
vaibhavatlanvaibhavatlan
committed
main logic
1 parent 0e1f294 commit 362da79

2 files changed

Lines changed: 262 additions & 2 deletions

File tree

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

Lines changed: 257 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from __future__ import annotations
66

77
from datetime import datetime
8-
from typing import TYPE_CHECKING, ClassVar, List, Optional, Set
8+
from typing import TYPE_CHECKING, ClassVar, List, Optional, Set, overload
99

1010
from pydantic.v1 import Field, validator
1111

@@ -15,6 +15,8 @@
1515
alpha_DQRuleAlertPriority,
1616
alpha_DQRuleStatus,
1717
alpha_DQSourceSyncStatus,
18+
alpha_DQRuleThresholdCompareOperator,
19+
alpha_DQRuleThresholdUnit,
1820
)
1921
from pyatlan.model.fields.atlan_fields import (
2022
KeywordField,
@@ -23,16 +25,125 @@
2325
RelationField,
2426
TextField,
2527
)
26-
from pyatlan.model.structs import alpha_DQRuleConfigArguments
28+
from pyatlan.model.structs import alpha_DQRuleConfigArguments, alpha_DQRuleThresholdObject
29+
from pyatlan.utils import init_guid, validate_required_fields
30+
# from pyatlan.client.atlan import AtlanClient
31+
import json
32+
import uuid
33+
import time
2734

2835
from .data_quality import DataQuality
36+
# from fluent_search import FluentSearch
2937

3038
if TYPE_CHECKING:
3139
from pyatlan.model.assets import Column
3240

3341

3442
class alpha_DQRule(DataQuality):
3543
"""Description"""
44+
@classmethod
45+
@init_guid
46+
def custom_sql_creator(
47+
cls,
48+
*,
49+
# client: AtlanClient,
50+
rule_name: str,
51+
asset: Asset,
52+
custom_sql: str,
53+
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
54+
threshold_value: float,
55+
alert_priority: alpha_DQRuleAlertPriority,
56+
dimension: alpha_DQDimension,
57+
description: Optional[str] = None,
58+
) -> alpha_DQRule:
59+
validate_required_fields(
60+
["rule_name", "asset", "threshold_compare_operator", "threshold_value", "alert_priority", "dimension", "custom_sql"],
61+
[rule_name, asset, threshold_compare_operator, threshold_value, alert_priority,dimension, custom_sql],
62+
)
63+
64+
attributes = alpha_DQRule.Attributes.create(
65+
# client=client,
66+
rule_name=rule_name,
67+
rule_type="Custom SQL",
68+
asset=asset,
69+
threshold_compare_operator=threshold_compare_operator,
70+
threshold_value=threshold_value,
71+
alert_priority=alert_priority,
72+
dimension=dimension,
73+
custom_sql=custom_sql,
74+
description=description,
75+
column_qualified_name=None,
76+
threshold_unit=None,
77+
)
78+
return cls(attributes=attributes)
79+
80+
@classmethod
81+
@init_guid
82+
def table_level_rule_creator(
83+
cls,
84+
*,
85+
# client:AtlanClient,
86+
rule_type: str,
87+
asset: Asset,
88+
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
89+
threshold_value: float,
90+
alert_priority: alpha_DQRuleAlertPriority,
91+
) -> alpha_DQRule:
92+
validate_required_fields(
93+
["rule_type", "asset", "threshold_compare_operator", "threshold_value", "alert_priority"],
94+
[rule_type, asset, threshold_compare_operator, threshold_value, alert_priority],
95+
)
96+
97+
attributes = alpha_DQRule.Attributes.create(
98+
client=client,
99+
rule_type=rule_type,
100+
asset=asset,
101+
threshold_compare_operator=threshold_compare_operator,
102+
threshold_value=threshold_value,
103+
alert_priority=alert_priority,
104+
rule_name=None,
105+
column_qualified_name=None,
106+
threshold_unit=None,
107+
dimension=None,
108+
custom_sql=None,
109+
description=None
110+
)
111+
return cls(attributes=attributes)
112+
113+
@classmethod
114+
@init_guid
115+
def column_level_rule_creator(
116+
cls,
117+
*,
118+
# client:AtlanClient,
119+
rule_type: str,
120+
asset: Asset,
121+
column_qualified_name: str,
122+
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
123+
threshold_value: float,
124+
alert_priority: alpha_DQRuleAlertPriority,
125+
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
126+
) -> alpha_DQRule:
127+
validate_required_fields(
128+
["rule_type", "asset", "column_qualified_name", "threshold_compare_operator", "threshold_value", "alert_priority"],
129+
[rule_type, asset, column_qualified_name, threshold_compare_operator, threshold_value, alert_priority],
130+
)
131+
132+
attributes = alpha_DQRule.Attributes.create(
133+
client=client,
134+
rule_type=rule_type,
135+
asset=asset,
136+
column_qualified_name=column_qualified_name,
137+
threshold_compare_operator=threshold_compare_operator,
138+
threshold_value=threshold_value,
139+
alert_priority=alert_priority,
140+
threshold_unit=threshold_unit,
141+
rule_name=None,
142+
dimension=None,
143+
custom_sql=None,
144+
description=None
145+
)
146+
return cls(attributes=attributes)
36147

37148
type_name: str = Field(default="alpha_DQRule", allow_mutation=False)
38149

@@ -700,6 +811,149 @@ class Attributes(DataQuality.Attributes):
700811
alpha_dq_rule_reference_datasets: Optional[List[Asset]] = Field(
701812
default=None, description=""
702813
) # relationship
814+
815+
@staticmethod
816+
def _generate_config_arguments_raw(
817+
*,
818+
is_alert_enabled: bool = True,
819+
custom_sql: Optional[str] = None,
820+
display_name: Optional[str] = None,
821+
dimension: Optional[alpha_DQDimension] = None,
822+
compare_operator: alpha_DQRuleThresholdCompareOperator,
823+
threshold_value: float,
824+
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
825+
column_qualified_name: Optional[str] = None,
826+
dq_priority: alpha_DQRuleAlertPriority,
827+
description: Optional[str] = None
828+
) -> str:
829+
830+
config = {
831+
"isAlertEnabled": is_alert_enabled,
832+
"alpha_dqRuleTemplateConfigThresholdObject": {
833+
"alpha_dqRuleTemplateConfigThresholdCompareOperator": compare_operator,
834+
"alpha_dqRuleTemplateConfigThresholdValue": threshold_value,
835+
"alpha_dqRuleTemplateConfigThresholdUnit": threshold_unit,
836+
},
837+
"alpha_dqRuleTemplateAdvancedSettings.dqPriority": dq_priority,
838+
}
839+
840+
if column_qualified_name is not None:
841+
config["alpha_dqRuleTemplateConfigBaseColumnQualifiedName"] = column_qualified_name
842+
843+
if description is not None:
844+
config["alpha_dqRuleTemplateConfigUserDescription"] = description
845+
846+
if custom_sql is not None:
847+
config["alpha_dqRuleTemplateConfigCustomSQL"] = custom_sql
848+
849+
if display_name is not None:
850+
config["alpha_dqRuleTemplateConfigDisplayName"] = display_name
851+
852+
if dimension is not None:
853+
config["alpha_dqRuleTemplateConfigDimension"] = dimension
854+
855+
return json.dumps(config)
856+
857+
@staticmethod
858+
def _generate_uuid():
859+
d = int(time.time() * 1000)
860+
random_bytes = uuid.uuid4().bytes
861+
rand_index = 0
862+
863+
def replace_char(c):
864+
nonlocal d, rand_index
865+
r = (d + random_bytes[rand_index % 16]) % 16
866+
rand_index += 1
867+
d = d // 16
868+
if c == 'x':
869+
return hex(r)[2:]
870+
elif c == 'y':
871+
return hex((r & 0x3) | 0x8)[2:] # y -> 8 to b
872+
else:
873+
return c
874+
875+
template = 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'
876+
uuid_str = ''.join(replace_char(c) if c in 'xy' else c for c in template)
877+
return uuid_str
878+
879+
@classmethod
880+
@init_guid
881+
def create(
882+
cls,
883+
*,
884+
# client: AtlanClient,
885+
rule_name:str,
886+
rule_type: str,
887+
asset: Asset,
888+
threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
889+
threshold_value: float,
890+
alert_priority: alpha_DQRuleAlertPriority,
891+
column_qualified_name: Optional[str] = None,
892+
threshold_unit: Optional[alpha_DQRuleThresholdUnit] = None,
893+
dimension: Optional[alpha_DQDimension] = None,
894+
custom_sql: Optional[str] = None,
895+
description: Optional[str] = None,
896+
) -> alpha_DQRule.Attributes:
897+
# request = (
898+
# FluentSearch()
899+
# .where(Asset.TYPE_NAME.eq("alpha_DQRuleTemplate"))
900+
# .include_on_results(Asset.NAME)
901+
# .include_on_results(Asset.DISPLAY_NAME)
902+
# .include_on_results(Asset.QUALIFIED_NAME)
903+
# ).to_request() #
904+
# for result in client.asset.search(request):
905+
# if result.display_name == rule_type:
906+
# template_rule_name = result.name
907+
# template_qualified_name = result.qualified_name
908+
# if dimension is None:
909+
# dimension = result.alpha_dq_rule_dimension
910+
911+
config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw(
912+
is_alert_enabled=True,
913+
custom_sql=custom_sql,
914+
display_name=rule_name,
915+
dimension=dimension,
916+
compare_operator=threshold_compare_operator,
917+
threshold_value=threshold_value,
918+
threshold_unit=threshold_unit,
919+
column_qualified_name = column_qualified_name,
920+
dq_priority=alert_priority,
921+
description=description,
922+
)
923+
924+
925+
attr_dq = alpha_DQRule.Attributes(
926+
name="",
927+
alpha_dq_rule_config_arguments=alpha_DQRuleConfigArguments(
928+
alpha_dq_rule_threshold_object=alpha_DQRuleThresholdObject(
929+
alpha_dq_rule_threshold_compare_operator=threshold_compare_operator,
930+
alpha_dq_rule_threshold_value=threshold_value,
931+
alpha_dq_rule_threshold_unit=threshold_unit
932+
),
933+
alpha_dq_rule_config_arguments_raw=config_arguments_raw,
934+
),
935+
alpha_dq_rule_base_dataset_qualified_name=asset.qualified_name,
936+
alpha_dq_rule_alert_priority=alert_priority,
937+
alpha_dq_rule_source_sync_status= alpha_DQSourceSyncStatus.IN_PROGRESS,
938+
alpha_dq_rule_status=alpha_DQRuleStatus.ACTIVE,
939+
alpha_dq_rule_base_dataset= asset,
940+
qualified_name=f"{asset.qualified_name}/rule/{str(cls._generate_uuid())}",
941+
alpha_dq_rule_dimension=dimension,
942+
# alpha_dq_rule_template_name=template_rule_name,
943+
# alpha_dq_rule_template=alpha_DQRuleTemplate.ref_by_qualified_name(qualified_name=template_qualified_name),
944+
)
945+
946+
if column_qualified_name is not None:
947+
attr_dq.alpha_dq_rule_base_column_qualified_name=column_qualified_name
948+
attr_dq.alpha_dq_rule_base_column=Column.ref_by_qualified_name(qualified_name=column_qualified_name)
949+
950+
if rule_type == "Custom SQL":
951+
attr_dq.alpha_dq_rule_custom_s_q_l = custom_sql
952+
attr_dq.display_name = rule_name
953+
if description is not None:
954+
attr_dq.user_description = description
955+
956+
return attr_dq
703957

704958
attributes: alpha_DQRule.Attributes = Field(
705959
default_factory=lambda: alpha_DQRule.Attributes(),
@@ -713,3 +967,4 @@ class Attributes(DataQuality.Attributes):
713967

714968
from .alpha__d_q_rule_template import alpha_DQRuleTemplate # noqa: E402, F401
715969
from .asset import Asset # noqa: E402, F401
970+
# from .column import co

‎pyatlan/model/enums.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2974,3 +2974,8 @@ class alpha_DQSourceSyncStatus(str, Enum):
29742974
FAILURE = "FAILURE"
29752975
IN_PROGRESS = "IN_PROGRESS"
29762976
WAITING_FOR_SCHEDULE = "WAITING_FOR_SCHEDULE"
2977+
2978+
class alpha_DQRuleThresholdCompareOperator(str, Enum):
2979+
EQUAL = "EQ"
2980+
GREATER_THAN_EQUAL = "GTE"
2981+
LESS_THAN_EQUAL = "LTE"

0 commit comments

Comments
 (0)