Merge pull request #251 from atlanhq/DVX-267

DVX-267: Fixes lineage response not including custom metadata

pyatlan/client/asset.py

@@ -235,6 +235,10 @@ def get_lineage_list(
         )
         if "entities" in raw_json:
             try:
+                for entity in raw_json["entities"]:
+                    unflatten_custom_metadata_for_entity(
+                        entity=entity, attributes=lineage_request.attributes
+                    )
                 assets = parse_obj_as(List[Asset], raw_json["entities"])
                 has_more = parse_obj_as(bool, raw_json["hasMore"])
             except ValidationError as err:

pyatlan/model/lineage.py

@@ -335,7 +335,9 @@ def __init__(
         size: StrictInt = 10,
         exclude_meanings: StrictBool = True,
         exclude_atlan_tags: StrictBool = True,
-        includes_on_results: Optional[Union[List[AtlanField], AtlanField]] = None,
+        includes_on_results: Optional[
+            Union[List[str], str, List[AtlanField], AtlanField]
+        ] = None,
         includes_in_results: Optional[Union[List[LineageFilter], LineageFilter]] = None,
         includes_condition: FilterList.Condition = FilterList.Condition.AND,
         where_assets: Optional[Union[List[LineageFilter], LineageFilter]] = None,
@@ -370,8 +372,9 @@ def __init__(
         self._direction: LineageDirection = direction
         self._exclude_atlan_tags: bool = exclude_atlan_tags
         self._exclude_meanings: bool = exclude_meanings
-
-        self._includes_on_results: List[AtlanField] = self._to_list(includes_on_results)
+        self._includes_on_results: List[Union[str, AtlanField]] = self._to_list(
+            includes_on_results
+        )
         self._includes_in_results: List[LineageFilter] = self._to_list(
             includes_in_results
         )
@@ -443,11 +446,11 @@ def exclude_meanings(self, exclude_meanings: StrictBool) -> "FluentLineage":
         clone._exclude_meanings = exclude_meanings
         return clone
 
-    def include_on_results(self, field: AtlanField) -> "FluentLineage":
+    def include_on_results(self, field: Union[str, AtlanField]) -> "FluentLineage":
         """Adds the include_on_results to traverse the lineage.
         :param field: attributes to retrieve for each asset in the lineage results
         :returns: the FluentLineage with this include_on_results criterion added"""
-        validate_type(name="field", _type=AtlanField, value=field)
+        validate_type(name="field", _type=(str, AtlanField), value=field)
         clone = self._clone()
         clone._includes_on_results.append(field)
         return clone
@@ -567,7 +570,8 @@ def request(self) -> LineageListRequest:
             request.entity_filters = FilterList(condition=self._includes_condition, criteria=criteria)  # type: ignore
         if self._includes_on_results:
             request.attributes = [
-                field.atlan_field_name for field in self._includes_on_results
+                field.atlan_field_name if isinstance(field, AtlanField) else field
+                for field in self._includes_on_results
             ]
         if self._size:
             request.size = self._size

pyatlan/utils.py

@@ -299,24 +299,32 @@ def is_comparable_type(attribute_type: str, to: ComparisonCategory) -> bool:
     return to == ComparisonCategory.STRING
 
 
-def validate_type(name: str, _type: type, value):
+def validate_type(name: str, _type, value):
     """
     Validate that the given value is of the specified type.
 
-    :param name: the name of the variable to be used in error message
-    :_type: the type of the variable to be validated
-    :value: the value to be validated that it is of the specified type
+    :param name: the name of the variable to be used in the error message
+    :param _type: the type of the variable to be validated
+    :param value: the value to be validated that it is of the specified type
 
     """
     if _type is int:
         if isinstance(value, _type) and not isinstance(value, bool):
             return
+    elif isinstance(_type, tuple):
+        if any(isinstance(value, t) for t in _type):
+            return
     elif isinstance(value, _type):
         return
-    raise ErrorCode.INVALID_PARAMETER_TYPE.exception_with_parameters(
-        name, _type.__name__
+
+    type_name = (
+        ", ".join(t.__name__ for t in _type)
+        if isinstance(_type, tuple)
+        else _type.__name__
     )
 
+    raise ErrorCode.INVALID_PARAMETER_TYPE.exception_with_parameters(name, type_name)
+
 
 def move_struct(data):
     struct_names = {

setup.py

@@ -55,6 +55,7 @@ def read(file_name):
     url="https://github.com/atlanhq/atlan-python",
     license="Apache LICENSE 2.0",
     classifiers=[
+        "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
         "License :: OSI Approved :: Apache Software License",
         "Operating System :: OS Independent",

tests/unit/data/lineage_responses/lineage_list.json

@@ -0,0 +1,38 @@
+{
+    "entities": [
+      {
+        "typeName": "View",
+        "attributes": {
+          "qualifiedName": "test-qn",
+          "testcm1.testcm2": "test-cm-value",
+          "name": "test-name"
+        },
+        "guid": "test-guid",
+        "status": "ACTIVE",
+        "displayText": "test-display-text",
+        "classificationNames": [],
+        "meaningNames": [],
+        "meanings": [],
+        "isIncomplete": false,
+        "labels": [],
+        "createdBy": "service-account-apikey-1234",
+        "updatedBy": "service-account-apikey-1234",
+        "createTime": 1708145336284,
+        "updateTime": 1708146440004
+      }
+    ],
+    "hasMore": false,
+    "entityCount": 1,
+    "searchParameters": {
+      "guid": "test-guid",
+      "size": 10,
+      "from": 0,
+      "depth": 1,
+      "direction": "INPUT",
+      "attributes": [
+        "testcm1.testcm2"
+      ],
+      "excludeMeanings": false,
+      "excludeClassifications": false
+    }
+}

tests/unit/test_client.py

@@ -9,6 +9,7 @@
 
 from pyatlan.client.asset import AssetClient, Batch, CustomMetadataHandling
 from pyatlan.client.atlan import AtlanClient
+from pyatlan.client.common import ApiCaller
 from pyatlan.client.search_log import SearchLogClient
 from pyatlan.errors import AtlanError, ErrorCode, InvalidRequestError, NotFoundError
 from pyatlan.model.assets import (
@@ -18,7 +19,13 @@
     Table,
 )
 from pyatlan.model.core import Announcement, BulkRequest
-from pyatlan.model.enums import AnnouncementType, CertificateStatus, SaveSemantic
+from pyatlan.model.enums import (
+    AnnouncementType,
+    CertificateStatus,
+    LineageDirection,
+    SaveSemantic,
+)
+from pyatlan.model.lineage import LineageListRequest
 from pyatlan.model.response import AssetMutationResponse
 from pyatlan.model.search import Bool, Term
 from pyatlan.model.search_log import SearchLogRequest
@@ -55,10 +62,14 @@
 LOG_USERNAME = "userName"
 SEARCH_PARAMS = "searchParameters"
 SEARCH_COUNT = "approximateCount"
-SEARCH_RESPONSES_DIR = Path(__file__).parent / "data" / "search_log_responses"
+TEST_DATA_DIR = Path(__file__).parent / "data"
+SEARCH_RESPONSES_DIR = TEST_DATA_DIR / "search_log_responses"
 SL_MOST_RECENT_VIEWERS_JSON = "sl_most_recent_viewers.json"
 SL_MOST_VIEWED_ASSETS_JSON = "sl_most_viewed_assets.json"
 SL_DETAILED_LOG_ENTRIES_JSON = "sl_detailed_log_entries.json"
+CM_NAME = "testcm1.testcm2"
+LINEAGE_LIST_JSON = "lineage_list.json"
+LINEAGE_RESPONSES_DIR = TEST_DATA_DIR / "lineage_responses"
 TEST_ANNOUNCEMENT = Announcement(
     announcement_title="test-title",
     announcement_message="test-msg",
@@ -86,24 +97,40 @@ def mock_asset_client():
     return Mock(AssetClient)
 
 
-def load_json(filename):
-    with (SEARCH_RESPONSES_DIR / filename).open() as input_file:
+@pytest.fixture(scope="module")
+def mock_api_caller():
+    return Mock(spec=ApiCaller)
+
+
+@pytest.fixture()
+def mock_cm_cache():
+    with patch("pyatlan.model.custom_metadata.CustomMetadataCache") as cache:
+        yield cache
+
+
+def load_json(respones_dir, filename):
+    with (respones_dir / filename).open() as input_file:
         return load(input_file)
 
 
 @pytest.fixture()
 def sl_most_recent_viewers_json():
-    return load_json(SL_MOST_RECENT_VIEWERS_JSON)
+    return load_json(SEARCH_RESPONSES_DIR, SL_MOST_RECENT_VIEWERS_JSON)
 
 
 @pytest.fixture()
 def sl_most_viewed_assets_json():
-    return load_json(SL_MOST_VIEWED_ASSETS_JSON)
+    return load_json(SEARCH_RESPONSES_DIR, SL_MOST_VIEWED_ASSETS_JSON)
 
 
 @pytest.fixture()
 def sl_detailed_log_entries_json():
-    return load_json(SL_DETAILED_LOG_ENTRIES_JSON)
+    return load_json(SEARCH_RESPONSES_DIR, SL_DETAILED_LOG_ENTRIES_JSON)
+
+
+@pytest.fixture()
+def lineage_list_json():
+    return load_json(LINEAGE_RESPONSES_DIR, LINEAGE_LIST_JSON)
 
 
 @pytest.mark.parametrize(
@@ -1077,6 +1104,29 @@ def test_search_log_views_by_guid(mock_sl_api_call, sl_detailed_log_entries_json
     assert log_entries[0].request_relation_attributes
 
 
+def test_asset_get_lineage_list_response_with_custom_metadata(
+    mock_api_caller, mock_cm_cache, lineage_list_json
+):
+    client = AssetClient(mock_api_caller)
+    mock_cm_cache.get_name_for_id.return_value = CM_NAME
+    mock_api_caller._call_api.return_value = lineage_list_json
+
+    lineage_request = LineageListRequest(
+        guid="test-guid", depth=1, direction=LineageDirection.UPSTREAM
+    )
+    lineage_request.attributes = [CM_NAME]
+    lineage_response = client.get_lineage_list(lineage_request=lineage_request)
+    asset = next(iter(lineage_response))
+
+    assert asset
+    assert asset.type_name == "View"
+    assert asset.guid == "test-guid"
+    assert asset.qualified_name == "test-qn"
+    assert asset.attributes
+    assert asset.business_attributes
+    assert asset.business_attributes == {"testcm1": {"testcm2": "test-cm-value"}}
+
+
 @pytest.mark.parametrize(
     "test_method, test_kwargs, test_asset_types",
     [

tests/unit/test_lineage.py

@@ -37,22 +37,21 @@
 )
 from pyatlan.model.typedef import AttributeDef
 
-BASE_GUID_TARGET = "e44ed3a2-1de5-4f23-b3f1-6e005156fee9"
-
-DATA_DIR = Path(__file__).parent / "data"
-BASE_GUID = "75474eab-3105-4ef9-9f84-709e386a7d3e"
 TODAY = date.today()
+BASE_GUID = "75474eab-3105-4ef9-9f84-709e386a7d3e"
+BASE_GUID_TARGET = "e44ed3a2-1de5-4f23-b3f1-6e005156fee9"
+LINEAGE_RESPONSES_DIR = Path(__file__).parent / "data" / "lineage_responses"
 
 
 @pytest.fixture(scope="session")
-def lineage_response_json():
-    with (DATA_DIR / "lineage_response.json").open() as input_file:
+def lineage_json():
+    with (LINEAGE_RESPONSES_DIR / "lineage.json").open() as input_file:
         return json.load(input_file)
 
 
 @pytest.fixture(scope="session")
-def lineage_response(lineage_response_json):
-    return LineageResponse(**lineage_response_json)
+def lineage_response(lineage_json):
+    return LineageResponse(**lineage_json)
 
 
 @pytest.fixture(scope="session")
@@ -737,7 +736,7 @@ def sut(self) -> FluentLineage:
                 BAD_LINEAGE_FILTER_LIST,
                 GOOD_WHERE_ASSETS,
                 GOOD_WHERE_RELATIONSHIPS,
-                r"2 validation errors for Init\nincludes_on_results",
+                r"4 validation errors for Init\nincludes_on_results",
             ),
             (
                 GOOD_GUID,
@@ -826,7 +825,7 @@ def test_request_with_defaults(self, sut: FluentLineage):
             ),
         ],
     )
-    def test_reqeust(
+    def test_request(
         self,
         starting_guid,
         depth,
@@ -957,7 +956,7 @@ def validate_filter(filter_, filter_condition, results):
             (
                 "include_on_results",
                 1,
-                r"ATLAN-PYTHON-400-048 Invalid parameter type for field should be AtlanField",
+                r"ATLAN-PYTHON-400-048 Invalid parameter type for field should be str, AtlanField",
             ),
             (
                 "include_in_results",