APP-8059 : Added support for DQ rules schedule (#692)

* added dq rules schedule support

* added unit test

* added doc string and method name change

* made the suggested changes

* fixed the doc string

Author: vaibhavatlan

pyatlan/client/asset.py

@@ -84,6 +84,7 @@
     EntityStatus,
     SaveSemantic,
     SortOrder,
+    alpha_DQScheduleType,
 )
 from pyatlan.model.fields.atlan_fields import AtlanField
 from pyatlan.model.lineage import LineageDirection, LineageListRequest
@@ -2029,6 +2030,35 @@ def process_assets(
                     func(asset)
         return len(guids_processed)
 
+    @validate_arguments
+    def add_dq_rule_schedule(
+        self,
+        asset_type: Type[A],
+        asset_name: str,
+        asset_qualified_name: str,
+        schedule_crontab: str,
+        schedule_time_zone: str,
+    ) -> AssetMutationResponse:
+        """
+        Add a data quality rule schedule to an asset.
+
+        :param asset_type: the type of asset to update (e.g., Table)
+        :param asset_name: the name of the asset to update
+        :param asset_qualified_name: the qualified name of the asset to update
+        :param schedule_crontab: cron expression string defining the schedule for the DQ rules, e.g: `5 4 * * *`.
+        :param schedule_time_zone: timezone for the schedule, e.g: `Europe/Paris`.
+        :returns: the result of the save
+        :raises AtlanError: on any API communication issue
+        """
+        updated_asset = asset_type.updater(
+            qualified_name=asset_qualified_name, name=asset_name
+        )
+        updated_asset.alpha_asset_d_q_schedule_time_zone = schedule_time_zone
+        updated_asset.alpha_asset_d_q_schedule_crontab = schedule_crontab
+        updated_asset.alpha_asset_d_q_schedule_type = alpha_DQScheduleType.CRON
+        response = self.save(updated_asset)
+        return response
+
 
 class SearchResults(ABC, Iterable):
     """

pyatlan/model/utils.py

@@ -25,7 +25,7 @@ def encoders():
     }
 
 
-def convert_with_fixed_prefix(input_str, fixed_prefix="alpha_dq"):
+def convert_with_fixed_prefix(input_str, fixed_prefix):
     prefix = fixed_prefix
     remaining = input_str[len(prefix) + 1 :]
     parts = remaining.split("_")
@@ -39,7 +39,9 @@ def to_camel_case(value: str) -> str:
     if value == "__root__":
         return value
     if value.startswith("alpha_dq"):
-        return convert_with_fixed_prefix(value)
+        return convert_with_fixed_prefix(input_str=value, fixed_prefix="alpha_dq")
+    if value.startswith("alpha_asset"):
+        return convert_with_fixed_prefix(input_str=value, fixed_prefix="alpha_asset")
     if value in CAMEL_CASE_OVERRIDES:
         return CAMEL_CASE_OVERRIDES[value]
     value = "".join(word.capitalize() for word in value.split("_"))

tests/unit/test_client.py

@@ -48,6 +48,7 @@
     LineageDirection,
     SaveSemantic,
     SortOrder,
+    alpha_DQScheduleType,
 )
 from pyatlan.model.fluent_search import CompoundQuery, FluentSearch
 from pyatlan.model.group import GroupRequest
@@ -2716,3 +2717,50 @@ def test_get_by_qualified_name_asset_not_found(mock_api_caller):
             )
 
         mock_execute.assert_called_once()
+
+
+def test_add_dq_rule_schedule(mock_api_caller):
+    asset_client = AssetClient(mock_api_caller)
+    asset_type = Table
+    asset_name = "Test Table"
+    asset_qualified_name = "test/qualified/name"
+    schedule_cron_string = "0 0 * * *"
+    schedule_time_zone = "UTC"
+
+    updated_table = Table()
+    updated_table.guid = "test-guid-123"
+    updated_table.alpha_asset_d_q_schedule_time_zone = schedule_time_zone
+    updated_table.alpha_asset_d_q_schedule_crontab = schedule_cron_string
+    updated_table.alpha_asset_d_q_schedule_type = alpha_DQScheduleType.CRON
+
+    mock_response = Mock(spec=AssetMutationResponse)
+
+    with patch.object(
+        asset_type, "updater", return_value=updated_table
+    ) as mock_updater:
+        with patch.object(
+            asset_client, "save", return_value=mock_response
+        ) as mock_save:
+            result = asset_client.add_dq_rule_schedule(
+                asset_type=asset_type,
+                asset_name=asset_name,
+                asset_qualified_name=asset_qualified_name,
+                schedule_crontab=schedule_cron_string,
+                schedule_time_zone=schedule_time_zone,
+            )
+
+            mock_updater.assert_called_once_with(
+                qualified_name=asset_qualified_name,
+                name=asset_name,
+            )
+            assert (
+                updated_table.alpha_asset_d_q_schedule_time_zone == schedule_time_zone
+            )
+            assert (
+                updated_table.alpha_asset_d_q_schedule_crontab == schedule_cron_string
+            )
+            assert (
+                updated_table.alpha_asset_d_q_schedule_type == alpha_DQScheduleType.CRON
+            )
+            mock_save.assert_called_once_with(updated_table)
+            assert result == mock_response