[generator] Updated templates as per latest refactoring changes

Author: Aryamanz29

pyatlan/generator/templates/imports.jinja2

@@ -127,3 +127,5 @@ from pyatlan.model.data_mesh import DataProductsAssetsDSL
 from pyatlan.model.contract import DataContractSpec
 from pyatlan.model.lineage_ref import LineageRef
 from pyatlan.model.utils import construct_object_key
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient

pyatlan/generator/templates/methods/asset/badge.jinja2

@@ -4,6 +4,7 @@
     def creator(
         cls,
         *,
+        client: AtlanClient,
         name: StrictStr,
         cm_name: str,
         cm_attribute: str,
@@ -12,6 +13,7 @@
         return cls(
             status=EntityStatus.ACTIVE,
             attributes=Badge.Attributes.create(
+                client=client,
                 name=name,
                 cm_name=cm_name,
                 cm_attribute=cm_attribute,
@@ -24,6 +26,7 @@
     def create(
         cls,
         *,
+        client: AtlanClient,
         name: StrictStr,
         cm_name: str,
         cm_attribute: str,
@@ -38,6 +41,7 @@
             stacklevel=2,
         )
         return cls.creator(
+            client=client,
             name=name,
             cm_name=cm_name,
             cm_attribute=cm_attribute,

pyatlan/generator/templates/methods/asset/connection.jinja2

@@ -4,6 +4,7 @@
     def creator(
         cls,
         *,
+        client: AtlanClient,
         name: str,
         connector_type: AtlanConnectorType,
         admin_users: Optional[List[str]] = None,
@@ -12,7 +13,9 @@
         host: Optional[str] = None,
         port: Optional[int] = None,
     ) -> Connection:
-        validate_required_fields(["name", "connector_type"], [name, connector_type])
+        validate_required_fields(
+            ["client", "name", "connector_type"], [client, name, connector_type]
+        )
         if not admin_users and not admin_groups and not admin_roles:
             raise ValueError(
                 "One of admin_user, admin_groups or admin_roles is required"
@@ -23,6 +26,9 @@
             connector_name=connector_type.value,
             category=connector_type.category.value,
         )
+        client.user_cache.validate_names(names=admin_users or [])
+        client.role_cache.validate_idstrs(idstrs=admin_roles or [])
+        client.group_cache.validate_aliases(aliases=admin_groups or [])
         attr.admin_users = set() if admin_users is None else set(admin_users)
         attr.admin_groups = set() if admin_groups is None else set(admin_groups)
         attr.admin_roles = set() if admin_roles is None else set(admin_roles)
@@ -35,6 +41,7 @@
     def create(
         cls,
         *,
+        client: AtlanClient,
         name: str,
         connector_type: AtlanConnectorType,
         admin_users: Optional[List[str]] = None,
@@ -50,6 +57,7 @@
             stacklevel=2,
         )
         return cls.creator(
+            client=client,
             name=name,
             connector_type=connector_type,
             admin_users=admin_users,

pyatlan/generator/templates/methods/asset/data_product.jinja2

@@ -85,18 +85,15 @@
             name=name,
         )
 
-    def get_assets(self, client: Optional[AtlanClient] = None):
+    def get_assets(self, client: AtlanClient):
         """
         Reterieves list of all assets linked to the provided data product.
 
-        :param client: connectivity to an Atlan tenant (optional). If not provided, the default client will be used.
+        :param client: connectivity to an Atlan tenant.
         :raises InvalidRequestError: if DataProduct asset DSL cannot be found (does not exist) in Atlan
         :raises AtlanError: if there is an issue interacting with the API
         :returns: instance of `IndexSearchResults` with list of all assets linked to the provided data product
         """
-        from pyatlan.client.atlan import AtlanClient
-
-        client = AtlanClient.get_current_client() if not client else client
         dp_dsl = self.data_product_assets_d_s_l
         if not dp_dsl:
             raise ErrorCode.MISSING_DATA_PRODUCT_ASSET_DSL.exception_with_parameters()

pyatlan/generator/templates/methods/asset/purpose.jinja2

@@ -23,6 +23,7 @@
     def create_metadata_policy(
         cls,
         *,
+        client: AtlanClient,
         name: str,
         purpose_id: str,
         policy_type: AuthPolicyType,
@@ -32,8 +33,8 @@
         all_users: bool = False,
     ) -> AuthPolicy:
         validate_required_fields(
-            ["name", "purpose_id", "policy_type", "actions"],
-            [name, purpose_id, policy_type, actions],
+            ["client", "name", "purpose_id", "policy_type", "actions"],
+            [client, name, purpose_id, policy_type, actions],
         )
         target_found = False
         policy = AuthPolicy._AuthPolicy__create(name=name)  # type: ignore[attr-defined]
@@ -51,12 +52,8 @@
             policy.policy_groups = {"public"}
         else:
             if policy_groups:
-                from pyatlan.client.atlan import AtlanClient
-
                 for group_name in policy_groups:
-                    if not AtlanClient.get_current_client().group_cache.get_id_for_name(
-                        group_name
-                    ):
+                    if not client.group_cache.get_id_for_name(group_name):
                         raise ValueError(
                             f"Provided group name {group_name} was not found in Atlan."
                         )
@@ -65,12 +62,8 @@
             else:
                 policy.policy_groups = None
             if policy_users:
-                from pyatlan.client.atlan import AtlanClient
-
                 for username in policy_users:
-                    if not AtlanClient.get_current_client().user_cache.get_id_for_name(
-                        username
-                    ):
+                    if not client.user_cache.get_id_for_name(username):
                         raise ValueError(
                             f"Provided username {username} was not found in Atlan."
                         )
@@ -87,6 +80,7 @@
     def create_data_policy(
         cls,
         *,
+        client: AtlanClient,
         name: str,
         purpose_id: str,
         policy_type: AuthPolicyType,
@@ -95,7 +89,8 @@
         all_users: bool = False,
     ) -> AuthPolicy:
         validate_required_fields(
-            ["name", "purpose_id", "policy_type"], [name, purpose_id, policy_type]
+            ["client", "name", "purpose_id", "policy_type"],
+            [client, name, purpose_id, policy_type],
         )
         policy = AuthPolicy._AuthPolicy__create(name=name)  # type: ignore[attr-defined]
         policy.policy_actions = {DataAction.SELECT.value}
@@ -112,12 +107,8 @@
             policy.policy_groups = {"public"}
         else:
             if policy_groups:
-                from pyatlan.client.atlan import AtlanClient
-
                 for group_name in policy_groups:
-                    if not AtlanClient.get_current_client().group_cache.get_id_for_name(
-                        group_name
-                    ):
+                    if not client.group_cache.get_id_for_name(group_name):
                         raise ValueError(
                             f"Provided group name {group_name} was not found in Atlan."
                         )
@@ -126,12 +117,8 @@
             else:
                 policy.policy_groups = None
             if policy_users:
-                from pyatlan.client.atlan import AtlanClient
-
                 for username in policy_users:
-                    if not AtlanClient.get_current_client().user_cache.get_id_for_name(
-                        username
-                    ):
+                    if not client.user_cache.get_id_for_name(username):
                         raise ValueError(
                             f"Provided username {username} was not found in Atlan."
                         )

pyatlan/generator/templates/methods/attribute/badge.jinja2

@@ -4,18 +4,16 @@
         def create(
             cls,
             *,
+            client: AtlanClient,
             name: StrictStr,
             cm_name: str,
             cm_attribute: str,
             badge_conditions: List[BadgeCondition],
         ) -> Badge.Attributes:
             validate_required_fields(
-                ["name", "cm_name", "cm_attribute", "badge_conditions"],
-                [name, cm_name, cm_attribute, badge_conditions],
+                ["client", "name", "cm_name", "cm_attribute", "badge_conditions"],
+                [client, name, cm_name, cm_attribute, badge_conditions],
             )
-            from pyatlan.client.atlan import AtlanClient
-
-            client = AtlanClient.get_current_client()
             cm_id = client.custom_metadata_cache.get_id_for_name(cm_name)
             cm_attr_id = client.custom_metadata_cache.get_attr_id_for_name(
                 set_name=cm_name, attr_name=cm_attribute

pyatlan/generator/templates/methods/attribute/collection.jinja2

@@ -7,7 +7,7 @@
             client: AtlanClient,
             name: str,
         ) -> Collection.Attributes:
-            validate_required_fields(["name"], [name])
+            validate_required_fields(["client", "name"], [client, name])
             return Collection.Attributes(
                 name=name,
                 qualified_name=Collection._generate_qualified_name(client),

pyatlan/generator/templates/methods/attribute/connection.jinja2

@@ -1,32 +1,2 @@
 
         is_loaded: bool = Field(default=True)
-
-        @validator("admin_users")
-        def admin_users_valid(cls, admin_users, values):
-            from pyatlan.client.atlan import AtlanClient
-
-            if values.get("is_loaded", False):
-                AtlanClient.get_current_client().user_cache.validate_names(
-                    names=admin_users
-                )
-            return admin_users
-
-        @validator("admin_roles")
-        def admin_roles_valid(cls, admin_roles, values):
-            from pyatlan.client.atlan import AtlanClient
-
-            if values.get("is_loaded", False):
-                AtlanClient.get_current_client().role_cache.validate_idstrs(
-                    idstrs=admin_roles
-                )
-            return admin_roles
-
-        @validator("admin_groups")
-        def admin_groups_valid(cls, admin_groups, values):
-            from pyatlan.client.atlan import AtlanClient
-
-            if values.get("is_loaded", False):
-                AtlanClient.get_current_client().group_cache.validate_aliases(
-                    aliases=admin_groups
-                )
-                return admin_groups

pyatlan/generator/templates/referenceable_attributes.jinja2

@@ -90,7 +90,7 @@
 
     type_name: str = Field(default="Referenceable", description='Name of the type definition that defines this instance.'
     )
-    _metadata_proxy: CustomMetadataProxy = PrivateAttr()
+    _metadata_proxy: CustomMetadataProxy = PrivateAttr(default=None)
     attributes: {{entity_def.name}}.Attributes = Field(
         default_factory = lambda : {{entity_def.name}}.Attributes(),
         description='Map of attributes in the instance and their values. The specific keys of this map will vary '

pyatlan/generator/templates/referenceable_methods.jinja2

@@ -2,9 +2,6 @@
     def __init__(__pydantic_self__, **data: Any) -> None:
         super().__init__(**data)
         __pydantic_self__.__fields_set__.update(["attributes", "type_name"])
-        __pydantic_self__._metadata_proxy = CustomMetadataProxy(
-            __pydantic_self__.business_attributes
-        )
 
     @root_validator(pre=True)
     def parse_custom_attributes(cls, values):
@@ -24,21 +21,39 @@
         return values
 
     def json(self, *args, **kwargs) -> str:
-        if self._metadata_proxy and self._metadata_proxy.business_attributes:
+        if not self._metadata_proxy and kwargs.get("client"):
+            self._metadata_proxy = CustomMetadataProxy(
+                client=kwargs.get("client"),  # type: ignore[arg-type]
+                business_attributes=self.business_attributes,
+            )
             self.business_attributes = self._metadata_proxy.business_attributes
         return super().json(**kwargs)
 
     def validate_required(self):
         if not self.create_time or self.created_by:
             self.attributes.validate_required()
 
-    def get_custom_metadata(self, name: str) -> CustomMetadataDict:
+    def get_custom_metadata(self, client: AtlanClient, name: str) -> CustomMetadataDict:
+        if not self._metadata_proxy:
+            self._metadata_proxy = CustomMetadataProxy(
+                business_attributes=self.business_attributes, client=client
+            )
         return self._metadata_proxy.get_custom_metadata(name=name)
 
-    def set_custom_metadata(self, custom_metadata: CustomMetadataDict):
+    def set_custom_metadata(
+        self, client: AtlanClient, custom_metadata: CustomMetadataDict
+    ):
+        if not self._metadata_proxy:
+            self._metadata_proxy = CustomMetadataProxy(
+                business_attributes=self.business_attributes, client=client
+            )
         return self._metadata_proxy.set_custom_metadata(custom_metadata=custom_metadata)
 
-    def flush_custom_metadata(self):
+    def flush_custom_metadata(self, client: AtlanClient):
+        if not self._metadata_proxy:
+            self._metadata_proxy = CustomMetadataProxy(
+                business_attributes=self.business_attributes, client=client
+            )
         self.business_attributes = self._metadata_proxy.business_attributes
 
     @classmethod
@@ -59,13 +74,4 @@
 
     @property
     def atlan_tag_names(self) -> List[str]:
-        from pyatlan.client.atlan import AtlanClient
-        from pyatlan.model.constants import DELETED_
-
-        if self.classification_names:
-            return [
-                AtlanClient.get_current_client().atlan_tag_cache.get_name_for_id(tag_id)
-                or DELETED_
-                for tag_id in self.classification_names
-            ]
-        return []
+        return self.classification_names or []