Creator coming to closure

Author: vaibhavatlan

pyatlan/errors.py

@@ -932,6 +932,13 @@ class ErrorCode(Enum):
         "Verify the API token provided is a valid username for that token.",
         NotFoundError,
     )
+    DQ_RULE_NOT_FOUND = (
+        404,
+        "ATLAN-PYTHON-404-029",
+        "DQ rule with type {0} was not found.",
+        "Verify you have provided a valid DQ rule type.",
+        NotFoundError,
+    )
     CONFLICT_PASSTHROUGH = (
         409,
         "ATLAN-PYTHON-409-000",

pyatlan/model/assets/core/alpha__d_q_rule.py

@@ -7,6 +7,11 @@
 from datetime import datetime
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Set, overload
 
+from pyatlan.utils import init_guid, validate_required_fields,ErrorCode
+import json
+import uuid
+import time
+
 from pydantic.v1 import Field, validator
 
 from pyatlan.model.enums import (
@@ -25,18 +30,15 @@
     RelationField,
     TextField,
 )
+
 from pyatlan.model.structs import alpha_DQRuleConfigArguments, alpha_DQRuleThresholdObject
-from pyatlan.utils import init_guid, validate_required_fields
-# from pyatlan.client.atlan import AtlanClient
-import json
-import uuid
-import time
 
 from .data_quality import DataQuality
-# from fluent_search import FluentSearch
 
 if TYPE_CHECKING:
     from pyatlan.model.assets import Column
+    from pyatlan.client.atlan import AtlanClient
+
 
 
 class alpha_DQRule(DataQuality):
@@ -46,7 +48,7 @@ class alpha_DQRule(DataQuality):
     def custom_sql_creator(
         cls,
         *,
-        # client: AtlanClient,
+        client: AtlanClient,
         rule_name: str,
         asset: Asset,
         custom_sql: str,
@@ -61,8 +63,8 @@ def custom_sql_creator(
             [rule_name, asset, threshold_compare_operator, threshold_value, alert_priority,dimension, custom_sql],
         )
 
-        attributes = alpha_DQRule.Attributes.create(
-            # client=client,
+        attributes = alpha_DQRule.Attributes.creator(
+            client=client,
             rule_name=rule_name,
             rule_type="Custom SQL",
             asset=asset,
@@ -82,7 +84,7 @@ def custom_sql_creator(
     def table_level_rule_creator(
         cls,
         *,
-        # client:AtlanClient,
+        client:AtlanClient,
         rule_type: str,
         asset: Asset,
         threshold_compare_operator: alpha_DQRuleThresholdCompareOperator,
@@ -94,7 +96,7 @@ def table_level_rule_creator(
             [rule_type, asset, threshold_compare_operator, threshold_value, alert_priority],
         )
 
-        attributes = alpha_DQRule.Attributes.create(
+        attributes = alpha_DQRule.Attributes.creator(
             client=client,
             rule_type=rule_type,
             asset=asset,
@@ -115,7 +117,7 @@ def table_level_rule_creator(
     def column_level_rule_creator(
         cls,
         *,
-        # client:AtlanClient,
+        client:AtlanClient,
         rule_type: str,
         asset: Asset,
         column_qualified_name: str,
@@ -129,7 +131,7 @@ def column_level_rule_creator(
             [rule_type, asset, column_qualified_name, threshold_compare_operator, threshold_value, alert_priority],
         )
 
-        attributes = alpha_DQRule.Attributes.create(
+        attributes = alpha_DQRule.Attributes.creator(
             client=client,
             rule_type=rule_type,
             asset=asset,
@@ -878,10 +880,10 @@ def replace_char(c):
 
         @classmethod
         @init_guid
-        def create(
+        def creator(
             cls,
             *,
-            # client: AtlanClient,
+            client: AtlanClient,
             rule_name:str,
             rule_type: str,
             asset: Asset,
@@ -894,19 +896,34 @@ def create(
             custom_sql: Optional[str] = None,
             description: Optional[str] = None,
         ) -> alpha_DQRule.Attributes:
-            # request = (
-            #         FluentSearch()
-            #         .where(Asset.TYPE_NAME.eq("alpha_DQRuleTemplate"))
-            #         .include_on_results(Asset.NAME)
-            #         .include_on_results(Asset.DISPLAY_NAME)
-            #         .include_on_results(Asset.QUALIFIED_NAME)
-            # ).to_request()  # 
-            # for result in client.asset.search(request):
-            #     if result.display_name == rule_type:
-            #         template_rule_name = result.name
-            #         template_qualified_name = result.qualified_name
-            #         if dimension is None:
-            #             dimension = result.alpha_dq_rule_dimension
+            from pyatlan.model.fluent_search import FluentSearch
+
+            request = (
+                FluentSearch()
+                .where(Asset.TYPE_NAME.eq(alpha_DQRuleTemplate.__name__))
+                .include_on_results(alpha_DQRuleTemplate.NAME)
+                .include_on_results(alpha_DQRuleTemplate.QUALIFIED_NAME)
+                .include_on_results(alpha_DQRuleTemplate.DISPLAY_NAME)
+                .include_on_results(alpha_DQRuleTemplate.ALPHADQ_RULE_TEMPLATE_DIMENSION)
+                .include_on_results(alpha_DQRuleTemplate.ALPHADQ_RULE_TEMPLATE_CONFIG)
+            ).to_request()
+            for result in client.asset.search(request):
+                if result.display_name == rule_type:
+                    template_rule_name = result.name
+                    template_qualified_name = result.qualified_name
+                    
+                    if dimension is None:
+                        dimension = result.alpha_dq_rule_template_dimension
+
+                    if threshold_unit is None:
+                        threashold_object = result.alpha_dq_rule_template_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
+                # else:
+                #     raise ErrorCode.DQ_RULE_NOT_FOUND.exception_with_parameters(result.display_name)
 
             config_arguments_raw = alpha_DQRule.Attributes._generate_config_arguments_raw(
                 is_alert_enabled=True,
@@ -939,8 +956,8 @@ def create(
                 alpha_dq_rule_base_dataset= asset,
                 qualified_name=f"{asset.qualified_name}/rule/{str(cls._generate_uuid())}",
                 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),
+                alpha_dq_rule_template_name=template_rule_name,
+                alpha_dq_rule_template=alpha_DQRuleTemplate.ref_by_qualified_name(qualified_name=template_qualified_name),
             )
 
             if column_qualified_name is not None:

pyatlan/model/utils.py

@@ -25,11 +25,20 @@ def encoders():
     }
 
 
+def convert_with_fixed_prefix(input_str, fixed_prefix='alpha_dq'):
+    prefix = fixed_prefix
+    remaining = input_str[len(prefix)+1:]
+    parts = remaining.split('_')
+    camel_part = ''.join(word.capitalize() for word in parts)
+    return prefix + camel_part
+
 def to_camel_case(value: str) -> str:
     if not isinstance(value, str):
         raise ValueError("Value must be a string")
     if value == "__root__":
         return value
+    if value.startswith("alpha_dq"):
+        return convert_with_fixed_prefix(value)
     if value in CAMEL_CASE_OVERRIDES:
         return CAMEL_CASE_OVERRIDES[value]
     value = "".join(word.capitalize() for word in value.split("_"))