Merge pull request #731 from atlanhq/APP-8880

APP-8880: Ensure custom metadata is flushed before bulk `save()`

Author: Aryamanz29

pyatlan/client/aio/asset.py

@@ -353,14 +353,14 @@ async def save(
         :raises ApiError: if a connection was created and blocking until policies are synced overruns the retry limit
         """
 
-        query_params, request = Save.prepare_request(
+        query_params, request = await Save.prepare_request_async(
             entity=entity,
             replace_atlan_tags=replace_atlan_tags,
             replace_custom_metadata=replace_custom_metadata,
             overwrite_custom_metadata=overwrite_custom_metadata,
             append_atlan_tags=append_atlan_tags,
+            client=self._client,  # type: ignore[arg-type]
         )
-        Save.validate_and_flush_entities(request.entities, self._client)
         raw_json = await self._client._call_api(BULK_UPDATE, query_params, request)
         response = Save.process_response(raw_json)
         if connections_created := response.assets_created(Connection):

pyatlan/client/asset.py

@@ -447,8 +447,8 @@ def save(
             replace_custom_metadata=replace_custom_metadata,
             overwrite_custom_metadata=overwrite_custom_metadata,
             append_atlan_tags=append_atlan_tags,
+            client=self._client,  # type: ignore[arg-type]
         )
-        Save.validate_and_flush_entities(request.entities, self._client)
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
         response = Save.process_response(raw_json)
         if connections_created := response.assets_created(Connection):
@@ -564,7 +564,7 @@ def save_replacing_cm(
         query_params, request = Save.prepare_request_replacing_cm(
             entity=entity,
             replace_atlan_tags=replace_atlan_tags,
-            client=self._client,
+            client=self._client,  # type: ignore[arg-type]
         )
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
         return Save.process_response_replacing_cm(raw_json)

pyatlan/client/common/asset.py

@@ -61,6 +61,8 @@
 from pyatlan.utils import unflatten_custom_metadata_for_entity
 
 if TYPE_CHECKING:
+    from pyatlan.client.aio import AsyncAtlanClient
+    from pyatlan.client.atlan import AtlanClient
     from pyatlan.model.fluent_search import FluentSearch
 
 LOGGER = logging.getLogger(__name__)
@@ -676,6 +678,7 @@ def prepare_request(
         replace_custom_metadata: bool = False,
         overwrite_custom_metadata: bool = False,
         append_atlan_tags: bool = False,
+        client: Optional[AtlanClient] = None,
     ) -> tuple[Dict[str, Any], BulkRequest[Asset]]:
         """
         Prepare the request for saving assets.
@@ -685,6 +688,48 @@ def prepare_request(
         :param replace_custom_metadata: replaces any custom metadata with non-empty values provided
         :param overwrite_custom_metadata: overwrites any custom metadata, even with empty values
         :param append_atlan_tags: whether to add/update/remove AtlanTags during an update
+        :param client: the Atlan client instance for flushing custom metadata
+        :returns: tuple of (query_params, bulk_request)
+        """
+        query_params = {
+            "replaceTags": replace_atlan_tags,
+            "appendTags": append_atlan_tags,
+            "replaceBusinessAttributes": replace_custom_metadata,
+            "overwriteBusinessAttributes": overwrite_custom_metadata,
+        }
+
+        entities: List[Asset] = []
+        if isinstance(entity, list):
+            entities.extend(entity)
+        else:
+            entities.append(entity)
+
+        if not client:
+            raise ValueError(
+                "AtlanClient instance must be provided to validate and flush cm for assets."
+            )
+        # Validate and flush entities BEFORE creating the BulkRequest
+        Save.validate_and_flush_entities(entities, client)
+        return query_params, BulkRequest[Asset](entities=entities)
+
+    @staticmethod
+    async def prepare_request_async(
+        entity: Union[Asset, List[Asset]],
+        replace_atlan_tags: bool = False,
+        replace_custom_metadata: bool = False,
+        overwrite_custom_metadata: bool = False,
+        append_atlan_tags: bool = False,
+        client: Optional[AsyncAtlanClient] = None,
+    ) -> tuple[Dict[str, Any], BulkRequest[Asset]]:
+        """
+        Prepare the request for saving assets.
+
+        :param entity: one or more assets to save
+        :param replace_atlan_tags: whether to replace AtlanTags during an update
+        :param replace_custom_metadata: replaces any custom metadata with non-empty values provided
+        :param overwrite_custom_metadata: overwrites any custom metadata, even with empty values
+        :param append_atlan_tags: whether to add/update/remove AtlanTags during an update
+        :param client: Optional[AsyncAtlanClient] = None,
         :returns: tuple of (query_params, bulk_request)
         """
         query_params = {
@@ -700,10 +745,16 @@ def prepare_request(
         else:
             entities.append(entity)
 
+        if not client:
+            raise ValueError(
+                "AsyncAtlanClient instance must be provided to validate and flush cm for assets."
+            )
+        # Validate and flush entities BEFORE creating the BulkRequest
+        await Save.validate_and_flush_entities_async(entities, client)
         return query_params, BulkRequest[Asset](entities=entities)
 
     @staticmethod
-    def validate_and_flush_entities(entities: List[Asset], client) -> None:
+    def validate_and_flush_entities(entities: List[Asset], client: AtlanClient) -> None:
         """
         Validate required fields and flush custom metadata for each asset.
 
@@ -714,6 +765,20 @@ def validate_and_flush_entities(entities: List[Asset], client) -> None:
             asset.validate_required()
             asset.flush_custom_metadata(client=client)
 
+    @staticmethod
+    async def validate_and_flush_entities_async(
+        entities: List[Asset], client: AsyncAtlanClient
+    ) -> None:
+        """
+        Validate required fields and flush custom metadata for each asset.
+
+        :param entities: list of assets to validate and flush
+        :param client: the Atlan client instance
+        """
+        for asset in entities:
+            asset.validate_required()
+            await asset.flush_custom_metadata_async(client=client)
+
     @staticmethod
     def process_response(raw_json: Dict[str, Any]) -> AssetMutationResponse:
         """
@@ -750,7 +815,7 @@ def log_connections_finished():
     def prepare_request_replacing_cm(
         entity: Union[Asset, List[Asset]],
         replace_atlan_tags: bool = False,
-        client=None,
+        client: Optional[AtlanClient] = None,
     ) -> tuple[Dict[str, Any], BulkRequest[Asset]]:
         """
         Prepare the request for saving assets with replacing custom metadata.
@@ -773,12 +838,12 @@ def prepare_request_replacing_cm(
         else:
             entities.append(entity)
 
+        if not client:
+            raise ValueError(
+                "AtlanClient instance must be provided to validate and flush cm for assets."
+            )
         # Validate and flush entities BEFORE creating the BulkRequest
-        if client:
-            for asset in entities:
-                asset.validate_required()
-                asset.flush_custom_metadata(client=client)
-
+        Save.validate_and_flush_entities(entities, client)
         return query_params, BulkRequest[Asset](entities=entities)
 
     @staticmethod

tests/integration/custom_metadata_test.py

@@ -61,6 +61,7 @@
 
 
 CM_QUALITY = f"{MODULE_NAME}_DQ"
+CM_SQL = f"{MODULE_NAME}_SQL"
 CM_ATTR_QUALITY_COUNT = "Count"
 CM_ATTR_QUALITY_SQL = "SQL"
 CM_ATTR_QUALITY_TYPE = "Type"
@@ -432,6 +433,29 @@ def cm_dq(
     wait_for_successful_custometadatadef_purge(CM_QUALITY, client=client)
 
 
+@pytest.fixture(scope="module")
+def cm_sql(
+    client: AtlanClient,
+) -> Generator[CustomMetadataDef, None, None]:
+    attribute_defs = [
+        AttributeDef.create(
+            client=client,
+            display_name=f"{MODULE_NAME}_{CM_ATTR_QUALITY_SQL}",
+            attribute_type=AtlanCustomAttributePrimitiveType.SQL,
+        ),
+    ]
+    cm = create_custom_metadata(
+        client,
+        name=CM_SQL,
+        attribute_defs=attribute_defs,
+        logo="https://github.com/great-expectations/great_expectations/raw/develop/docs/docusaurus/static/img/"
+        "gx-mark-160.png",
+        locked=False,
+    )
+    yield cm
+    wait_for_successful_custometadatadef_purge(CM_SQL, client=client)
+
+
 def test_cm_dq(
     cm_dq: CustomMetadataDef,
 ):
@@ -1145,3 +1169,45 @@ def test_add_badge_cm_dq(
     assert (badges := response.assets_created(asset_type=Badge))
     assert len(badges) == 1
     client.asset.purge_by_guid(badges[0].guid)
+
+
+@pytest.mark.order()
+def test_save_merging_cm(
+    term: AtlasGlossaryTerm,
+    glossary: AtlasGlossary,
+    cm_sql: CustomMetadataDef,
+    client: AtlanClient,
+):
+    cm_sql_dict = CustomMetadataDict(client=client, name=CM_SQL)
+    cm_sql_dict[f"{MODULE_NAME}_{CM_ATTR_QUALITY_SQL}"] = "SELECT * FROM batman;"
+    assert term.qualified_name
+    assert term.name
+
+    to_update = AtlasGlossaryTerm.create_for_modification(
+        qualified_name=term.qualified_name, name=term.name, glossary_guid=glossary.guid
+    )
+    to_update.set_custom_metadata(custom_metadata=cm_sql_dict, client=client)
+    response = client.asset.save_merging_cm(to_update, replace_atlan_tags=False)
+    assert response
+    assert len(response.assets_deleted(asset_type=AtlasGlossaryTerm)) == 0
+    assert len(response.assets_created(asset_type=AtlasGlossaryTerm)) == 0
+    assert len(response.assets_updated(asset_type=AtlasGlossaryTerm)) == 1
+
+    t = response.assets_updated(asset_type=AtlasGlossaryTerm)[0]
+    assert isinstance(t, AtlasGlossaryTerm)
+    assert t.guid == term.guid
+    assert t.qualified_name == term.qualified_name
+    assert term.qualified_name
+
+    x = client.asset.get_by_qualified_name(
+        qualified_name=term.qualified_name,
+        asset_type=AtlasGlossaryTerm,
+        ignore_relationships=False,
+    )
+    assert x
+    assert not x.is_incomplete
+    assert x.qualified_name == term.qualified_name
+    retrieved_cm_sql = x.get_custom_metadata(client=client, name=CM_SQL)
+    assert retrieved_cm_sql
+    sql_attr_value = retrieved_cm_sql[f"{MODULE_NAME}_{CM_ATTR_QUALITY_SQL}"]
+    assert sql_attr_value == "SELECT * FROM batman;"