Fixed based on the suggestions

Author: vaibhavatlan

pyatlan/cache/template_config_cache.py

@@ -12,7 +12,7 @@
     from pyatlan.client.atlan import AtlanClient
 
 
-class TemplateConfigCache:
+class DQTemplateConfigCache:
     """
     Lazily-loaded cache for DQ rule template configurations to avoid multiple API calls.
     """
@@ -60,7 +60,8 @@ def _refresh_cache(self) -> None:
                     )
                 ).to_request()
 
-                for result in self.client.asset.search(request):
+                results = self.client.asset.search(request)
+                for result in results:
                     template_config = {
                         "name": result.name,
                         "qualified_name": result.qualified_name,

pyatlan/client/atlan.py

@@ -37,7 +37,7 @@
 from pyatlan.cache.group_cache import GroupCache
 from pyatlan.cache.role_cache import RoleCache
 from pyatlan.cache.source_tag_cache import SourceTagCache
-from pyatlan.cache.template_config_cache import TemplateConfigCache
+from pyatlan.cache.template_config_cache import DQTemplateConfigCache
 from pyatlan.cache.user_cache import UserCache
 from pyatlan.client.admin import AdminClient
 from pyatlan.client.asset import A, AssetClient, IndexSearchResults, LineageListResults
@@ -175,7 +175,9 @@ class AtlanClient(BaseSettings):
     _custom_metadata_cache: Optional[CustomMetadataCache] = PrivateAttr(default=None)
     _connection_cache: Optional[ConnectionCache] = PrivateAttr(default=None)
     _source_tag_cache: Optional[SourceTagCache] = PrivateAttr(default=None)
-    _template_config_cache: Optional[TemplateConfigCache] = PrivateAttr(default=None)
+    _dq_template_config_cache: Optional[DQTemplateConfigCache] = PrivateAttr(
+        default=None
+    )
 
     class Config:
         env_prefix = "atlan_"
@@ -350,10 +352,10 @@ def source_tag_cache(self) -> SourceTagCache:
         return self._source_tag_cache
 
     @property
-    def template_config_cache(self) -> TemplateConfigCache:
-        if self._template_config_cache is None:
-            self._template_config_cache = TemplateConfigCache(client=self)
-        return self._template_config_cache
+    def dq_template_config_cache(self) -> DQTemplateConfigCache:
+        if self._dq_template_config_cache is None:
+            self._dq_template_config_cache = DQTemplateConfigCache(client=self)
+        return self._dq_template_config_cache
 
     @classmethod
     def from_token_guid(cls, guid: str) -> AtlanClient:

pyatlan/model/assets/core/alpha__d_q_rule.py

@@ -61,6 +61,7 @@ def custom_sql_creator(
     ) -> alpha_DQRule:
         validate_required_fields(
             [
+                "client",
                 "rule_name",
                 "asset",
                 "threshold_compare_operator",
@@ -70,6 +71,7 @@ def custom_sql_creator(
                 "custom_sql",
             ],
             [
+                client,
                 rule_name,
                 asset,
                 threshold_compare_operator,
@@ -110,13 +112,15 @@ def table_level_rule_creator(
     ) -> alpha_DQRule:
         validate_required_fields(
             [
+                "client",
                 "rule_type",
                 "asset",
                 "threshold_compare_operator",
                 "threshold_value",
                 "alert_priority",
             ],
             [
+                client,
                 rule_type,
                 asset,
                 threshold_compare_operator,
@@ -159,6 +163,7 @@ def column_level_rule_creator(
     ) -> alpha_DQRule:
         validate_required_fields(
             [
+                "client",
                 "rule_type",
                 "asset",
                 "column",
@@ -167,6 +172,7 @@ def column_level_rule_creator(
                 "alert_priority",
             ],
             [
+                client,
                 rule_type,
                 asset,
                 column,
@@ -212,8 +218,8 @@ def updater(
         from pyatlan.model.fluent_search import FluentSearch
 
         validate_required_fields(
-            ["qualified_name"],
-            [qualified_name],
+            ["client", "qualified_name"],
+            [client, qualified_name],
         )
         request = (
             FluentSearch()
@@ -235,6 +241,11 @@ def updater(
 
         results = client.asset.search(request)
 
+        if results.count != 1:
+            raise ValueError(
+                f"Expected exactly 1 asset for qualified_name: {qualified_name}, "
+                f"but found: {results.count}"
+            )
         search_result = results.current_page()[0]
 
         retrieved_custom_sql = search_result.alpha_dq_rule_custom_s_q_l  # type: ignore[attr-defined]
@@ -246,9 +257,27 @@ def updater(
         retrieved_asset = search_result.alpha_dq_rule_base_dataset  # type: ignore[attr-defined]
         retrieved_template_rule_name = search_result.alpha_dq_rule_template_name  # type: ignore[attr-defined]
         retrieved_template = search_result.alpha_dq_rule_template  # type: ignore[attr-defined]
-        retrieved_threshold_compare_operator = search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_compare_operator  # type: ignore[attr-defined]
-        retrieved_threshold_value = search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_value  # type: ignore[attr-defined]
-        retrieved_threshold_unit = search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_unit  # type: ignore[attr-defined]
+        retrieved_threshold_compare_operator = (
+            search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_compare_operator  # type: ignore[attr-defined]
+            if search_result.alpha_dq_rule_config_arguments is not None  # type: ignore[attr-defined]
+            and search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object  # type: ignore[attr-defined]
+            is not None
+            else None
+        )
+        retrieved_threshold_value = (
+            search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_value  # type: ignore[attr-defined]
+            if search_result.alpha_dq_rule_config_arguments is not None  # type: ignore[attr-defined]
+            and search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object  # type: ignore[attr-defined]
+            is not None
+            else None
+        )  # type: ignore[attr-defined]
+        retrieved_threshold_unit = (
+            search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object.alpha_dq_rule_threshold_unit  # type: ignore[attr-defined]
+            if search_result.alpha_dq_rule_config_arguments is not None  # type: ignore[attr-defined]
+            and search_result.alpha_dq_rule_config_arguments.alpha_dq_rule_threshold_object  # type: ignore[attr-defined]
+            is not None
+            else None
+        )  # type: ignore[attr-defined]
 
         config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw(
             is_alert_enabled=True,
@@ -1053,30 +1082,32 @@ def creator(
             custom_sql: Optional[str] = None,
             description: Optional[str] = None,
         ) -> alpha_DQRule.Attributes:
-            template_config = client.template_config_cache.get_template_config(
+            template_config = client.dq_template_config_cache.get_template_config(
                 rule_type
             )
 
             if template_config is None:
                 raise ErrorCode.DQ_RULE_NOT_FOUND.exception_with_parameters(rule_type)
 
-            template_rule_name = template_config["name"]
-            template_qualified_name = template_config["qualified_name"]
+            template_rule_name = template_config.get("name")
+            template_qualified_name = template_config.get("qualified_name")
 
             if dimension is None:
-                dimension = template_config["dimension"]
+                dimension = template_config.get("dimension")
 
             if threshold_unit is None:
-                threashold_object = template_config[
-                    "config"
-                ].alpha_dq_rule_template_config_threshold_object
-                threashold_object_json = json.loads(threashold_object)
-                properties = threashold_object_json.get("properties", {})
-                threshold_unit_field = properties.get(
-                    "alpha_dqRuleTemplateConfigThresholdUnit", {}
-                )
-                default_value = threshold_unit_field.get("default")
-                threshold_unit = default_value
+                config = template_config.get("config")
+                if config is not None:
+                    threashold_object = (
+                        config.alpha_dq_rule_template_config_threshold_object
+                    )
+                    threashold_object_json = json.loads(threashold_object)
+                    properties = threashold_object_json.get("properties", {})
+                    threshold_unit_field = properties.get(
+                        "alpha_dqRuleTemplateConfigThresholdUnit", {}
+                    )
+                    default_value = threshold_unit_field.get("default")
+                    threshold_unit = default_value
 
             config_arguments_raw = (
                 alpha_DQRule.Attributes._generate_config_arguments_raw(
@@ -1112,7 +1143,7 @@ def creator(
                 alpha_dq_rule_dimension=dimension,
                 alpha_dq_rule_template_name=template_rule_name,
                 alpha_dq_rule_template=alpha_DQRuleTemplate.ref_by_qualified_name(
-                    qualified_name=template_qualified_name
+                    qualified_name=template_qualified_name  # type: ignore
                 ),
             )
 

tests/unit/model/alpha__d_q_rule_test.py

@@ -26,7 +26,7 @@
 @pytest.fixture
 def mock_client():
     client = Mock()
-    client.template_config_cache = Mock()
+    client.dq_template_config_cache = Mock()
 
     # Create a proper config object with a JSON string for threshold_object
     config = Mock()
@@ -40,7 +40,7 @@ def mock_client():
         }
     )
 
-    client.template_config_cache.get_template_config.return_value = {
+    client.dq_template_config_cache.get_template_config.return_value = {
         "name": "Test Template",
         "qualified_name": "test/template/123",
         "dimension": alpha_DQDimension.COMPLETENESS,