BLDX-721: Fixed backend bugs and added DataContract.save() / delete() methods

pyatlan/model/assets/core/data_contract.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from json import JSONDecodeError, loads
-from typing import ClassVar, List, Optional, Union, overload
+from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Type, Union, overload
 
 from pydantic.v1 import Field, validator
 
@@ -21,18 +21,32 @@
 
 from .catalog import Catalog
 
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+    from pyatlan.model.response import AssetMutationResponse
+
 
 class DataContract(Catalog):
     """Description"""
 
     @overload
     @classmethod
-    def creator(cls, asset_qualified_name: str, contract_json: str) -> DataContract: ...
+    def creator(
+        cls,
+        *,
+        asset_qualified_name: str,
+        asset_type: Type[Asset],
+        contract_json: str,
+    ) -> DataContract: ...
 
     @overload
     @classmethod
     def creator(
-        cls, asset_qualified_name: str, contract_spec: Union[DataContractSpec, str]
+        cls,
+        *,
+        asset_qualified_name: str,
+        asset_type: Type[Asset],
+        contract_spec: Union[DataContractSpec, str],
     ) -> DataContract: ...
 
     @classmethod
@@ -41,6 +55,7 @@ def creator(
         cls,
         *,
         asset_qualified_name: str,
+        asset_type: Type[Asset],
         contract_json: Optional[str] = None,
         contract_spec: Optional[Union[DataContractSpec, str]] = None,
     ) -> DataContract:
@@ -50,11 +65,71 @@ def creator(
         )
         attributes = DataContract.Attributes.creator(
             asset_qualified_name=asset_qualified_name,
+            asset_type=asset_type,
             contract_json=contract_json,
             contract_spec=contract_spec,
         )
         return cls(attributes=attributes)
 
+    @staticmethod
+    def save(
+        client: "AtlanClient",
+        contract: "DataContract",
+    ) -> "AssetMutationResponse":
+        """Save a DataContract.
+
+        The contract's ``data_contract_asset_latest`` relationship (set by
+        ``creator()``) links the contract to the governed asset automatically.
+
+        :param client: connectivity to an Atlan tenant
+        :param contract: DataContract to save (from ``DataContract.creator()``)
+        :returns: the result of the save
+        """
+        return client.asset.save(contract)
+
+    @staticmethod
+    def delete(
+        client: "AtlanClient",
+        contract_guid: str,
+    ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]":
+        """Delete (purge) a DataContract and clean up the linked asset.
+
+        Retrieves the contract to find the linked asset, clears
+        ``hasContract``, ``dataContractLatest``, and
+        ``dataContractLatestCertified`` on it, then hard-deletes the contract.
+
+        :param client: connectivity to an Atlan tenant
+        :param contract_guid: GUID of the DataContract to delete
+        :returns: tuple of (contract delete response, asset update response)
+        """
+        from pyatlan.model.assets.core.asset import Asset
+        from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset
+
+        contract = client.asset.get_by_guid(
+            contract_guid,
+            asset_type=DataContract,
+            attributes=[DataContract.DATA_CONTRACT_ASSET_LATEST],
+            related_attributes=[Asset.NAME, Asset.QUALIFIED_NAME, Asset.TYPE_NAME],
+        )
+        linked_asset = contract.data_contract_asset_latest
+        if not linked_asset or not linked_asset.guid:
+            raise ValueError(
+                "Cannot determine the linked asset for this contract. "
+                "Ensure the contract has a valid data_contract_asset_latest relationship."
+            )
+
+        delete_response = client.asset.purge_by_guid(contract_guid)
+        asset_update = IndistinctAsset()
+        asset_update.type_name = linked_asset.type_name
+        asset_update.guid = linked_asset.guid
+        asset_update.qualified_name = linked_asset.qualified_name
+        asset_update.name = linked_asset.name
+        asset_update.has_contract = False
+        asset_update.data_contract_latest = None  # type: ignore[assignment]
+        asset_update.data_contract_latest_certified = None  # type: ignore[assignment]
+        asset_response = client.asset.save(asset_update)
+        return delete_response, asset_response
+
     type_name: str = Field(default="DataContract", allow_mutation=False)
 
     @validator("type_name")
@@ -261,6 +336,7 @@ def creator(
             cls,
             *,
             asset_qualified_name: str,
+            asset_type: Type[Asset],
             contract_json: Optional[str] = None,
             contract_spec: Optional[Union[DataContractSpec, str]] = None,
         ) -> DataContract.Attributes:
@@ -309,6 +385,9 @@ def creator(
                 qualified_name=f"{asset_qualified_name}/contract",
                 data_contract_json=contract_json,
                 data_contract_spec=contract_spec,  # type: ignore[arg-type]
+                data_contract_asset_latest=asset_type.ref_by_qualified_name(
+                    asset_qualified_name
+                ),
             )
 
     attributes: DataContract.Attributes = Field(

pyatlan_v9/model/assets/asset.py

@@ -31,6 +31,7 @@
 from .anomalo_related import RelatedAnomaloCheck
 from .app_related import RelatedApplication, RelatedApplicationField
 from .asset_related import RelatedAsset
+from .catalog_related import RelatedCatalog
 from .data_mesh_related import RelatedDataProduct
 from .data_quality_related import RelatedDataQualityRule, RelatedMetric
 from .gtc_related import RelatedAtlasGlossaryTerm
@@ -990,6 +991,12 @@ class Asset(Referenceable):
     soda_checks: Union[List[RelatedSodaCheck], None, UnsetType] = UNSET
     """"""
 
+    data_contract_latest: Union[RelatedCatalog, None, UnsetType] = UNSET
+    """Latest data contract for this asset."""
+
+    data_contract_latest_certified: Union[RelatedCatalog, None, UnsetType] = UNSET
+    """Latest certified data contract for this asset."""
+
     def __post_init__(self) -> None:
         if self.type_name is UNSET:
             self.type_name = "Asset"
@@ -2043,6 +2050,12 @@ class AssetRelationshipAttributes(ReferenceableRelationshipAttributes):
     soda_checks: Union[List[RelatedSodaCheck], None, UnsetType] = UNSET
     """"""
 
+    data_contract_latest: Union[RelatedCatalog, None, UnsetType] = UNSET
+    """Latest data contract for this asset."""
+
+    data_contract_latest_certified: Union[RelatedCatalog, None, UnsetType] = UNSET
+    """Latest certified data contract for this asset."""
+
 
 class AssetNested(ReferenceableNested):
     """Asset in nested API format for high-performance serialization."""
@@ -2081,6 +2094,8 @@ class AssetNested(ReferenceableNested):
     "readme",
     "schema_registry_subjects",
     "soda_checks",
+    "data_contract_latest",
+    "data_contract_latest_certified",
 ]
 
 

pyatlan_v9/model/assets/data_contract.py

@@ -7,7 +7,7 @@
 
 import re
 from json import JSONDecodeError, loads
-from typing import Union
+from typing import TYPE_CHECKING, Tuple, Type, Union
 
 from msgspec import UNSET, UnsetType
 
@@ -32,6 +32,10 @@
 )
 from .catalog_related import RelatedCatalog
 
+if TYPE_CHECKING:
+    from pyatlan_v9.client.atlan import AtlanClient
+    from pyatlan_v9.model.response import AssetMutationResponse
+
 
 @register_asset
 class DataContract(Catalog):
@@ -63,28 +67,118 @@ class DataContract(Catalog):
     data_contract_previous_version: Union[RelatedCatalog, None, UnsetType] = UNSET
     """Previous version in this contract chain."""
 
+    def __post_init__(self) -> None:
+        self.type_name = "DataContract"
+
     @classmethod
     @init_guid
     def creator(
         cls,
         *,
         asset_qualified_name: str,
+        asset_type: Type["Asset"],
         contract_json: Union[str, None] = None,
         contract_spec: Union[DataContractSpec, str, None] = None,
     ) -> "DataContract":
-        """Create a new DataContract asset."""
+        """Create a new DataContract asset.
+
+        :param asset_qualified_name: qualified name of the asset this contract governs
+        :param asset_type: the type of the governed asset (e.g. ``Table``, ``View``)
+        :param contract_json: deprecated JSON representation of the contract
+        :param contract_spec: YAML contract spec (string or ``DataContractSpec``)
+        """
         attrs = DataContract.Attributes.creator(
             asset_qualified_name=asset_qualified_name,
             contract_json=contract_json,
             contract_spec=contract_spec,
         )
+        asset_ref = RelatedAsset(qualified_name=asset_qualified_name)
+        asset_ref.type_name = asset_type.__name__
         return cls(
             name=attrs.name,
             qualified_name=attrs.qualified_name,
             data_contract_json=attrs.data_contract_json,
             data_contract_spec=attrs.data_contract_spec,
+            data_contract_asset_latest=asset_ref,
         )
 
+    @staticmethod
+    def save(
+        client: "AtlanClient",
+        contract: "DataContract",
+    ) -> "AssetMutationResponse":
+        """Save a DataContract.
+
+        The contract's ``data_contract_asset_latest`` relationship (set by
+        ``creator()``) links the contract to the governed asset automatically.
+
+        :param client: connectivity to an Atlan tenant
+        :param contract: DataContract to save (from ``DataContract.creator()``)
+        :returns: the result of the save
+        """
+        return client.asset.save(contract)
+
+    @staticmethod
+    def delete(
+        client: "AtlanClient",
+        contract_guid: str,
+    ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]":
+        """Delete (purge) a DataContract and clean up the linked asset.
+
+        Retrieves the contract to find the linked asset, clears
+        ``hasContract``, ``dataContractLatest``, and
+        ``dataContractLatestCertified`` on it, then hard-deletes the contract.
+
+        :param client: connectivity to an Atlan tenant
+        :param contract_guid: GUID of the DataContract to delete
+        :returns: tuple of (contract delete response, asset update response)
+        """
+        from pyatlan.client.constants import BULK_UPDATE
+        from pyatlan_v9.client.asset import _parse_mutation_response
+        from pyatlan_v9.model.assets.asset import Asset
+        from pyatlan_v9.model.assets.referenceable import Referenceable
+
+        contract = client.asset.get_by_guid(
+            contract_guid,
+            asset_type=DataContract,
+            attributes=["dataContractAssetLatest"],
+            related_attributes=[
+                Asset.NAME,
+                Referenceable.QUALIFIED_NAME,
+                Referenceable.TYPE_NAME,
+            ],
+        )
+        linked_asset = contract.data_contract_asset_latest
+        if not linked_asset or not linked_asset.guid:
+            raise ValueError(
+                "Cannot determine the linked asset for this contract. "
+                "Ensure the contract has a valid data_contract_asset_latest relationship."
+            )
+
+        delete_response = client.asset.purge_by_guid(contract_guid)
+
+        # Build raw payload because v9 Asset doesn't model these fields
+        asset_payload = {
+            "entities": [
+                {
+                    "guid": linked_asset.guid,
+                    "typeName": linked_asset.type_name,
+                    "attributes": {
+                        "qualifiedName": linked_asset.qualified_name,
+                        "name": linked_asset.name,
+                        "hasContract": False,
+                    },
+                    "relationshipAttributes": {
+                        "dataContractLatest": None,
+                        "dataContractLatestCertified": None,
+                    },
+                }
+            ]
+        }
+        raw_json = client._call_api(BULK_UPDATE, {}, asset_payload)
+        asset_response = _parse_mutation_response(raw_json)
+        return delete_response, asset_response
+
     @classmethod
     def updater(cls, *, qualified_name: str, name: str) -> "DataContract":
         """Create a DataContract instance for update operations."""

pyatlan_v9/model/assets/entity.py

@@ -123,6 +123,9 @@ class Entity(msgspec.Struct, kw_only=True, omit_defaults=True, rename="camel"):
     status: Union[str, UnsetType] = UNSET
     """Entity status (ACTIVE, DELETED, PURGED)."""
 
+    delete_handler: Union[str, UnsetType] = UNSET
+    """Handler used for deletion of this entity (e.g. HARD, SOFT, PURGE)."""
+
     version: Union[int, UnsetType] = UNSET
     """Version number of this entity."""
 

pyatlan_v9/model/assets/referenceable.py

@@ -290,6 +290,7 @@ class ReferenceableNested(
     guid: Union[Any, UnsetType] = UNSET
     type_name: Union[Any, UnsetType] = UNSET
     status: Union[Any, UnsetType] = UNSET
+    delete_handler: Union[Any, UnsetType] = UNSET
     version: Union[Any, UnsetType] = UNSET
     create_time: Union[Any, UnsetType] = UNSET
     update_time: Union[Any, UnsetType] = UNSET
@@ -368,6 +369,7 @@ def _referenceable_to_nested(referenceable: Referenceable) -> ReferenceableNeste
         guid=referenceable.guid,
         type_name=referenceable.type_name,
         status=referenceable.status,
+        delete_handler=referenceable.delete_handler,
         version=referenceable.version,
         create_time=referenceable.create_time,
         update_time=referenceable.update_time,
@@ -413,6 +415,7 @@ def _referenceable_from_nested(nested: ReferenceableNested) -> Referenceable:
         guid=nested.guid,
         type_name=nested.type_name,
         status=nested.status,
+        delete_handler=nested.delete_handler,
         version=nested.version,
         create_time=nested.create_time,
         update_time=nested.update_time,