improved error handling for get_assets() method

Author: vaibhavatlan

pyatlan/errors.py

@@ -629,6 +629,13 @@ class ErrorCode(Enum):
         "Set `test` to `True` or remove `username` and `password`.",
         InvalidRequestError,
     )
+    MISSING_DP_Asset_DSL = (
+        400,
+        "ATLAN-PYTHON-400-072",
+        "No DataProduct Aseet DSL was found.",
+        "You must provide a DataProduct asset DSL when retrieving DataProduct assets.",
+        NotFoundError,
+    )
     AUTHENTICATION_PASSTHROUGH = (
         401,
         "ATLAN-PYTHON-401-000",

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

@@ -91,6 +91,7 @@
 
         :param client: connectivity to an Atlan tenant (optional). If not provided, the default client will be used.
 
+        :raises NotFoundError: if DataProduct asset DSL cannot be found (does not exist) in Atlan
         :raises AtlanError: if there is an issue interacting with the API
         :returns: instance of `IndexSearchResults` with list of all assets linked to the provided data product
 
@@ -99,7 +100,10 @@
 
         client = AtlanClient.get_current_client() if not client else client
         dp_dsl = self.data_product_assets_d_s_l
-        json_object = json.loads(dp_dsl) if dp_dsl else {}
-        request = IndexSearchRequest(**json_object.get("query", {}))
-        response = client.asset.search(request)
+        if dp_dsl:
+            json_object = json.loads(dp_dsl) if dp_dsl else {}
+            request = IndexSearchRequest(**json_object.get("query", {}))
+            response = client.asset.search(request)
+        else:
+            raise ErrorCode.MISSING_DP_Asset_DSL.exception_with_parameters()
         return response

pyatlan/model/assets/core/data_product.py

@@ -11,6 +11,7 @@
 
 from pydantic.v1 import Field, StrictStr, validator
 
+from pyatlan.errors import ErrorCode
 from pyatlan.model.data_mesh import DataProductsAssetsDSL
 from pyatlan.model.enums import (
     DataProductCriticality,
@@ -129,6 +130,7 @@ def get_assets(self, client: Optional[AtlanClient] = None):
 
         :param client: connectivity to an Atlan tenant (optional). If not provided, the default client will be used.
 
+        :raises NotFoundError: if DataProduct asset DSL cannot be found (does not exist) in Atlan
         :raises AtlanError: if there is an issue interacting with the API
         :returns: instance of `IndexSearchResults` with list of all assets linked to the provided data product
 
@@ -137,9 +139,12 @@ def get_assets(self, client: Optional[AtlanClient] = None):
 
         client = AtlanClient.get_current_client() if not client else client
         dp_dsl = self.data_product_assets_d_s_l
-        json_object = json.loads(dp_dsl) if dp_dsl else {}
-        request = IndexSearchRequest(**json_object.get("query", {}))
-        response = client.asset.search(request)
+        if dp_dsl:
+            json_object = json.loads(dp_dsl) if dp_dsl else {}
+            request = IndexSearchRequest(**json_object.get("query", {}))
+            response = client.asset.search(request)
+        else:
+            raise ErrorCode.MISSING_DP_Asset_DSL.exception_with_parameters()
         return response
 
     type_name: str = Field(default="DataProduct", allow_mutation=False)

tests/integration/data_mesh_test.py

@@ -6,6 +6,7 @@
 
 from pyatlan.client.asset import IndexSearchResults
 from pyatlan.client.atlan import AtlanClient
+from pyatlan.errors import NotFoundError
 from pyatlan.model.assets import (
     Asset,
     AtlasGlossary,
@@ -502,6 +503,26 @@ def test_product_get_assets(client: AtlanClient, product: DataProduct):
     assert isinstance(asset_list, IndexSearchResults)
 
 
+@pytest.mark.order(after="test_retrieve_product")
+def test_dp_dsl_missing(client: AtlanClient, product: DataProduct):
+    assert product.guid
+    test_product = (
+        FluentSearch.select()
+        .where(Asset.GUID.eq(product.guid))
+        .where(Asset.TYPE_NAME.eq("DataProduct"))
+        .to_request()
+    )
+    retrieved_dp = client.asset.search(test_product).current_page()[0]
+    assert retrieved_dp
+    assert retrieved_dp.data_product_assets_d_s_l is None  # type: ignore
+
+    with pytest.raises(
+        NotFoundError,
+        match="No DataProduct Aseet DSL was found. Suggestion: You must provide a DataProduct asset DSL when retrieving DataProduct assets.",
+    ):
+        retrieved_dp.get_assets(client)  # type: ignore
+
+
 @pytest.mark.order(after="test_retrieve_contract")
 def test_delete_contract(client: AtlanClient, contract: DataContract):
     response = client.asset.purge_by_guid(contract.guid)

tests/unit/model/data_product_test.py

@@ -3,6 +3,8 @@
 
 import pytest
 
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.errors import NotFoundError
 from pyatlan.model.assets import AtlasGlossary, DataProduct
 from pyatlan.model.enums import CertificateStatus, DataProductStatus
 from pyatlan.model.fluent_search import CompoundQuery, FluentSearch
@@ -144,6 +146,19 @@ def test_create_for_modification():
     _assert_product(test_product)
 
 
+def test_get_assets_with_missing_dp_asset_dsl_raise_not_found_error():
+    data_product = DataProduct()
+    client = AtlanClient()
+    data_product.attributes.name = DATA_PRODUCT_NAME
+    data_product.parent_domain_qualified_name = DATA_PRODUCT_QUALIFIED_NAME
+    data_product.data_product_assets_d_s_l = None
+    with pytest.raises(
+        NotFoundError,
+        match="No DataProduct Aseet DSL was found. Suggestion: You must provide a DataProduct asset DSL when retrieving DataProduct assets.",
+    ):
+        data_product.get_assets(client)
+
+
 def test_trim_to_required():
     test_product = DataProduct.create_for_modification(
         qualified_name=DATA_PRODUCT_QUALIFIED_NAME,