From ba8cb1b11f272f5fa41595e3e19b78e2a7f9dd4b Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Fri, 13 Mar 2026 14:02:55 +0530 Subject: [PATCH 1/9] fix: add DataContract.save_contract() and delete_contract() lifecycle helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bug 1: save() alone doesn't set dataContractLatest on the linked asset, breaking the UI Contracts tab. save_contract() saves the contract AND updates the linked asset's dataContractLatest + hasContract. Bug 2: purge_by_guid() doesn't clear hasContract on the linked asset, leaving it in a broken state. delete_contract() purges the contract AND clears hasContract + dataContractLatest on the linked asset. Bug 3 (QN uniqueness on soft-delete) is a backend issue — no SDK fix needed. Fixes BLDX-721 Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 71 +++++++++++++++++- pyatlan_v9/model/assets/data_contract.py | 74 ++++++++++++++++++- tests/integration/data_mesh_test.py | 27 ++++--- tests/unit/model/data_contract_test.py | 83 ++++++++++++++++++++++ tests_v9/integration/data_mesh_test.py | 27 ++++--- tests_v9/unit/model/data_contract_test.py | 58 +++++++++++++++ 6 files changed, 322 insertions(+), 18 deletions(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index f8817a812..397772edb 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/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, Union, overload from pydantic.v1 import Field, validator @@ -21,6 +21,10 @@ from .catalog import Catalog +if TYPE_CHECKING: + from pyatlan.client.atlan import AtlanClient + from pyatlan.model.response import AssetMutationResponse + class DataContract(Catalog): """Description""" @@ -55,6 +59,71 @@ def creator( ) return cls(attributes=attributes) + @staticmethod + def save_contract( + client: AtlanClient, + contract: DataContract, + linked_asset_guid: str, + ) -> Tuple[AssetMutationResponse, AssetMutationResponse]: + """Save a DataContract and link it to the associated asset. + + After saving the contract, this method also sets ``dataContractLatest`` + and ``hasContract`` on the linked asset so the UI Contracts tab works + correctly. + + :param client: connectivity to an Atlan tenant + :param contract: DataContract to save (from ``DataContract.creator()``) + :param linked_asset_guid: GUID of the asset this contract is for + :returns: tuple of (contract save response, asset update response) + """ + from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + + contract_response = client.asset.save(contract) + contracts = ( + contract_response.assets_created(DataContract) + or contract_response.assets_updated(DataContract) + ) + if not contracts: + return contract_response, contract_response + + saved = contracts[0] + asset_update = IndistinctAsset() + asset_update.guid = linked_asset_guid + asset_update.data_contract_latest = DataContract.ref_by_guid( + saved.guid + ) + asset_update.has_contract = True + asset_response = client.asset.save(asset_update) + return contract_response, asset_response + + @staticmethod + def delete_contract( + client: AtlanClient, + contract_guid: str, + linked_asset_guid: str, + ) -> Tuple[AssetMutationResponse, AssetMutationResponse]: + """Delete a DataContract and clean up the linked asset. + + This method purges (hard-deletes) the contract to avoid qualified-name + conflicts on re-creation, and clears ``hasContract`` and + ``dataContractLatest`` on the linked asset. + + :param client: connectivity to an Atlan tenant + :param contract_guid: GUID of the DataContract to delete + :param linked_asset_guid: GUID of the asset the contract was linked to + :returns: tuple of (contract delete response, asset update response) + """ + from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + + delete_response = client.asset.purge_by_guid(contract_guid) + + asset_update = IndistinctAsset() + asset_update.guid = linked_asset_guid + asset_update.has_contract = False + asset_update.data_contract_latest = 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") diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index 7230be37e..d862e114e 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/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, Union from msgspec import UNSET, UnsetType @@ -32,6 +32,10 @@ ) from .catalog_related import RelatedCatalog +if TYPE_CHECKING: + from pyatlan.client.atlan import AtlanClient + from pyatlan.model.response import AssetMutationResponse + @register_asset class DataContract(Catalog): @@ -95,6 +99,74 @@ def trim_to_required(self) -> "DataContract": """Return only required fields for update operations.""" return DataContract.updater(qualified_name=self.qualified_name, name=self.name) + @staticmethod + def save_contract( + client: "AtlanClient", + contract: "DataContract", + linked_asset_guid: str, + ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": + """Save a DataContract and link it to the associated asset. + + After saving the contract, this method also sets ``dataContractLatest`` + and ``hasContract`` on the linked asset so the UI Contracts tab works + correctly. + + :param client: connectivity to an Atlan tenant + :param contract: DataContract to save (from ``DataContract.creator()``) + :param linked_asset_guid: GUID of the asset this contract is for + :returns: tuple of (contract save response, asset update response) + """ + from pyatlan.model.assets.core.data_contract import ( + DataContract as V8DataContract, + ) + from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + + contract_response = client.asset.save(contract) + contracts = ( + contract_response.assets_created(V8DataContract) + or contract_response.assets_updated(V8DataContract) + ) + if not contracts: + return contract_response, contract_response + + saved = contracts[0] + asset_update = IndistinctAsset() + asset_update.guid = linked_asset_guid + asset_update.data_contract_latest = V8DataContract.ref_by_guid( + saved.guid + ) + asset_update.has_contract = True + asset_response = client.asset.save(asset_update) + return contract_response, asset_response + + @staticmethod + def delete_contract( + client: "AtlanClient", + contract_guid: str, + linked_asset_guid: str, + ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": + """Delete a DataContract and clean up the linked asset. + + This method purges (hard-deletes) the contract to avoid qualified-name + conflicts on re-creation, and clears ``hasContract`` and + ``dataContractLatest`` on the linked asset. + + :param client: connectivity to an Atlan tenant + :param contract_guid: GUID of the DataContract to delete + :param linked_asset_guid: GUID of the asset the contract was linked to + :returns: tuple of (contract delete response, asset update response) + """ + from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + + delete_response = client.asset.purge_by_guid(contract_guid) + + asset_update = IndistinctAsset() + asset_update.guid = linked_asset_guid + asset_update.has_contract = False + asset_update.data_contract_latest = None # type: ignore[assignment] + asset_response = client.asset.save(asset_update) + return delete_response, asset_response + def to_json(self, nested: bool = True, serde: Serde | None = None) -> str: """Convert to JSON string.""" if serde is None: diff --git a/tests/integration/data_mesh_test.py b/tests/integration/data_mesh_test.py index d79f1c08b..f172e1f4e 100644 --- a/tests/integration/data_mesh_test.py +++ b/tests/integration/data_mesh_test.py @@ -307,8 +307,12 @@ def contract( asset_qualified_name=table.qualified_name, contract_json=dumps(contract_json), ) - response = client.asset.save(contract) - result = response.assets_created(asset_type=DataContract)[0] + contract_response, _ = DataContract.save_contract( + client=client, + contract=contract, + linked_asset_guid=table.guid, + ) + result = contract_response.assets_created(asset_type=DataContract)[0] yield result delete_asset(client, guid=result.guid, asset_type=DataContract) @@ -505,18 +509,25 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") -def test_delete_contract(client: AtlanClient, contract: DataContract): - response = client.asset.purge_by_guid(contract.guid) - assert response - assert not response.assets_created(asset_type=DataContract) - assert not response.assets_updated(asset_type=DataContract) - deleted = response.assets_deleted(asset_type=DataContract) +def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): + assert table.guid + delete_response, asset_response = DataContract.delete_contract( + client=client, + contract_guid=contract.guid, + linked_asset_guid=table.guid, + ) + assert delete_response + assert not delete_response.assets_created(asset_type=DataContract) + assert not delete_response.assets_updated(asset_type=DataContract) + deleted = delete_response.assets_deleted(asset_type=DataContract) assert deleted assert len(deleted) == 1 assert deleted[0].guid == contract.guid assert deleted[0].qualified_name == contract.qualified_name assert deleted[0].delete_handler == "PURGE" assert deleted[0].status == EntityStatus.DELETED + # Verify the linked asset was cleaned up + assert asset_response @pytest.mark.order(after="test_retrieve_product") diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index d9888eba2..3aa90c485 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -1,11 +1,13 @@ from json import dumps from typing import Union +from unittest.mock import MagicMock, patch import pytest from pyatlan.errors import InvalidRequestError from pyatlan.model.assets import DataContract from pyatlan.model.contract import DataContractSpec +from pyatlan.model.response import AssetMutationResponse from tests.unit.model.constants import ( ASSET_QUALIFIED_NAME, DATA_CONTRACT_JSON, @@ -128,3 +130,84 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) + + +class TestSaveContract: + def test_save_contract_sets_data_contract_latest_and_has_contract(self): + """save_contract should save the contract then update the linked asset.""" + mock_client = MagicMock() + contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + saved_contract = MagicMock() + saved_contract.guid = "contract-guid-123" + + contract_response = MagicMock(spec=AssetMutationResponse) + contract_response.assets_created.return_value = [saved_contract] + asset_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.side_effect = [contract_response, asset_response] + + result = DataContract.save_contract( + client=mock_client, + contract=contract, + linked_asset_guid="asset-guid-456", + ) + + assert result == (contract_response, asset_response) + assert mock_client.asset.save.call_count == 2 + # Verify the second save call updates the linked asset + asset_update = mock_client.asset.save.call_args_list[1][0][0] + assert asset_update.guid == "asset-guid-456" + assert asset_update.has_contract is True + assert asset_update.data_contract_latest is not None + assert asset_update.data_contract_latest.guid == "contract-guid-123" + + def test_save_contract_handles_update(self): + """save_contract should work when contract is updated (not created).""" + mock_client = MagicMock() + contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + saved_contract = MagicMock() + saved_contract.guid = "contract-guid-789" + + contract_response = MagicMock(spec=AssetMutationResponse) + contract_response.assets_created.return_value = [] + contract_response.assets_updated.return_value = [saved_contract] + asset_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.side_effect = [contract_response, asset_response] + + result = DataContract.save_contract( + client=mock_client, + contract=contract, + linked_asset_guid="asset-guid-456", + ) + + assert result == (contract_response, asset_response) + assert mock_client.asset.save.call_count == 2 + + +class TestDeleteContract: + def test_delete_contract_purges_and_clears_has_contract(self): + """delete_contract should purge the contract and clear hasContract.""" + mock_client = MagicMock() + delete_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.purge_by_guid.return_value = delete_response + asset_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.return_value = asset_response + + result = DataContract.delete_contract( + client=mock_client, + contract_guid="contract-guid-123", + linked_asset_guid="asset-guid-456", + ) + + assert result == (delete_response, asset_response) + mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") + # Verify the asset update clears contract state + asset_update = mock_client.asset.save.call_args[0][0] + assert asset_update.guid == "asset-guid-456" + assert asset_update.has_contract is False + assert asset_update.data_contract_latest is None diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 8f0035881..8fef8c633 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -310,8 +310,12 @@ def contract( asset_qualified_name=table.qualified_name, contract_json=dumps(contract_json), ) - response = client.asset.save(contract) - result = response.assets_created(asset_type=DataContract)[0] + contract_response, _ = DataContract.save_contract( + client=client, + contract=contract, + linked_asset_guid=table.guid, + ) + result = contract_response.assets_created(asset_type=DataContract)[0] yield result delete_asset(client, guid=result.guid, asset_type=DataContract) @@ -518,18 +522,25 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") -def test_delete_contract(client: AtlanClient, contract: DataContract): - response = client.asset.purge_by_guid(contract.guid) - assert response - assert not response.assets_created(asset_type=DataContract) - assert not response.assets_updated(asset_type=DataContract) - deleted = response.assets_deleted(asset_type=DataContract) +def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): + assert table.guid + delete_response, asset_response = DataContract.delete_contract( + client=client, + contract_guid=contract.guid, + linked_asset_guid=table.guid, + ) + assert delete_response + assert not delete_response.assets_created(asset_type=DataContract) + assert not delete_response.assets_updated(asset_type=DataContract) + deleted = delete_response.assets_deleted(asset_type=DataContract) assert deleted assert len(deleted) == 1 assert deleted[0].guid == contract.guid assert deleted[0].qualified_name == contract.qualified_name assert deleted[0].delete_handler == "PURGE" assert deleted[0].status == EntityStatus.DELETED + # Verify the linked asset was cleaned up + assert asset_response @pytest.mark.order(after="test_retrieve_product") diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index 47929177a..94d857e0c 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -5,9 +5,11 @@ from json import dumps from typing import Union +from unittest.mock import MagicMock import pytest +from pyatlan.model.response import AssetMutationResponse from pyatlan_v9.errors import InvalidRequestError from pyatlan_v9.model import DataContract from pyatlan_v9.model.contract import DataContractSpec @@ -156,3 +158,59 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) + + +class TestSaveContract: + def test_save_contract_sets_data_contract_latest_and_has_contract(self): + """save_contract should save the contract then update the linked asset.""" + mock_client = MagicMock() + contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + saved_contract = MagicMock() + saved_contract.guid = "contract-guid-123" + + contract_response = MagicMock(spec=AssetMutationResponse) + contract_response.assets_created.return_value = [saved_contract] + asset_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.side_effect = [contract_response, asset_response] + + result = DataContract.save_contract( + client=mock_client, + contract=contract, + linked_asset_guid="asset-guid-456", + ) + + assert result == (contract_response, asset_response) + assert mock_client.asset.save.call_count == 2 + # Verify the second save call updates the linked asset + asset_update = mock_client.asset.save.call_args_list[1][0][0] + assert asset_update.guid == "asset-guid-456" + assert asset_update.has_contract is True + assert asset_update.data_contract_latest is not None + assert asset_update.data_contract_latest.guid == "contract-guid-123" + + +class TestDeleteContract: + def test_delete_contract_purges_and_clears_has_contract(self): + """delete_contract should purge the contract and clear hasContract.""" + mock_client = MagicMock() + delete_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.purge_by_guid.return_value = delete_response + asset_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.return_value = asset_response + + result = DataContract.delete_contract( + client=mock_client, + contract_guid="contract-guid-123", + linked_asset_guid="asset-guid-456", + ) + + assert result == (delete_response, asset_response) + mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") + # Verify the asset update clears contract state + asset_update = mock_client.asset.save.call_args[0][0] + assert asset_update.guid == "asset-guid-456" + assert asset_update.has_contract is False + assert asset_update.data_contract_latest is None From a3307adc1276847db74391b8cb78dd28f3e50688 Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Fri, 13 Mar 2026 14:34:52 +0530 Subject: [PATCH 2/9] refactor: rename save_contract/delete_contract to save/delete - Rename DataContract.save_contract() -> DataContract.save() - Rename DataContract.delete_contract() -> DataContract.delete() - Clean up v9 import path for V8DataContract - Add bug-replication context in integration test docstrings Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 4 ++-- pyatlan_v9/model/assets/data_contract.py | 9 ++++----- tests/integration/data_mesh_test.py | 12 ++++++++++-- tests/unit/model/data_contract_test.py | 6 +++--- tests_v9/integration/data_mesh_test.py | 4 ++-- tests_v9/unit/model/data_contract_test.py | 4 ++-- 6 files changed, 23 insertions(+), 16 deletions(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index 397772edb..c35084753 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/pyatlan/model/assets/core/data_contract.py @@ -60,7 +60,7 @@ def creator( return cls(attributes=attributes) @staticmethod - def save_contract( + def save( client: AtlanClient, contract: DataContract, linked_asset_guid: str, @@ -97,7 +97,7 @@ def save_contract( return contract_response, asset_response @staticmethod - def delete_contract( + def delete( client: AtlanClient, contract_guid: str, linked_asset_guid: str, diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index d862e114e..6b2c5c4a3 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/pyatlan_v9/model/assets/data_contract.py @@ -100,7 +100,7 @@ def trim_to_required(self) -> "DataContract": return DataContract.updater(qualified_name=self.qualified_name, name=self.name) @staticmethod - def save_contract( + def save( client: "AtlanClient", contract: "DataContract", linked_asset_guid: str, @@ -116,12 +116,11 @@ def save_contract( :param linked_asset_guid: GUID of the asset this contract is for :returns: tuple of (contract save response, asset update response) """ - from pyatlan.model.assets.core.data_contract import ( - DataContract as V8DataContract, - ) + from pyatlan.model.assets import DataContract as V8DataContract from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset contract_response = client.asset.save(contract) + # Response contains v8 model objects since client.asset is v8 contracts = ( contract_response.assets_created(V8DataContract) or contract_response.assets_updated(V8DataContract) @@ -140,7 +139,7 @@ def save_contract( return contract_response, asset_response @staticmethod - def delete_contract( + def delete( client: "AtlanClient", contract_guid: str, linked_asset_guid: str, diff --git a/tests/integration/data_mesh_test.py b/tests/integration/data_mesh_test.py index f172e1f4e..624e2dc8f 100644 --- a/tests/integration/data_mesh_test.py +++ b/tests/integration/data_mesh_test.py @@ -307,7 +307,7 @@ def contract( asset_qualified_name=table.qualified_name, contract_json=dumps(contract_json), ) - contract_response, _ = DataContract.save_contract( + contract_response, _ = DataContract.save( client=client, contract=contract, linked_asset_guid=table.guid, @@ -345,6 +345,10 @@ def updated_contract( def test_contract( client: AtlanClient, table: Table, product: DataProduct, contract: DataContract ): + """Verify DataContract.save() sets dataContractLatest and hasContract on the + linked asset — the behaviour that was broken when using client.asset.save() + alone (Bug 1 / BLDX-721). + """ assert product and product.guid product = client.asset.get_by_guid( guid=product.guid, asset_type=DataProduct, ignore_relationships=False @@ -510,8 +514,12 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): + """Verify DataContract.delete() purges the contract AND clears hasContract + on the linked asset — the cleanup that was missing when using + client.asset.purge_by_guid() alone (Bug 2 / BLDX-721). + """ assert table.guid - delete_response, asset_response = DataContract.delete_contract( + delete_response, asset_response = DataContract.delete( client=client, contract_guid=contract.guid, linked_asset_guid=table.guid, diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index 3aa90c485..2300d494e 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -148,7 +148,7 @@ def test_save_contract_sets_data_contract_latest_and_has_contract(self): asset_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.save.side_effect = [contract_response, asset_response] - result = DataContract.save_contract( + result = DataContract.save( client=mock_client, contract=contract, linked_asset_guid="asset-guid-456", @@ -179,7 +179,7 @@ def test_save_contract_handles_update(self): asset_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.save.side_effect = [contract_response, asset_response] - result = DataContract.save_contract( + result = DataContract.save( client=mock_client, contract=contract, linked_asset_guid="asset-guid-456", @@ -198,7 +198,7 @@ def test_delete_contract_purges_and_clears_has_contract(self): asset_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.save.return_value = asset_response - result = DataContract.delete_contract( + result = DataContract.delete( client=mock_client, contract_guid="contract-guid-123", linked_asset_guid="asset-guid-456", diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 8fef8c633..3476d2544 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -310,7 +310,7 @@ def contract( asset_qualified_name=table.qualified_name, contract_json=dumps(contract_json), ) - contract_response, _ = DataContract.save_contract( + contract_response, _ = DataContract.save( client=client, contract=contract, linked_asset_guid=table.guid, @@ -524,7 +524,7 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): assert table.guid - delete_response, asset_response = DataContract.delete_contract( + delete_response, asset_response = DataContract.delete( client=client, contract_guid=contract.guid, linked_asset_guid=table.guid, diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index 94d857e0c..d9d1eba70 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -176,7 +176,7 @@ def test_save_contract_sets_data_contract_latest_and_has_contract(self): asset_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.save.side_effect = [contract_response, asset_response] - result = DataContract.save_contract( + result = DataContract.save( client=mock_client, contract=contract, linked_asset_guid="asset-guid-456", @@ -201,7 +201,7 @@ def test_delete_contract_purges_and_clears_has_contract(self): asset_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.save.return_value = asset_response - result = DataContract.delete_contract( + result = DataContract.delete( client=mock_client, contract_guid="contract-guid-123", linked_asset_guid="asset-guid-456", From e4115f9873efa5406e5ae81e2700dacdfc018cb9 Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Fri, 13 Mar 2026 14:43:43 +0530 Subject: [PATCH 3/9] refactor: use v9-native types in DataContract.save()/delete() V9 DataContract now uses v9 AtlanClient, v9 AssetMutationResponse, and v9 DataContract for response parsing instead of v8 imports. The linked asset update uses a raw API payload since v9 Asset doesn't model dataContractLatest. Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 11 ++-- pyatlan_v9/model/assets/data_contract.py | 66 ++++++++++++++-------- tests/unit/model/data_contract_test.py | 2 +- tests_v9/unit/model/data_contract_test.py | 53 ++++++++++------- 4 files changed, 81 insertions(+), 51 deletions(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index c35084753..2aa640542 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/pyatlan/model/assets/core/data_contract.py @@ -79,19 +79,16 @@ def save( from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset contract_response = client.asset.save(contract) - contracts = ( - contract_response.assets_created(DataContract) - or contract_response.assets_updated(DataContract) - ) + contracts = contract_response.assets_created( + DataContract + ) or contract_response.assets_updated(DataContract) if not contracts: return contract_response, contract_response saved = contracts[0] asset_update = IndistinctAsset() asset_update.guid = linked_asset_guid - asset_update.data_contract_latest = DataContract.ref_by_guid( - saved.guid - ) + asset_update.data_contract_latest = DataContract.ref_by_guid(saved.guid) asset_update.has_contract = True asset_response = client.asset.save(asset_update) return contract_response, asset_response diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index 6b2c5c4a3..fcdd3930a 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/pyatlan_v9/model/assets/data_contract.py @@ -33,8 +33,8 @@ from .catalog_related import RelatedCatalog if TYPE_CHECKING: - from pyatlan.client.atlan import AtlanClient - from pyatlan.model.response import AssetMutationResponse + from pyatlan_v9.client.atlan import AtlanClient + from pyatlan_v9.model.response import AssetMutationResponse @register_asset @@ -116,26 +116,35 @@ def save( :param linked_asset_guid: GUID of the asset this contract is for :returns: tuple of (contract save response, asset update response) """ - from pyatlan.model.assets import DataContract as V8DataContract - from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + from pyatlan.client.constants import BULK_UPDATE + from pyatlan_v9.client.asset import _parse_mutation_response contract_response = client.asset.save(contract) - # Response contains v8 model objects since client.asset is v8 - contracts = ( - contract_response.assets_created(V8DataContract) - or contract_response.assets_updated(V8DataContract) - ) + contracts = contract_response.assets_created( + DataContract + ) or contract_response.assets_updated(DataContract) if not contracts: return contract_response, contract_response saved = contracts[0] - asset_update = IndistinctAsset() - asset_update.guid = linked_asset_guid - asset_update.data_contract_latest = V8DataContract.ref_by_guid( - saved.guid - ) - asset_update.has_contract = True - asset_response = client.asset.save(asset_update) + # Build raw payload because v9 Asset doesn't model dataContractLatest + asset_payload = { + "entities": [ + { + "guid": linked_asset_guid, + "typeName": "IndistinctAsset", + "attributes": {"hasContract": True}, + "relationshipAttributes": { + "dataContractLatest": { + "guid": saved.guid, + "typeName": "DataContract", + } + }, + } + ] + } + raw_json = client._call_api(BULK_UPDATE, {}, asset_payload) + asset_response = _parse_mutation_response(raw_json) return contract_response, asset_response @staticmethod @@ -155,15 +164,28 @@ def delete( :param linked_asset_guid: GUID of the asset the contract was linked to :returns: tuple of (contract delete response, asset update response) """ - from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + from pyatlan.client.constants import BULK_UPDATE + from pyatlan_v9.client.asset import _parse_mutation_response delete_response = client.asset.purge_by_guid(contract_guid) - asset_update = IndistinctAsset() - asset_update.guid = linked_asset_guid - asset_update.has_contract = False - asset_update.data_contract_latest = None # type: ignore[assignment] - asset_response = client.asset.save(asset_update) + # Build raw payload because v9 Asset doesn't model dataContractLatest + asset_payload = { + "entities": [ + { + "guid": linked_asset_guid, + "typeName": "IndistinctAsset", + "attributes": { + "hasContract": False, + }, + "relationshipAttributes": { + "dataContractLatest": None, + }, + } + ] + } + raw_json = client._call_api(BULK_UPDATE, {}, asset_payload) + asset_response = _parse_mutation_response(raw_json) return delete_response, asset_response def to_json(self, nested: bool = True, serde: Serde | None = None) -> str: diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index 2300d494e..3c7aa1dc1 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -1,6 +1,6 @@ from json import dumps from typing import Union -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index d9d1eba70..213779dfd 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -9,7 +9,7 @@ import pytest -from pyatlan.model.response import AssetMutationResponse +from pyatlan_v9.model.response import AssetMutationResponse from pyatlan_v9.errors import InvalidRequestError from pyatlan_v9.model import DataContract from pyatlan_v9.model.contract import DataContractSpec @@ -162,7 +162,7 @@ def test_trim_to_required(): class TestSaveContract: def test_save_contract_sets_data_contract_latest_and_has_contract(self): - """save_contract should save the contract then update the linked asset.""" + """save should save the contract then update the linked asset via raw API.""" mock_client = MagicMock() contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, @@ -173,8 +173,8 @@ def test_save_contract_sets_data_contract_latest_and_has_contract(self): contract_response = MagicMock(spec=AssetMutationResponse) contract_response.assets_created.return_value = [saved_contract] - asset_response = MagicMock(spec=AssetMutationResponse) - mock_client.asset.save.side_effect = [contract_response, asset_response] + mock_client.asset.save.return_value = contract_response + mock_client._call_api.return_value = {"mutatedEntities": {}} result = DataContract.save( client=mock_client, @@ -182,24 +182,32 @@ def test_save_contract_sets_data_contract_latest_and_has_contract(self): linked_asset_guid="asset-guid-456", ) - assert result == (contract_response, asset_response) - assert mock_client.asset.save.call_count == 2 - # Verify the second save call updates the linked asset - asset_update = mock_client.asset.save.call_args_list[1][0][0] - assert asset_update.guid == "asset-guid-456" - assert asset_update.has_contract is True - assert asset_update.data_contract_latest is not None - assert asset_update.data_contract_latest.guid == "contract-guid-123" + assert result[0] == contract_response + mock_client.asset.save.assert_called_once() + # Verify the raw API call for the linked asset update + mock_client._call_api.assert_called_once() + call_args = mock_client._call_api.call_args + payload = call_args[0][2] + entity = payload["entities"][0] + assert entity["guid"] == "asset-guid-456" + assert entity["attributes"]["hasContract"] is True + assert ( + entity["relationshipAttributes"]["dataContractLatest"]["guid"] + == "contract-guid-123" + ) + assert ( + entity["relationshipAttributes"]["dataContractLatest"]["typeName"] + == "DataContract" + ) class TestDeleteContract: def test_delete_contract_purges_and_clears_has_contract(self): - """delete_contract should purge the contract and clear hasContract.""" + """delete should purge the contract and clear hasContract via raw API.""" mock_client = MagicMock() delete_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.purge_by_guid.return_value = delete_response - asset_response = MagicMock(spec=AssetMutationResponse) - mock_client.asset.save.return_value = asset_response + mock_client._call_api.return_value = {"mutatedEntities": {}} result = DataContract.delete( client=mock_client, @@ -207,10 +215,13 @@ def test_delete_contract_purges_and_clears_has_contract(self): linked_asset_guid="asset-guid-456", ) - assert result == (delete_response, asset_response) + assert result[0] == delete_response mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") - # Verify the asset update clears contract state - asset_update = mock_client.asset.save.call_args[0][0] - assert asset_update.guid == "asset-guid-456" - assert asset_update.has_contract is False - assert asset_update.data_contract_latest is None + # Verify the raw API call clears contract state + mock_client._call_api.assert_called_once() + call_args = mock_client._call_api.call_args + payload = call_args[0][2] + entity = payload["entities"][0] + assert entity["guid"] == "asset-guid-456" + assert entity["attributes"]["hasContract"] is False + assert entity["relationshipAttributes"]["dataContractLatest"] is None From 21d9102e49aaf1571e61330fcb2d2e52e3f6ecb8 Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Fri, 13 Mar 2026 17:01:34 +0530 Subject: [PATCH 4/9] refactor: set dataContractAssetLatest in creator, remove save()/delete() The creator() now accepts asset_type and sets data_contract_asset_latest via ref_by_qualified_name so the relationship is established on normal client.asset.save(). This removes the need for DataContract.save() and DataContract.delete() helper methods. Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 84 ++-------------- pyatlan_v9/model/assets/data_contract.py | 107 +++------------------ tests/integration/data_mesh_test.py | 20 +--- tests/unit/model/data_contract_test.py | 104 ++++---------------- tests_v9/integration/data_mesh_test.py | 15 +-- tests_v9/unit/model/data_contract_test.py | 90 ++++------------- 6 files changed, 66 insertions(+), 354 deletions(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index 2aa640542..e4a3c23e1 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/pyatlan/model/assets/core/data_contract.py @@ -5,7 +5,7 @@ from __future__ import annotations from json import JSONDecodeError, loads -from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Union, overload +from typing import ClassVar, List, Optional, Type, Union from pydantic.v1 import Field, validator @@ -21,30 +21,17 @@ 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: ... - - @overload - @classmethod - def creator( - cls, asset_qualified_name: str, contract_spec: Union[DataContractSpec, str] - ) -> DataContract: ... - @classmethod @init_guid def creator( cls, *, asset_qualified_name: str, + asset_type: Type[Asset], contract_json: Optional[str] = None, contract_spec: Optional[Union[DataContractSpec, str]] = None, ) -> DataContract: @@ -54,73 +41,12 @@ 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, - linked_asset_guid: str, - ) -> Tuple[AssetMutationResponse, AssetMutationResponse]: - """Save a DataContract and link it to the associated asset. - - After saving the contract, this method also sets ``dataContractLatest`` - and ``hasContract`` on the linked asset so the UI Contracts tab works - correctly. - - :param client: connectivity to an Atlan tenant - :param contract: DataContract to save (from ``DataContract.creator()``) - :param linked_asset_guid: GUID of the asset this contract is for - :returns: tuple of (contract save response, asset update response) - """ - from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset - - contract_response = client.asset.save(contract) - contracts = contract_response.assets_created( - DataContract - ) or contract_response.assets_updated(DataContract) - if not contracts: - return contract_response, contract_response - - saved = contracts[0] - asset_update = IndistinctAsset() - asset_update.guid = linked_asset_guid - asset_update.data_contract_latest = DataContract.ref_by_guid(saved.guid) - asset_update.has_contract = True - asset_response = client.asset.save(asset_update) - return contract_response, asset_response - - @staticmethod - def delete( - client: AtlanClient, - contract_guid: str, - linked_asset_guid: str, - ) -> Tuple[AssetMutationResponse, AssetMutationResponse]: - """Delete a DataContract and clean up the linked asset. - - This method purges (hard-deletes) the contract to avoid qualified-name - conflicts on re-creation, and clears ``hasContract`` and - ``dataContractLatest`` on the linked asset. - - :param client: connectivity to an Atlan tenant - :param contract_guid: GUID of the DataContract to delete - :param linked_asset_guid: GUID of the asset the contract was linked to - :returns: tuple of (contract delete response, asset update response) - """ - from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset - - delete_response = client.asset.purge_by_guid(contract_guid) - - asset_update = IndistinctAsset() - asset_update.guid = linked_asset_guid - asset_update.has_contract = False - asset_update.data_contract_latest = 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") @@ -327,6 +253,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: @@ -375,6 +302,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( diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index fcdd3930a..c82d83d74 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/pyatlan_v9/model/assets/data_contract.py @@ -7,7 +7,7 @@ import re from json import JSONDecodeError, loads -from typing import TYPE_CHECKING, Tuple, Union +from typing import Type, Union from msgspec import UNSET, UnsetType @@ -32,10 +32,6 @@ ) 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): @@ -73,20 +69,30 @@ 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, ) @classmethod @@ -99,95 +105,6 @@ def trim_to_required(self) -> "DataContract": """Return only required fields for update operations.""" return DataContract.updater(qualified_name=self.qualified_name, name=self.name) - @staticmethod - def save( - client: "AtlanClient", - contract: "DataContract", - linked_asset_guid: str, - ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": - """Save a DataContract and link it to the associated asset. - - After saving the contract, this method also sets ``dataContractLatest`` - and ``hasContract`` on the linked asset so the UI Contracts tab works - correctly. - - :param client: connectivity to an Atlan tenant - :param contract: DataContract to save (from ``DataContract.creator()``) - :param linked_asset_guid: GUID of the asset this contract is for - :returns: tuple of (contract save response, asset update response) - """ - from pyatlan.client.constants import BULK_UPDATE - from pyatlan_v9.client.asset import _parse_mutation_response - - contract_response = client.asset.save(contract) - contracts = contract_response.assets_created( - DataContract - ) or contract_response.assets_updated(DataContract) - if not contracts: - return contract_response, contract_response - - saved = contracts[0] - # Build raw payload because v9 Asset doesn't model dataContractLatest - asset_payload = { - "entities": [ - { - "guid": linked_asset_guid, - "typeName": "IndistinctAsset", - "attributes": {"hasContract": True}, - "relationshipAttributes": { - "dataContractLatest": { - "guid": saved.guid, - "typeName": "DataContract", - } - }, - } - ] - } - raw_json = client._call_api(BULK_UPDATE, {}, asset_payload) - asset_response = _parse_mutation_response(raw_json) - return contract_response, asset_response - - @staticmethod - def delete( - client: "AtlanClient", - contract_guid: str, - linked_asset_guid: str, - ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": - """Delete a DataContract and clean up the linked asset. - - This method purges (hard-deletes) the contract to avoid qualified-name - conflicts on re-creation, and clears ``hasContract`` and - ``dataContractLatest`` on the linked asset. - - :param client: connectivity to an Atlan tenant - :param contract_guid: GUID of the DataContract to delete - :param linked_asset_guid: GUID of the asset the contract was linked to - :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 - - delete_response = client.asset.purge_by_guid(contract_guid) - - # Build raw payload because v9 Asset doesn't model dataContractLatest - asset_payload = { - "entities": [ - { - "guid": linked_asset_guid, - "typeName": "IndistinctAsset", - "attributes": { - "hasContract": False, - }, - "relationshipAttributes": { - "dataContractLatest": None, - }, - } - ] - } - raw_json = client._call_api(BULK_UPDATE, {}, asset_payload) - asset_response = _parse_mutation_response(raw_json) - return delete_response, asset_response - def to_json(self, nested: bool = True, serde: Serde | None = None) -> str: """Convert to JSON string.""" if serde is None: diff --git a/tests/integration/data_mesh_test.py b/tests/integration/data_mesh_test.py index 624e2dc8f..121d790ed 100644 --- a/tests/integration/data_mesh_test.py +++ b/tests/integration/data_mesh_test.py @@ -305,13 +305,10 @@ def contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) - contract_response, _ = DataContract.save( - client=client, - contract=contract, - linked_asset_guid=table.guid, - ) + contract_response = client.asset.save(contract) result = contract_response.assets_created(asset_type=DataContract)[0] yield result delete_asset(client, guid=result.guid, asset_type=DataContract) @@ -514,16 +511,9 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): - """Verify DataContract.delete() purges the contract AND clears hasContract - on the linked asset — the cleanup that was missing when using - client.asset.purge_by_guid() alone (Bug 2 / BLDX-721). - """ + """Verify purge_by_guid deletes the contract.""" assert table.guid - delete_response, asset_response = DataContract.delete( - client=client, - contract_guid=contract.guid, - linked_asset_guid=table.guid, - ) + delete_response = client.asset.purge_by_guid(contract.guid) assert delete_response assert not delete_response.assets_created(asset_type=DataContract) assert not delete_response.assets_updated(asset_type=DataContract) @@ -534,8 +524,6 @@ def test_delete_contract(client: AtlanClient, table: Table, contract: DataContra assert deleted[0].qualified_name == contract.qualified_name assert deleted[0].delete_handler == "PURGE" assert deleted[0].status == EntityStatus.DELETED - # Verify the linked asset was cleaned up - assert asset_response @pytest.mark.order(after="test_retrieve_product") diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index 3c7aa1dc1..86c1bf158 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -1,13 +1,11 @@ from json import dumps from typing import Union -from unittest.mock import MagicMock import pytest from pyatlan.errors import InvalidRequestError -from pyatlan.model.assets import DataContract +from pyatlan.model.assets import DataContract, Table from pyatlan.model.contract import DataContractSpec -from pyatlan.model.response import AssetMutationResponse from tests.unit.model.constants import ( ASSET_QUALIFIED_NAME, DATA_CONTRACT_JSON, @@ -44,6 +42,7 @@ def test_creator_with_missing_parameters_raise_value_error( with pytest.raises(ValueError, match=message): DataContract.creator( # type: ignore asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, contract_spec=contract_spec, ) @@ -70,6 +69,7 @@ def test_creator_with_invalid_contract_json_raises_error( with pytest.raises(InvalidRequestError, match=error_msg): DataContract.creator( asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, ) @@ -77,6 +77,7 @@ def test_creator_with_invalid_contract_json_raises_error( def test_creator_atttributes_with_required_parameters(): attributes = DataContract.Attributes.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_json=dumps(DATA_CONTRACT_JSON), ) _assert_contract(attributes, is_json=True) @@ -85,6 +86,7 @@ def test_creator_atttributes_with_required_parameters(): def test_creator_with_required_parameters_json(): test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_json=dumps(DATA_CONTRACT_JSON), ) _assert_contract(test_contract) @@ -93,6 +95,7 @@ def test_creator_with_required_parameters_json(): def test_creator_with_required_parameters_spec_str(): test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR, ) _assert_contract(test_contract) @@ -101,6 +104,7 @@ def test_creator_with_required_parameters_spec_str(): def test_creator_with_required_parameters_spec_str_without_dataset(): test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR_WITHOUT_DATASET, ) # Ensure the default contract name is extracted from the table's qualified name (QN). @@ -111,11 +115,24 @@ def test_creator_with_required_parameters_spec_model(): spec = DataContractSpec.from_yaml(DATA_CONTRACT_SPEC_STR) test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=spec, ) _assert_contract(test_contract) +def test_creator_sets_data_contract_asset_latest(): + """Creator should set data_contract_asset_latest with the correct asset ref.""" + test_contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + ref = test_contract.data_contract_asset_latest + assert ref is not None + assert ref.unique_attributes == {"qualifiedName": ASSET_QUALIFIED_NAME} + + def test_create_for_modification(): test_contract = DataContract.create_for_modification( name=DATA_CONTRACT_NAME, @@ -130,84 +147,3 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) - - -class TestSaveContract: - def test_save_contract_sets_data_contract_latest_and_has_contract(self): - """save_contract should save the contract then update the linked asset.""" - mock_client = MagicMock() - contract = DataContract.creator( - asset_qualified_name=ASSET_QUALIFIED_NAME, - contract_json=dumps(DATA_CONTRACT_JSON), - ) - saved_contract = MagicMock() - saved_contract.guid = "contract-guid-123" - - contract_response = MagicMock(spec=AssetMutationResponse) - contract_response.assets_created.return_value = [saved_contract] - asset_response = MagicMock(spec=AssetMutationResponse) - mock_client.asset.save.side_effect = [contract_response, asset_response] - - result = DataContract.save( - client=mock_client, - contract=contract, - linked_asset_guid="asset-guid-456", - ) - - assert result == (contract_response, asset_response) - assert mock_client.asset.save.call_count == 2 - # Verify the second save call updates the linked asset - asset_update = mock_client.asset.save.call_args_list[1][0][0] - assert asset_update.guid == "asset-guid-456" - assert asset_update.has_contract is True - assert asset_update.data_contract_latest is not None - assert asset_update.data_contract_latest.guid == "contract-guid-123" - - def test_save_contract_handles_update(self): - """save_contract should work when contract is updated (not created).""" - mock_client = MagicMock() - contract = DataContract.creator( - asset_qualified_name=ASSET_QUALIFIED_NAME, - contract_json=dumps(DATA_CONTRACT_JSON), - ) - saved_contract = MagicMock() - saved_contract.guid = "contract-guid-789" - - contract_response = MagicMock(spec=AssetMutationResponse) - contract_response.assets_created.return_value = [] - contract_response.assets_updated.return_value = [saved_contract] - asset_response = MagicMock(spec=AssetMutationResponse) - mock_client.asset.save.side_effect = [contract_response, asset_response] - - result = DataContract.save( - client=mock_client, - contract=contract, - linked_asset_guid="asset-guid-456", - ) - - assert result == (contract_response, asset_response) - assert mock_client.asset.save.call_count == 2 - - -class TestDeleteContract: - def test_delete_contract_purges_and_clears_has_contract(self): - """delete_contract should purge the contract and clear hasContract.""" - mock_client = MagicMock() - delete_response = MagicMock(spec=AssetMutationResponse) - mock_client.asset.purge_by_guid.return_value = delete_response - asset_response = MagicMock(spec=AssetMutationResponse) - mock_client.asset.save.return_value = asset_response - - result = DataContract.delete( - client=mock_client, - contract_guid="contract-guid-123", - linked_asset_guid="asset-guid-456", - ) - - assert result == (delete_response, asset_response) - mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") - # Verify the asset update clears contract state - asset_update = mock_client.asset.save.call_args[0][0] - assert asset_update.guid == "asset-guid-456" - assert asset_update.has_contract is False - assert asset_update.data_contract_latest is None diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 3476d2544..3e2a67bff 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -308,13 +308,10 @@ def contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) - contract_response, _ = DataContract.save( - client=client, - contract=contract, - linked_asset_guid=table.guid, - ) + contract_response = client.asset.save(contract) result = contract_response.assets_created(asset_type=DataContract)[0] yield result delete_asset(client, guid=result.guid, asset_type=DataContract) @@ -524,11 +521,7 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): assert table.guid - delete_response, asset_response = DataContract.delete( - client=client, - contract_guid=contract.guid, - linked_asset_guid=table.guid, - ) + delete_response = client.asset.purge_by_guid(contract.guid) assert delete_response assert not delete_response.assets_created(asset_type=DataContract) assert not delete_response.assets_updated(asset_type=DataContract) @@ -539,8 +532,6 @@ def test_delete_contract(client: AtlanClient, table: Table, contract: DataContra assert deleted[0].qualified_name == contract.qualified_name assert deleted[0].delete_handler == "PURGE" assert deleted[0].status == EntityStatus.DELETED - # Verify the linked asset was cleaned up - assert asset_response @pytest.mark.order(after="test_retrieve_product") diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index 213779dfd..00bb1ce4b 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -5,13 +5,11 @@ from json import dumps from typing import Union -from unittest.mock import MagicMock import pytest -from pyatlan_v9.model.response import AssetMutationResponse from pyatlan_v9.errors import InvalidRequestError -from pyatlan_v9.model import DataContract +from pyatlan_v9.model import DataContract, Table from pyatlan_v9.model.contract import DataContractSpec from tests_v9.unit.model.constants import ( ASSET_QUALIFIED_NAME, @@ -50,6 +48,7 @@ def test_creator_with_missing_parameters_raise_value_error( with pytest.raises(ValueError, match=message): DataContract.creator( # type: ignore[arg-type] asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, contract_spec=contract_spec, ) @@ -77,6 +76,7 @@ def test_creator_with_invalid_contract_json_raises_error( with pytest.raises(InvalidRequestError, match=error_msg): DataContract.creator( asset_qualified_name=asset_qualified_name, + asset_type=Table, contract_json=contract_json, ) @@ -94,6 +94,7 @@ def test_creator_with_required_parameters_json(): """Test DataContract.creator for JSON payload.""" test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_json=dumps(DATA_CONTRACT_JSON), ) _assert_contract(test_contract) @@ -103,6 +104,7 @@ def test_creator_with_required_parameters_spec_str(): """Test DataContract.creator for YAML string payload.""" test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR, ) _assert_contract(test_contract) @@ -112,6 +114,7 @@ def test_creator_with_required_parameters_spec_str_without_dataset(): """Test creator defaults contract name from asset QN when dataset is absent.""" test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=DATA_CONTRACT_SPEC_STR_WITHOUT_DATASET, ) _assert_contract(test_contract, contract_name=DATA_CONTRACT_NAME_DEFAULT) @@ -122,11 +125,25 @@ def test_creator_with_required_parameters_spec_model(): spec = DataContractSpec.from_yaml(DATA_CONTRACT_SPEC_STR) test_contract = DataContract.creator( asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, contract_spec=spec, ) _assert_contract(test_contract) +def test_creator_sets_data_contract_asset_latest(): + """Creator should set data_contract_asset_latest with the correct asset ref.""" + test_contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + ref = test_contract.data_contract_asset_latest + assert ref is not None + assert ref.type_name == "Table" + assert ref.unique_attributes == {"qualifiedName": ASSET_QUALIFIED_NAME} + + @pytest.mark.parametrize( "qualified_name, name, message", [ @@ -158,70 +175,3 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) - - -class TestSaveContract: - def test_save_contract_sets_data_contract_latest_and_has_contract(self): - """save should save the contract then update the linked asset via raw API.""" - mock_client = MagicMock() - contract = DataContract.creator( - asset_qualified_name=ASSET_QUALIFIED_NAME, - contract_json=dumps(DATA_CONTRACT_JSON), - ) - saved_contract = MagicMock() - saved_contract.guid = "contract-guid-123" - - contract_response = MagicMock(spec=AssetMutationResponse) - contract_response.assets_created.return_value = [saved_contract] - mock_client.asset.save.return_value = contract_response - mock_client._call_api.return_value = {"mutatedEntities": {}} - - result = DataContract.save( - client=mock_client, - contract=contract, - linked_asset_guid="asset-guid-456", - ) - - assert result[0] == contract_response - mock_client.asset.save.assert_called_once() - # Verify the raw API call for the linked asset update - mock_client._call_api.assert_called_once() - call_args = mock_client._call_api.call_args - payload = call_args[0][2] - entity = payload["entities"][0] - assert entity["guid"] == "asset-guid-456" - assert entity["attributes"]["hasContract"] is True - assert ( - entity["relationshipAttributes"]["dataContractLatest"]["guid"] - == "contract-guid-123" - ) - assert ( - entity["relationshipAttributes"]["dataContractLatest"]["typeName"] - == "DataContract" - ) - - -class TestDeleteContract: - def test_delete_contract_purges_and_clears_has_contract(self): - """delete should purge the contract and clear hasContract via raw API.""" - mock_client = MagicMock() - delete_response = MagicMock(spec=AssetMutationResponse) - mock_client.asset.purge_by_guid.return_value = delete_response - mock_client._call_api.return_value = {"mutatedEntities": {}} - - result = DataContract.delete( - client=mock_client, - contract_guid="contract-guid-123", - linked_asset_guid="asset-guid-456", - ) - - assert result[0] == delete_response - mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") - # Verify the raw API call clears contract state - mock_client._call_api.assert_called_once() - call_args = mock_client._call_api.call_args - payload = call_args[0][2] - entity = payload["entities"][0] - assert entity["guid"] == "asset-guid-456" - assert entity["attributes"]["hasContract"] is False - assert entity["relationshipAttributes"]["dataContractLatest"] is None From 64d308a0548919d95d6857e2167e3cf369611c89 Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Fri, 13 Mar 2026 17:51:49 +0530 Subject: [PATCH 5/9] feat: restore save()/delete() on DataContract with cleanup logic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit save() delegates to client.asset.save() — the relationship is already set by creator(). delete() purges the contract and clears hasContract, dataContractLatest, and dataContractLatestCertified on the linked asset. Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 51 +++++++++++++++++- pyatlan_v9/model/assets/data_contract.py | 62 +++++++++++++++++++++- tests/integration/data_mesh_test.py | 11 ++-- tests/unit/model/data_contract_test.py | 44 +++++++++++++++ tests_v9/integration/data_mesh_test.py | 9 +++- tests_v9/unit/model/data_contract_test.py | 46 ++++++++++++++++ 6 files changed, 216 insertions(+), 7 deletions(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index e4a3c23e1..c2c67b11b 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/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, Type, Union +from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Type, Union from pydantic.v1 import Field, validator @@ -21,6 +21,10 @@ from .catalog import Catalog +if TYPE_CHECKING: + from pyatlan.client.atlan import AtlanClient + from pyatlan.model.response import AssetMutationResponse + class DataContract(Catalog): """Description""" @@ -47,6 +51,51 @@ def creator( ) 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, + linked_asset_guid: str, + ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": + """Delete (purge) a DataContract and clean up the linked asset. + + Uses hard-delete to avoid qualified-name conflicts on re-creation, + and clears ``hasContract``, ``dataContractLatest``, and + ``dataContractLatestCertified`` on the linked asset. + + :param client: connectivity to an Atlan tenant + :param contract_guid: GUID of the DataContract to delete + :param linked_asset_guid: GUID of the asset the contract was linked to + :returns: tuple of (contract delete response, asset update response) + """ + from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + + delete_response = client.asset.purge_by_guid(contract_guid) + + asset_update = IndistinctAsset() + asset_update.guid = linked_asset_guid + 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") diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index c82d83d74..ec3d2b6bb 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/pyatlan_v9/model/assets/data_contract.py @@ -7,7 +7,7 @@ import re from json import JSONDecodeError, loads -from typing import Type, 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): @@ -95,6 +99,62 @@ def creator( 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, + linked_asset_guid: str, + ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": + """Delete (purge) a DataContract and clean up the linked asset. + + Uses hard-delete to avoid qualified-name conflicts on re-creation, + and clears ``hasContract``, ``dataContractLatest``, and + ``dataContractLatestCertified`` on the linked asset. + + :param client: connectivity to an Atlan tenant + :param contract_guid: GUID of the DataContract to delete + :param linked_asset_guid: GUID of the asset the contract was linked to + :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 + + 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": "IndistinctAsset", + "attributes": {"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.""" diff --git a/tests/integration/data_mesh_test.py b/tests/integration/data_mesh_test.py index 121d790ed..baacbdf03 100644 --- a/tests/integration/data_mesh_test.py +++ b/tests/integration/data_mesh_test.py @@ -308,7 +308,7 @@ def contract( asset_type=Table, contract_json=dumps(contract_json), ) - contract_response = client.asset.save(contract) + contract_response = DataContract.save(client=client, contract=contract) result = contract_response.assets_created(asset_type=DataContract)[0] yield result delete_asset(client, guid=result.guid, asset_type=DataContract) @@ -511,10 +511,15 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): - """Verify purge_by_guid deletes the contract.""" + """Verify delete() purges the contract and cleans up the linked asset.""" assert table.guid - delete_response = client.asset.purge_by_guid(contract.guid) + delete_response, asset_response = DataContract.delete( + client=client, + contract_guid=contract.guid, + linked_asset_guid=table.guid, + ) assert delete_response + assert asset_response assert not delete_response.assets_created(asset_type=DataContract) assert not delete_response.assets_updated(asset_type=DataContract) deleted = delete_response.assets_deleted(asset_type=DataContract) diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index 86c1bf158..13492343e 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -1,11 +1,13 @@ from json import dumps from typing import Union +from unittest.mock import MagicMock import pytest from pyatlan.errors import InvalidRequestError from pyatlan.model.assets import DataContract, Table from pyatlan.model.contract import DataContractSpec +from pyatlan.model.response import AssetMutationResponse from tests.unit.model.constants import ( ASSET_QUALIFIED_NAME, DATA_CONTRACT_JSON, @@ -147,3 +149,45 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) + + +class TestSaveContract: + def test_save_delegates_to_asset_client(self): + """save() should delegate to client.asset.save().""" + mock_client = MagicMock() + contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + expected_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.return_value = expected_response + + result = DataContract.save(client=mock_client, contract=contract) + + assert result == expected_response + mock_client.asset.save.assert_called_once_with(contract) + + +class TestDeleteContract: + def test_delete_purges_and_clears_linked_asset(self): + """delete() should purge the contract and clear linked asset state.""" + mock_client = MagicMock() + delete_response = MagicMock(spec=AssetMutationResponse) + asset_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.purge_by_guid.return_value = delete_response + mock_client.asset.save.return_value = asset_response + + result = DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + linked_asset_guid="asset-guid-456", + ) + + assert result == (delete_response, asset_response) + mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") + asset_update = mock_client.asset.save.call_args[0][0] + assert asset_update.guid == "asset-guid-456" + assert asset_update.has_contract is False + assert asset_update.data_contract_latest is None + assert asset_update.data_contract_latest_certified is None diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 3e2a67bff..6b1471c1e 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -311,7 +311,7 @@ def contract( asset_type=Table, contract_json=dumps(contract_json), ) - contract_response = client.asset.save(contract) + contract_response = DataContract.save(client=client, contract=contract) result = contract_response.assets_created(asset_type=DataContract)[0] yield result delete_asset(client, guid=result.guid, asset_type=DataContract) @@ -521,8 +521,13 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): @pytest.mark.order(after="test_retrieve_contract") def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): assert table.guid - delete_response = client.asset.purge_by_guid(contract.guid) + delete_response, asset_response = DataContract.delete( + client=client, + contract_guid=contract.guid, + linked_asset_guid=table.guid, + ) assert delete_response + assert asset_response assert not delete_response.assets_created(asset_type=DataContract) assert not delete_response.assets_updated(asset_type=DataContract) deleted = delete_response.assets_deleted(asset_type=DataContract) diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index 00bb1ce4b..48a016795 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -5,12 +5,14 @@ from json import dumps from typing import Union +from unittest.mock import MagicMock import pytest from pyatlan_v9.errors import InvalidRequestError from pyatlan_v9.model import DataContract, Table from pyatlan_v9.model.contract import DataContractSpec +from pyatlan_v9.model.response import AssetMutationResponse from tests_v9.unit.model.constants import ( ASSET_QUALIFIED_NAME, DATA_CONTRACT_JSON, @@ -175,3 +177,47 @@ def test_trim_to_required(): qualified_name=DATA_CONTRACT_QUALIFIED_NAME, ).trim_to_required() _assert_contract(test_contract, False) + + +class TestSaveContract: + def test_save_delegates_to_asset_client(self): + """save() should delegate to client.asset.save().""" + mock_client = MagicMock() + contract = DataContract.creator( + asset_qualified_name=ASSET_QUALIFIED_NAME, + asset_type=Table, + contract_json=dumps(DATA_CONTRACT_JSON), + ) + expected_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.save.return_value = expected_response + + result = DataContract.save(client=mock_client, contract=contract) + + assert result == expected_response + mock_client.asset.save.assert_called_once_with(contract) + + +class TestDeleteContract: + def test_delete_purges_and_clears_linked_asset(self): + """delete() should purge the contract and clear linked asset state.""" + mock_client = MagicMock() + delete_response = MagicMock(spec=AssetMutationResponse) + mock_client.asset.purge_by_guid.return_value = delete_response + mock_client._call_api.return_value = {"mutatedEntities": {}} + + result = DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + linked_asset_guid="asset-guid-456", + ) + + assert result[0] == delete_response + mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") + # Verify the raw API call clears contract state + mock_client._call_api.assert_called_once() + payload = mock_client._call_api.call_args[0][2] + entity = payload["entities"][0] + assert entity["guid"] == "asset-guid-456" + assert entity["attributes"]["hasContract"] is False + assert entity["relationshipAttributes"]["dataContractLatest"] is None + assert entity["relationshipAttributes"]["dataContractLatestCertified"] is None From 560dcbd14588b894e30ea38ef3a160f2ade2ba23 Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Sun, 15 Mar 2026 16:40:37 +0530 Subject: [PATCH 6/9] refactor: remove linked_asset_guid from delete(), auto-retrieve linked asset delete() now retrieves the contract by GUID to find the linked asset automatically, eliminating the need for callers to pass linked_asset_guid. Also fixes IndistinctAsset type_name/qualifiedName/name for server compat. Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 25 ++++++++++++---- pyatlan_v9/model/assets/data_contract.py | 30 ++++++++++++++----- tests/integration/data_mesh_test.py | 2 +- tests/unit/model/data_contract_test.py | 34 ++++++++++++++++++++-- tests_v9/integration/data_mesh_test.py | 1 - tests_v9/unit/model/data_contract_test.py | 31 ++++++++++++++++++-- 6 files changed, 101 insertions(+), 22 deletions(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index c2c67b11b..1a745ec2d 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/pyatlan/model/assets/core/data_contract.py @@ -71,25 +71,38 @@ def save( def delete( client: "AtlanClient", contract_guid: str, - linked_asset_guid: str, ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": """Delete (purge) a DataContract and clean up the linked asset. - Uses hard-delete to avoid qualified-name conflicts on re-creation, - and clears ``hasContract``, ``dataContractLatest``, and - ``dataContractLatestCertified`` on 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 - :param linked_asset_guid: GUID of the asset the contract was linked to :returns: tuple of (contract delete response, asset update response) """ from pyatlan.model.assets.core.indistinct_asset import IndistinctAsset + contract = client.asset.get_by_guid( + contract_guid, + asset_type=DataContract, + ignore_relationships=False, + ) + 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.guid = linked_asset_guid + 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 or linked_asset.qualified_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] diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index ec3d2b6bb..d5b6b0624 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/pyatlan_v9/model/assets/data_contract.py @@ -119,31 +119,45 @@ def save( def delete( client: "AtlanClient", contract_guid: str, - linked_asset_guid: str, ) -> "Tuple[AssetMutationResponse, AssetMutationResponse]": """Delete (purge) a DataContract and clean up the linked asset. - Uses hard-delete to avoid qualified-name conflicts on re-creation, - and clears ``hasContract``, ``dataContractLatest``, and - ``dataContractLatestCertified`` on 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 - :param linked_asset_guid: GUID of the asset the contract was linked to :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 + contract = client.asset.get_by_guid( + contract_guid, + asset_type=DataContract, + ignore_relationships=False, + ) + 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": "IndistinctAsset", - "attributes": {"hasContract": False}, + "guid": linked_asset.guid, + "typeName": linked_asset.type_name, + "attributes": { + "qualifiedName": linked_asset.qualified_name, + "name": linked_asset.name or linked_asset.qualified_name, + "hasContract": False, + }, "relationshipAttributes": { "dataContractLatest": None, "dataContractLatestCertified": None, diff --git a/tests/integration/data_mesh_test.py b/tests/integration/data_mesh_test.py index baacbdf03..870d88e5c 100644 --- a/tests/integration/data_mesh_test.py +++ b/tests/integration/data_mesh_test.py @@ -331,6 +331,7 @@ def updated_contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) response = client.asset.save(contract) @@ -516,7 +517,6 @@ def test_delete_contract(client: AtlanClient, table: Table, contract: DataContra delete_response, asset_response = DataContract.delete( client=client, contract_guid=contract.guid, - linked_asset_guid=table.guid, ) assert delete_response assert asset_response diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index 13492343e..e8c755f0b 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -170,9 +170,20 @@ def test_save_delegates_to_asset_client(self): class TestDeleteContract: - def test_delete_purges_and_clears_linked_asset(self): - """delete() should purge the contract and clear linked asset state.""" + def test_delete_retrieves_contract_and_clears_linked_asset(self): + """delete() should retrieve the contract, find the linked asset, clear state, and purge.""" mock_client = MagicMock() + + # Mock the contract retrieved by get_by_guid + mock_linked_asset = MagicMock() + mock_linked_asset.type_name = "Table" + mock_linked_asset.guid = "asset-guid-456" + mock_linked_asset.qualified_name = "default/test/table-qn" + mock_linked_asset.name = "test-table" + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = mock_linked_asset + mock_client.asset.get_by_guid.return_value = mock_contract + delete_response = MagicMock(spec=AssetMutationResponse) asset_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.purge_by_guid.return_value = delete_response @@ -181,13 +192,30 @@ def test_delete_purges_and_clears_linked_asset(self): result = DataContract.delete( client=mock_client, contract_guid="contract-guid-123", - linked_asset_guid="asset-guid-456", ) assert result == (delete_response, asset_response) + mock_client.asset.get_by_guid.assert_called_once_with( + "contract-guid-123", + asset_type=DataContract, + ignore_relationships=False, + ) mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") asset_update = mock_client.asset.save.call_args[0][0] assert asset_update.guid == "asset-guid-456" assert asset_update.has_contract is False assert asset_update.data_contract_latest is None assert asset_update.data_contract_latest_certified is None + + def test_delete_raises_when_no_linked_asset(self): + """delete() should raise ValueError when contract has no linked asset.""" + mock_client = MagicMock() + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = None + mock_client.asset.get_by_guid.return_value = mock_contract + + with pytest.raises(ValueError, match="Cannot determine the linked asset"): + DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + ) diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 6b1471c1e..4ddc7bf14 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -524,7 +524,6 @@ def test_delete_contract(client: AtlanClient, table: Table, contract: DataContra delete_response, asset_response = DataContract.delete( client=client, contract_guid=contract.guid, - linked_asset_guid=table.guid, ) assert delete_response assert asset_response diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index 48a016795..3e62fd882 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -198,9 +198,17 @@ def test_save_delegates_to_asset_client(self): class TestDeleteContract: - def test_delete_purges_and_clears_linked_asset(self): - """delete() should purge the contract and clear linked asset state.""" + def test_delete_retrieves_contract_and_clears_linked_asset(self): + """delete() should retrieve the contract, find the linked asset, clear state, and purge.""" mock_client = MagicMock() + + # Mock the contract retrieved by get_by_guid + mock_linked_asset = MagicMock() + mock_linked_asset.guid = "asset-guid-456" + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = mock_linked_asset + mock_client.asset.get_by_guid.return_value = mock_contract + delete_response = MagicMock(spec=AssetMutationResponse) mock_client.asset.purge_by_guid.return_value = delete_response mock_client._call_api.return_value = {"mutatedEntities": {}} @@ -208,10 +216,14 @@ def test_delete_purges_and_clears_linked_asset(self): result = DataContract.delete( client=mock_client, contract_guid="contract-guid-123", - linked_asset_guid="asset-guid-456", ) assert result[0] == delete_response + mock_client.asset.get_by_guid.assert_called_once_with( + "contract-guid-123", + asset_type=DataContract, + ignore_relationships=False, + ) mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") # Verify the raw API call clears contract state mock_client._call_api.assert_called_once() @@ -221,3 +233,16 @@ def test_delete_purges_and_clears_linked_asset(self): assert entity["attributes"]["hasContract"] is False assert entity["relationshipAttributes"]["dataContractLatest"] is None assert entity["relationshipAttributes"]["dataContractLatestCertified"] is None + + def test_delete_raises_when_no_linked_asset(self): + """delete() should raise ValueError when contract has no linked asset.""" + mock_client = MagicMock() + mock_contract = MagicMock(spec=DataContract) + mock_contract.data_contract_asset_latest = None + mock_client.asset.get_by_guid.return_value = mock_contract + + with pytest.raises(ValueError, match="Cannot determine the linked asset"): + DataContract.delete( + client=mock_client, + contract_guid="contract-guid-123", + ) From 13b157e9a2f77f735364be9501fb48092ce4a94a Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Sun, 15 Mar 2026 16:54:36 +0530 Subject: [PATCH 7/9] fix: restore overload signatures for DataContract.creator() with asset_type Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index 1a745ec2d..00d1a1052 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/pyatlan/model/assets/core/data_contract.py @@ -5,7 +5,7 @@ from __future__ import annotations from json import JSONDecodeError, loads -from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Type, Union +from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Type, Union, overload from pydantic.v1 import Field, validator @@ -29,6 +29,26 @@ class DataContract(Catalog): """Description""" + @overload + @classmethod + def creator( + cls, + *, + asset_qualified_name: str, + asset_type: Type[Asset], + contract_json: str, + ) -> DataContract: ... + + @overload + @classmethod + def creator( + cls, + *, + asset_qualified_name: str, + asset_type: Type[Asset], + contract_spec: Union[DataContractSpec, str], + ) -> DataContract: ... + @classmethod @init_guid def creator( From 6924d16294ee2d178c97487c70e6d3e7ea0d3c3e Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Sun, 15 Mar 2026 16:55:19 +0530 Subject: [PATCH 8/9] fix: add missing asset_type param in v9 integration test updated_contract fixture Co-Authored-By: Claude Opus 4.6 --- tests_v9/integration/data_mesh_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 4ddc7bf14..0c5a5e0f5 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -334,6 +334,7 @@ def updated_contract( } contract = DataContract.creator( asset_qualified_name=table.qualified_name, + asset_type=Table, contract_json=dumps(contract_json), ) response = client.asset.save(contract) From 791ffff1d8499b4c6bbc4036b967c23f44c482af Mon Sep 17 00:00:00 2001 From: Aryamanz29 Date: Mon, 16 Mar 2026 16:17:04 +0530 Subject: [PATCH 9/9] fix: v9 DataContract type resolution and add data_contract_latest relationship - Add __post_init__ to DataContract to set type_name="DataContract" (was inheriting "Catalog" from parent, causing from_atlas_format to resolve as Catalog instead of DataContract) - Add data_contract_latest and data_contract_latest_certified relationship fields to Asset, AssetRelationshipAttributes, and _ASSET_REL_FIELDS - Add delete_handler field to Entity and ReferenceableNested with conversion support - Restore full data_contract_latest assertions in v9 integration tests Co-Authored-By: Claude Opus 4.6 --- pyatlan/model/assets/core/data_contract.py | 7 +++-- pyatlan_v9/model/assets/asset.py | 15 ++++++++++ pyatlan_v9/model/assets/data_contract.py | 14 +++++++-- pyatlan_v9/model/assets/entity.py | 3 ++ pyatlan_v9/model/assets/referenceable.py | 3 ++ tests/integration/data_mesh_test.py | 32 ++++++-------------- tests/unit/model/data_contract_test.py | 5 +++- tests_v9/integration/data_mesh_test.py | 35 +++++++--------------- tests_v9/unit/model/data_contract_test.py | 16 +++++++++- 9 files changed, 76 insertions(+), 54 deletions(-) diff --git a/pyatlan/model/assets/core/data_contract.py b/pyatlan/model/assets/core/data_contract.py index 00d1a1052..ae60e7b4a 100644 --- a/pyatlan/model/assets/core/data_contract.py +++ b/pyatlan/model/assets/core/data_contract.py @@ -102,12 +102,14 @@ def delete( :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, - ignore_relationships=False, + 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: @@ -117,12 +119,11 @@ def delete( ) 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 or 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] diff --git a/pyatlan_v9/model/assets/asset.py b/pyatlan_v9/model/assets/asset.py index aeb1fbe1b..5ce7a7690 100644 --- a/pyatlan_v9/model/assets/asset.py +++ b/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", ] diff --git a/pyatlan_v9/model/assets/data_contract.py b/pyatlan_v9/model/assets/data_contract.py index d5b6b0624..2a465d9b5 100644 --- a/pyatlan_v9/model/assets/data_contract.py +++ b/pyatlan_v9/model/assets/data_contract.py @@ -67,6 +67,9 @@ 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( @@ -132,11 +135,18 @@ def delete( """ 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, - ignore_relationships=False, + 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: @@ -155,7 +165,7 @@ def delete( "typeName": linked_asset.type_name, "attributes": { "qualifiedName": linked_asset.qualified_name, - "name": linked_asset.name or linked_asset.qualified_name, + "name": linked_asset.name, "hasContract": False, }, "relationshipAttributes": { diff --git a/pyatlan_v9/model/assets/entity.py b/pyatlan_v9/model/assets/entity.py index d7c5ba58f..201a44f1b 100644 --- a/pyatlan_v9/model/assets/entity.py +++ b/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.""" diff --git a/pyatlan_v9/model/assets/referenceable.py b/pyatlan_v9/model/assets/referenceable.py index 9949a9ed6..2413f5ded 100644 --- a/pyatlan_v9/model/assets/referenceable.py +++ b/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, diff --git a/tests/integration/data_mesh_test.py b/tests/integration/data_mesh_test.py index 870d88e5c..d615503f9 100644 --- a/tests/integration/data_mesh_test.py +++ b/tests/integration/data_mesh_test.py @@ -311,7 +311,7 @@ def contract( contract_response = DataContract.save(client=client, contract=contract) result = contract_response.assets_created(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + # Cleanup is handled by updated_contract fixture since it updates the same entity @pytest.fixture(scope="module") @@ -337,7 +337,14 @@ def updated_contract( response = client.asset.save(contract) result = response.assets_updated(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + delete_response, asset_response = DataContract.delete( + client=client, contract_guid=result.guid + ) + assert delete_response + assert asset_response + deleted = delete_response.assets_deleted(asset_type=DataContract) + assert deleted and len(deleted) == 1 + assert deleted[0].guid == result.guid def test_contract( @@ -510,27 +517,6 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): assert isinstance(asset_list, IndexSearchResults) -@pytest.mark.order(after="test_retrieve_contract") -def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): - """Verify delete() purges the contract and cleans up the linked asset.""" - assert table.guid - delete_response, asset_response = DataContract.delete( - client=client, - contract_guid=contract.guid, - ) - assert delete_response - assert asset_response - assert not delete_response.assets_created(asset_type=DataContract) - assert not delete_response.assets_updated(asset_type=DataContract) - deleted = delete_response.assets_deleted(asset_type=DataContract) - assert deleted - assert len(deleted) == 1 - assert deleted[0].guid == contract.guid - assert deleted[0].qualified_name == contract.qualified_name - assert deleted[0].delete_handler == "PURGE" - assert deleted[0].status == EntityStatus.DELETED - - @pytest.mark.order(after="test_retrieve_product") def test_delete_product(client: AtlanClient, product: DataProduct): response = client.asset.purge_by_guid(product.guid) diff --git a/tests/unit/model/data_contract_test.py b/tests/unit/model/data_contract_test.py index e8c755f0b..6d56ac412 100644 --- a/tests/unit/model/data_contract_test.py +++ b/tests/unit/model/data_contract_test.py @@ -195,10 +195,13 @@ def test_delete_retrieves_contract_and_clears_linked_asset(self): ) assert result == (delete_response, asset_response) + from pyatlan.model.assets import Asset + mock_client.asset.get_by_guid.assert_called_once_with( "contract-guid-123", asset_type=DataContract, - ignore_relationships=False, + attributes=[DataContract.DATA_CONTRACT_ASSET_LATEST], + related_attributes=[Asset.NAME, Asset.QUALIFIED_NAME, Asset.TYPE_NAME], ) mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") asset_update = mock_client.asset.save.call_args[0][0] diff --git a/tests_v9/integration/data_mesh_test.py b/tests_v9/integration/data_mesh_test.py index 0c5a5e0f5..dd9970517 100644 --- a/tests_v9/integration/data_mesh_test.py +++ b/tests_v9/integration/data_mesh_test.py @@ -314,7 +314,7 @@ def contract( contract_response = DataContract.save(client=client, contract=contract) result = contract_response.assets_created(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + # Cleanup is handled by updated_contract fixture since it updates the same entity @pytest.fixture(scope="module") @@ -340,7 +340,14 @@ def updated_contract( response = client.asset.save(contract) result = response.assets_updated(asset_type=DataContract)[0] yield result - delete_asset(client, guid=result.guid, asset_type=DataContract) + delete_response, asset_response = DataContract.delete( + client=client, contract_guid=result.guid + ) + assert delete_response + assert asset_response + deleted = delete_response.assets_deleted(asset_type=DataContract) + assert deleted and len(deleted) == 1 + assert deleted[0].guid == result.guid def test_contract( @@ -448,7 +455,7 @@ def test_update_product( assert len(products) == 1 assert products[ 0 - ].data_product_assets_d_s_l == DataProductsAssetsDSL.get_asset_selection(assets) + ].data_product_assets_dsl == DataProductsAssetsDSL.get_asset_selection(assets) # Test the product.updater() method without assets # (ensure asset selection remains unchanged) @@ -507,7 +514,7 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): product.guid, asset_type=DataProduct, ignore_relationships=False ) assert test_product - assert test_product.data_product_assets_d_s_l + assert test_product.data_product_assets_dsl asset_list = test_product.get_assets(client=client) assert asset_list.count and asset_list.count > 0 TOTAL_ASSETS = asset_list.count @@ -519,26 +526,6 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct): assert isinstance(asset_list, IndexSearchResults) -@pytest.mark.order(after="test_retrieve_contract") -def test_delete_contract(client: AtlanClient, table: Table, contract: DataContract): - assert table.guid - delete_response, asset_response = DataContract.delete( - client=client, - contract_guid=contract.guid, - ) - assert delete_response - assert asset_response - assert not delete_response.assets_created(asset_type=DataContract) - assert not delete_response.assets_updated(asset_type=DataContract) - deleted = delete_response.assets_deleted(asset_type=DataContract) - assert deleted - assert len(deleted) == 1 - assert deleted[0].guid == contract.guid - assert deleted[0].qualified_name == contract.qualified_name - assert deleted[0].delete_handler == "PURGE" - assert deleted[0].status == EntityStatus.DELETED - - @pytest.mark.order(after="test_retrieve_product") def test_delete_product(client: AtlanClient, product: DataProduct): response = client.asset.purge_by_guid(product.guid) diff --git a/tests_v9/unit/model/data_contract_test.py b/tests_v9/unit/model/data_contract_test.py index 3e62fd882..ed9024c43 100644 --- a/tests_v9/unit/model/data_contract_test.py +++ b/tests_v9/unit/model/data_contract_test.py @@ -204,7 +204,10 @@ def test_delete_retrieves_contract_and_clears_linked_asset(self): # Mock the contract retrieved by get_by_guid mock_linked_asset = MagicMock() + mock_linked_asset.type_name = "Table" mock_linked_asset.guid = "asset-guid-456" + mock_linked_asset.qualified_name = "default/test/table-qn" + mock_linked_asset.name = "test-table" mock_contract = MagicMock(spec=DataContract) mock_contract.data_contract_asset_latest = mock_linked_asset mock_client.asset.get_by_guid.return_value = mock_contract @@ -218,11 +221,19 @@ def test_delete_retrieves_contract_and_clears_linked_asset(self): contract_guid="contract-guid-123", ) + from pyatlan_v9.model.assets.asset import Asset + from pyatlan_v9.model.assets.referenceable import Referenceable + assert result[0] == delete_response mock_client.asset.get_by_guid.assert_called_once_with( "contract-guid-123", asset_type=DataContract, - ignore_relationships=False, + attributes=["dataContractAssetLatest"], + related_attributes=[ + Asset.NAME, + Referenceable.QUALIFIED_NAME, + Referenceable.TYPE_NAME, + ], ) mock_client.asset.purge_by_guid.assert_called_once_with("contract-guid-123") # Verify the raw API call clears contract state @@ -230,6 +241,9 @@ def test_delete_retrieves_contract_and_clears_linked_asset(self): payload = mock_client._call_api.call_args[0][2] entity = payload["entities"][0] assert entity["guid"] == "asset-guid-456" + assert entity["typeName"] == "Table" + assert entity["attributes"]["qualifiedName"] == "default/test/table-qn" + assert entity["attributes"]["name"] == "test-table" assert entity["attributes"]["hasContract"] is False assert entity["relationshipAttributes"]["dataContractLatest"] is None assert entity["relationshipAttributes"]["dataContractLatestCertified"] is None