[generator] Fixed generator as per latest async refactoring changes

Author: Aryamanz29

pyatlan/generator/class_generator.py

@@ -30,13 +30,15 @@
     ("array<int>", "Set[int]"),
     ("array<float>", "Set[float]"),
     ("array<long>", "Set[int]"),
+    ("array<double>", "Set[float]"),
     ("icon_type", "IconType"),
     ("string", "str"),
     ("date", "datetime"),
     ("array", "List"),
     ("boolean", "bool"),
     ("float", "float"),
     ("long", "int"),
+    ("double", "float"),
     ("__internal", "Internal"),
     ("certificate_status", "CertificateStatus"),
     ("map", "Dict"),
@@ -64,6 +66,7 @@
     "long": "int",
     "date": "int",
     "float": "float",
+    "double": "float",
     "string,string": "str, str",
     "string,long": "str, int",
 }
@@ -229,6 +232,11 @@ class AssetInfo:
         "Purpose",
         "alpha_DQRule",
     }
+    _ASSETS_REQUIRE_ASYNC_CLIENT = {
+        "Badge",
+        "Connection",
+        "Referenceable",
+    }
     _CORE_ASSETS = {
         "Referenceable",
         "Asset",
@@ -1002,6 +1010,9 @@ def filter_attributes_of_custom_entity_type():
     )
     args = parser.parse_args()
 
+    # Ensure docs directory exists
+    DOCS_DIR.mkdir(parents=True, exist_ok=True)
+
     type_defs = get_type_defs(args.typedefs_file)
     filter_attributes_of_custom_entity_type()
     AssetInfo.sub_type_names_to_ignore = type_defs.custom_entity_def_names

pyatlan/generator/templates/imports.jinja2

@@ -115,7 +115,9 @@ from pyatlan.model.structs import (
     alpha_DQRuleTemplateConfig,
     ResponseValue,
     FormField,
-    DatabricksAIModelVersionMetric
+    DatabricksAIModelVersionMetric,
+    AssetHistogram,
+    AppWorkflowRunStep,
 )
 from pyatlan.utils import (
     init_guid,

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

@@ -21,6 +21,28 @@
             ),
         )
 
+    @classmethod
+    @init_guid
+    async def creator_async(
+        cls,
+        *,
+        client: AsyncAtlanClient,
+        name: StrictStr,
+        cm_name: str,
+        cm_attribute: str,
+        badge_conditions: List[BadgeCondition],
+    ) -> Badge:
+        return cls(
+            status=EntityStatus.ACTIVE,
+            attributes=await Badge.Attributes.create_async(
+                client=client,
+                name=name,
+                cm_name=cm_name,
+                cm_attribute=cm_attribute,
+                badge_conditions=badge_conditions,
+            ),
+        )
+
     @classmethod
     @init_guid
     def create(

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

@@ -36,6 +36,61 @@
         attr.port = port
         return cls(attributes=attr)
 
+    @classmethod
+    @init_guid
+    async def creator_async(
+        cls,
+        *,
+        client: AsyncAtlanClient,
+        name: str,
+        connector_type: AtlanConnectorType,
+        admin_users: Optional[List[str]] = None,
+        admin_groups: Optional[List[str]] = None,
+        admin_roles: Optional[List[str]] = None,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+    ) -> Connection:
+        """
+        Async version of creator() method for creating a new connection.
+
+        :param client: async client to use for validation
+        :param name: name for the connection
+        :param connector_type: type of connector this connection is for
+        :param admin_users: list of admin usernames for the connection
+        :param admin_groups: list of admin group names for the connection
+        :param admin_roles: list of admin role GUIDs for the connection
+        :param host: hostname for the connection
+        :param port: port number for the connection
+        :returns: the new connection object
+        :raises ValueError: if required parameters are missing or invalid
+        """
+        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"
+            )
+        attr = cls.Attributes(
+            name=name,
+            qualified_name=connector_type.to_qualified_name(),
+            connector_name=connector_type.value,
+            category=connector_type.category.value,
+        )
+
+        # Use async cache validation methods
+        await client.user_cache.validate_names(names=admin_users or [])
+        await client.role_cache.validate_idstrs(idstrs=admin_roles or [])
+        await 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)
+
+        attr.host = host
+        attr.port = port
+        return cls(attributes=attr)
+
     @classmethod
     @init_guid
     def create(

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

@@ -24,3 +24,29 @@
                 badge_metadata_attribute=f"{cm_id}.{cm_attr_id}",
                 badge_conditions=badge_conditions,
             )
+
+        @classmethod
+        @init_guid
+        async def create_async(
+            cls,
+            *,
+            client: AsyncAtlanClient,
+            name: StrictStr,
+            cm_name: str,
+            cm_attribute: str,
+            badge_conditions: List[BadgeCondition],
+        ) -> Badge.Attributes:
+            validate_required_fields(
+                ["client", "name", "cm_name", "cm_attribute", "badge_conditions"],
+                [client, name, cm_name, cm_attribute, badge_conditions],
+            )
+            cm_id = await client.custom_metadata_cache.get_id_for_name(cm_name)
+            cm_attr_id = await client.custom_metadata_cache.get_attr_id_for_name(
+                set_name=cm_name, attr_name=cm_attribute
+            )
+            return Badge.Attributes(
+                name=name,
+                qualified_name=f"badges/global/{cm_id}.{cm_attr_id}",
+                badge_metadata_attribute=f"{cm_id}.{cm_attr_id}",
+                badge_conditions=badge_conditions,
+            )

pyatlan/generator/templates/module.jinja2

@@ -43,6 +43,10 @@ if TYPE_CHECKING:
     {% if asset_info.name == 'alpha_DQRule' %}
     from pyatlan.model.assets import Column
     {% endif %}
+    {% if asset_info.name in asset_info._ASSETS_REQUIRE_ASYNC_CLIENT %}
+    from pyatlan.client.aio.client import AsyncAtlanClient
+    from pyatlan.model.aio import AsyncCustomMetadataProxy
+    {% endif %}
 {% endif %}
 
 {% if asset_info.name == 'Asset' %}

pyatlan/generator/templates/referenceable_attributes.jinja2

@@ -96,6 +96,7 @@
     type_name: str = Field(default="Referenceable", description='Name of the type definition that defines this instance.'
     )
     _metadata_proxy: CustomMetadataProxy = PrivateAttr(default=None)
+    _async_metadata_proxy: AsyncCustomMetadataProxy = 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

@@ -21,11 +21,16 @@
         return values
 
     def json(self, *args, **kwargs) -> str:
+        from pyatlan.model.aio.custom_metadata import AsyncCustomMetadataProxy
         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._async_metadata_proxy = AsyncCustomMetadataProxy(
+                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)
 
@@ -56,6 +61,72 @@
             )
         self.business_attributes = self._metadata_proxy.business_attributes
 
+    async def flush_custom_metadata_async(self, client):
+        """
+        Async version of flush_custom_metadata for use with AsyncAtlanClient.
+        This properly handles async metadata proxy and sets business_attributes.
+        """
+        from pyatlan.model.aio.custom_metadata import AsyncCustomMetadataProxy
+
+        if not self._async_metadata_proxy:
+            self._async_metadata_proxy = AsyncCustomMetadataProxy(
+                business_attributes=self.business_attributes, client=client
+            )
+        # Get the business attributes asynchronously and set them
+        business_attrs = await self._async_metadata_proxy.business_attributes()
+        self.business_attributes = business_attrs
+
+    async def get_custom_metadata_async(self, client: AsyncAtlanClient, name: str):
+        """
+        Async version of get_custom_metadata.
+
+        :param client: async Atlan client to use for the request
+        :param name: human-readable name of the custom metadata set to retrieve
+        :returns: the requested custom metadata set, or an empty one if none exists
+        """
+        from pyatlan.model.aio.custom_metadata import AsyncCustomMetadataProxy
+
+        if not self._async_metadata_proxy:
+            self._async_metadata_proxy = AsyncCustomMetadataProxy(
+                business_attributes=self.business_attributes, client=client
+            )
+        return await self._async_metadata_proxy.get_custom_metadata(name=name)
+
+    async def set_custom_metadata_async(
+        self, client: AsyncAtlanClient, custom_metadata
+    ):
+        """
+        Async version of set_custom_metadata.
+
+        :param client: async Atlan client to use for the request
+        :param custom_metadata: the custom metadata to set on this asset
+        """
+        from pyatlan.model.aio.custom_metadata import AsyncCustomMetadataProxy
+
+        if not self._async_metadata_proxy:
+            self._async_metadata_proxy = AsyncCustomMetadataProxy(
+                business_attributes=self.business_attributes, client=client
+            )
+        return await self._async_metadata_proxy.set_custom_metadata(
+            custom_metadata=custom_metadata
+        )
+
+    async def flush_custom_metadata_async(self, client: AsyncAtlanClient):
+        """
+        Async version of flush_custom_metadata.
+
+        :param client: async Atlan client to use for the request
+        """
+        from pyatlan.model.aio.custom_metadata import AsyncCustomMetadataProxy
+
+        if not self._async_metadata_proxy:
+            self._async_metadata_proxy = AsyncCustomMetadataProxy(
+                business_attributes=self.business_attributes, client=client
+            )
+        self.business_attributes = (
+            await self._async_metadata_proxy.business_attributes()
+        )
+
     @classmethod
     def __get_validators__(cls):
         yield cls._convert_to_real_type_

pyatlan/generator/templates/structs.jinja2

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2024 Atlan Pte. Ltd.
+# Copyright 2025 Atlan Pte. Ltd.
 
 from __future__ import annotations
 
@@ -16,13 +16,15 @@ from pyatlan.model.enums import (
     FormFieldDimension,
     FormFieldType,
     alpha_DQRuleThresholdUnit,
+    AppWorkflowRunStatus,
 )
 from pyatlan.model.utils import to_camel_case
 from pyatlan.utils import select_optional_set_fields, validate_required_fields
 
 if TYPE_CHECKING:
     from pyatlan.cache.source_tag_cache import SourceTagName
     from pyatlan.client.atlan import AtlanClient
+    from pyatlan.client.aio import AsyncAtlanClient
 
 
 class AtlanObject(BaseModel):
@@ -123,6 +125,49 @@ class {{struct.name}}(AtlanObject):
             ),
         )
 
+    @classmethod
+    async def by_name_async(
+        cls,
+        client: AsyncAtlanClient,
+        name: SourceTagName,
+        source_tag_values: List[SourceTagAttachmentValue],
+        source_tag_sync_timestamp: Optional[datetime] = None,
+        is_source_tag_synced: Optional[bool] = None,
+        source_tag_sync_error: Optional[str] = None,
+    ):
+        """
+        Async version of by_name that creates a source-synced tag attachment with
+        a particular value when the attachment is synced to the source.
+
+        :param client: async connectivity to an Atlan tenant
+        :param name: unique name of the source tag in Atlan
+        :param source_tag_values: value of the tag attachment from the source
+        :param is_source_tag_synced: whether the tag attachment has been synced at the source (True) or not (False)
+        :param source_tag_sync_timestamp: time (epoch) when the tag attachment was synced at the source, in milliseconds
+        :param source_tag_sync_error: error message if the tag attachment sync at the source failed
+        :returns: a SourceTagAttachment with the provided information
+        :raises AtlanError: on any error communicating via the underlying APIs
+        :raises NotFoundError: if the source-synced tag cannot be resolved
+        """
+        tag = await client.source_tag_cache.get_by_name(name)
+        tag_connector_name = AtlanConnectorType._get_connector_type_from_qualified_name(
+            tag.qualified_name or ""
+        )
+        return cls.of(
+            source_tag_name=tag.name,
+            source_tag_qualified_name=tag.qualified_name,
+            source_tag_guid=tag.guid,
+            source_tag_connector_name=tag_connector_name,
+            source_tag_values=source_tag_values,
+            **select_optional_set_fields(
+                dict(
+                    is_source_tag_synced=is_source_tag_synced,
+                    source_tag_sync_timestamp=source_tag_sync_timestamp,
+                    source_tag_sync_error=source_tag_sync_error,
+                )
+            ),
+        )
+
     @classmethod
     def by_qualified_name(
         cls,