Skip to content

Commit 5c32f47

Browse files
UdaySagar-GitUdaySagar-Git
committed
feat: add support for new dq templates
1 parent 967b1c5 commit 5c32f47

3 files changed

Lines changed: 100 additions & 3 deletions

File tree

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ def custom_sql_creator(
5959
threshold_value: int,
6060
alert_priority: DataQualityRuleAlertPriority,
6161
dimension: DataQualityDimension,
62+
custom_sql_return_type: Optional[DataQualityRuleCustomSQLReturnType] = None,
6263
description: Optional[str] = None,
6364
) -> DataQualityRule:
6465
validate_required_fields(
@@ -94,6 +95,7 @@ def custom_sql_creator(
9495
alert_priority=alert_priority,
9596
dimension=dimension,
9697
custom_sql=custom_sql,
98+
custom_sql_return_type=custom_sql_return_type,
9799
description=description,
98100
column=None,
99101
threshold_unit=None,
@@ -111,6 +113,7 @@ def table_level_rule_creator(
111113
threshold_compare_operator: DataQualityRuleThresholdCompareOperator,
112114
threshold_value: int,
113115
alert_priority: DataQualityRuleAlertPriority,
116+
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
114117
) -> DataQualityRule:
115118
validate_required_fields(
116119
[
@@ -140,7 +143,7 @@ def table_level_rule_creator(
140143
alert_priority=alert_priority,
141144
rule_name=None,
142145
column=None,
143-
threshold_unit=None,
146+
threshold_unit=threshold_unit,
144147
dimension=None,
145148
custom_sql=None,
146149
description=None,
@@ -249,6 +252,7 @@ def updater(
249252
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
250253
dimension: Optional[DataQualityDimension] = None,
251254
custom_sql: Optional[str] = None,
255+
custom_sql_return_type: Optional[DataQualityRuleCustomSQLReturnType] = None,
252256
rule_name: Optional[str] = None,
253257
description: Optional[str] = None,
254258
rule_conditions: Optional[str] = None,
@@ -271,6 +275,7 @@ def updater(
271275
.include_on_results(DataQualityRule.DQ_RULE_ALERT_PRIORITY)
272276
.include_on_results(DataQualityRule.DISPLAY_NAME)
273277
.include_on_results(DataQualityRule.DQ_RULE_CUSTOM_SQL)
278+
.include_on_results(DataQualityRule.DQ_RULE_CUSTOM_SQL_RETURN_TYPE)
274279
.include_on_results(DataQualityRule.USER_DESCRIPTION)
275280
.include_on_results(DataQualityRule.DQ_RULE_DIMENSION)
276281
.include_on_results(DataQualityRule.DQ_RULE_CONFIG_ARGUMENTS)
@@ -289,6 +294,9 @@ def updater(
289294
search_result = results.current_page()[0]
290295

291296
retrieved_custom_sql = search_result.dq_rule_custom_s_q_l # type: ignore[attr-defined]
297+
retrieved_custom_sql_return_type = (
298+
search_result.dq_rule_custom_s_q_l_return_type # type: ignore[attr-defined]
299+
)
292300
retrieved_rule_name = search_result.display_name
293301
retrieved_dimension = search_result.dq_rule_dimension # type: ignore[attr-defined]
294302
retrieved_column = search_result.dq_rule_base_column # type: ignore[attr-defined]
@@ -376,6 +384,9 @@ def updater(
376384
if custom_sql is not None:
377385
attr_dq.dq_rule_custom_s_q_l = custom_sql
378386
attr_dq.display_name = rule_name or retrieved_rule_name
387+
attr_dq.dq_rule_custom_s_q_l_return_type = (
388+
custom_sql_return_type or retrieved_custom_sql_return_type
389+
)
379390
if description is not None:
380391
attr_dq.user_description = description or retrieved_description
381392

@@ -1257,6 +1268,7 @@ def creator(
12571268
threshold_unit: Optional[DataQualityRuleThresholdUnit] = None,
12581269
dimension: Optional[DataQualityDimension] = None,
12591270
custom_sql: Optional[str] = None,
1271+
custom_sql_return_type: Optional[DataQualityRuleCustomSQLReturnType] = None,
12601272
description: Optional[str] = None,
12611273
rule_conditions: Optional[str] = None,
12621274
row_scope_filtering_enabled: Optional[bool] = False,
@@ -1316,6 +1328,8 @@ def creator(
13161328
if custom_sql is not None:
13171329
attr_dq.dq_rule_custom_s_q_l = custom_sql
13181330
attr_dq.display_name = rule_name
1331+
if custom_sql_return_type is not None:
1332+
attr_dq.dq_rule_custom_s_q_l_return_type = custom_sql_return_type
13191333
if description is not None:
13201334
attr_dq.user_description = description
13211335

‎pyatlan/model/dq_rule_conditions.py‎

Lines changed: 69 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,17 +22,33 @@ class DQCondition(AtlanObject):
2222
)
2323
min_value: Optional[int] = Field(default=None, description="")
2424
max_value: Optional[int] = Field(default=None, description="")
25+
reference_table: Optional[str] = Field(default=None, description="")
26+
reference_column: Optional[str] = Field(default=None, description="")
27+
target_table: Optional[str] = Field(default=None, description="")
28+
target_column: Optional[str] = Field(default=None, description="")
2529

2630
def __init__(
2731
self,
2832
type: DataQualityRuleTemplateConfigRuleConditions,
2933
value: Optional[Union[str, int, List[str], Dict[str, Any]]] = None,
3034
min_value: Optional[int] = None,
3135
max_value: Optional[int] = None,
36+
reference_table: Optional[str] = None,
37+
reference_column: Optional[str] = None,
38+
target_table: Optional[str] = None,
39+
target_column: Optional[str] = None,
3240
**kwargs,
3341
):
3442
super().__init__(
35-
type=type, value=value, min_value=min_value, max_value=max_value, **kwargs
43+
type=type,
44+
value=value,
45+
min_value=min_value,
46+
max_value=max_value,
47+
reference_table=reference_table,
48+
reference_column=reference_column,
49+
target_table=target_table,
50+
target_column=target_column,
51+
**kwargs,
3652
)
3753

3854
if (
@@ -60,6 +76,23 @@ def __init__(
6076
"min_value, max_value",
6177
"min_value <= max_value",
6278
)
79+
elif self.type == DataQualityRuleTemplateConfigRuleConditions.IN_LIST_REFERENCE:
80+
validate_required_fields(
81+
["reference_table", "reference_column"],
82+
[self.reference_table, self.reference_column],
83+
)
84+
elif self.type == DataQualityRuleTemplateConfigRuleConditions.ROW_COUNT_RECON:
85+
validate_required_fields(["target_table"], [self.target_table])
86+
elif self.type in [
87+
DataQualityRuleTemplateConfigRuleConditions.AVERAGE_RECON,
88+
DataQualityRuleTemplateConfigRuleConditions.SUM_RECON,
89+
DataQualityRuleTemplateConfigRuleConditions.DUPLICATE_COUNT_RECON,
90+
DataQualityRuleTemplateConfigRuleConditions.UNIQUE_COUNT_RECON,
91+
]:
92+
validate_required_fields(
93+
["target_table", "target_column"],
94+
[self.target_table, self.target_column],
95+
)
6396
else:
6497
validate_required_fields(["value"], [self.value])
6598
if self.type in [
@@ -81,6 +114,25 @@ def to_dict(self) -> Dict[str, Any]:
81114
== DataQualityRuleTemplateConfigRuleConditions.STRING_LENGTH_BETWEEN
82115
):
83116
result["value"] = {"minValue": self.min_value, "maxValue": self.max_value}
117+
elif self.type == DataQualityRuleTemplateConfigRuleConditions.IN_LIST_REFERENCE:
118+
result["value"] = {
119+
"reference_table": self.reference_table,
120+
"reference_column": self.reference_column,
121+
}
122+
elif self.type == DataQualityRuleTemplateConfigRuleConditions.ROW_COUNT_RECON:
123+
result["value"] = {
124+
"target_table": self.target_table,
125+
}
126+
elif self.type in [
127+
DataQualityRuleTemplateConfigRuleConditions.AVERAGE_RECON,
128+
DataQualityRuleTemplateConfigRuleConditions.SUM_RECON,
129+
DataQualityRuleTemplateConfigRuleConditions.DUPLICATE_COUNT_RECON,
130+
DataQualityRuleTemplateConfigRuleConditions.UNIQUE_COUNT_RECON,
131+
]:
132+
result["value"] = {
133+
"target_table": self.target_table,
134+
"target_column": self.target_column,
135+
}
84136
else:
85137
result["value"] = {"value": self.value}
86138

@@ -99,6 +151,10 @@ def add_condition(
99151
value: Optional[Union[str, int, List[str]]] = None,
100152
min_value: Optional[int] = None,
101153
max_value: Optional[int] = None,
154+
reference_table: Optional[str] = None,
155+
reference_column: Optional[str] = None,
156+
target_table: Optional[str] = None,
157+
target_column: Optional[str] = None,
102158
) -> DQRuleConditionsBuilder:
103159
"""
104160
Add a condition to the builder.
@@ -107,11 +163,22 @@ def add_condition(
107163
:param value: value of type str, int, or list depending on condition type
108164
:param min_value: minimum value for range-based conditions
109165
:param max_value: maximum value for range-based conditions
166+
:param reference_table: qualified name of the reference table for IN_LIST_REFERENCE condition
167+
:param reference_column: qualified name of the reference column for IN_LIST_REFERENCE condition
168+
:param target_table: qualified name of the target table for reconciliation conditions
169+
:param target_column: qualified name of the target column for reconciliation conditions (except ROW_COUNT_RECON)
110170
:returns: the builder for method chaining
111171
"""
112172
self._conditions.append(
113173
DQCondition(
114-
type=type, value=value, min_value=min_value, max_value=max_value
174+
type=type,
175+
value=value,
176+
min_value=min_value,
177+
max_value=max_value,
178+
reference_table=reference_table,
179+
reference_column=reference_column,
180+
target_table=target_table,
181+
target_column=target_column,
115182
)
116183
)
117184
return self

‎pyatlan/model/enums.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2468,6 +2468,16 @@ class DataQualityRuleTemplateConfigRuleConditions(str, Enum):
24682468
IN_LIST = "IN_LIST"
24692469
NOT_IN_LIST = "NOT_IN_LIST"
24702470

2471+
# Reference list conditions
2472+
IN_LIST_REFERENCE = "IN_LIST_REFERENCE"
2473+
2474+
# Reconciliation conditions
2475+
ROW_COUNT_RECON = "ROW_COUNT_RECON"
2476+
AVERAGE_RECON = "AVERAGE_RECON"
2477+
SUM_RECON = "SUM_RECON"
2478+
DUPLICATE_COUNT_RECON = "DUPLICATE_COUNT_RECON"
2479+
UNIQUE_COUNT_RECON = "UNIQUE_COUNT_RECON"
2480+
24712481

24722482
# **************************************
24732483
# CODE BELOW IS GENERATED NOT MODIFY **
@@ -2738,6 +2748,11 @@ class DataQualityRuleAlertPriority(str, Enum):
27382748
LOW = "LOW"
27392749

27402750

2751+
class DataQualityRuleCustomSQLResultType(str, Enum):
2752+
ROW_COUNT = "ROW_COUNT"
2753+
NUMERIC_VALUE = "NUMERIC_VALUE"
2754+
2755+
27412756
class DataQualityRuleCustomSQLReturnType(str, Enum):
27422757
ROW_COUNT = "ROW_COUNT"
27432758
NUMERIC_VALUE = "NUMERIC_VALUE"
@@ -2763,6 +2778,7 @@ class DataQualityRuleThresholdUnit(str, Enum):
27632778
WEEKS = "WEEKS"
27642779
MONTHS = "MONTHS"
27652780
YEARS = "YEARS"
2781+
ABSOLUTE = "ABSOLUTE"
27662782

27672783

27682784
class DataQualityScheduleType(str, Enum):

0 commit comments

Comments
 (0)