FT-843: Implement GET operation that takes a list of attributes and relationships

pyatlan/client/asset.py

@@ -351,13 +351,15 @@ def find_purposes_by_name(
             allow_multiple=True,
         )
 
-    @validate_arguments
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def get_by_qualified_name(
         self,
         qualified_name: str,
         asset_type: Type[A],
         min_ext_info: bool = False,
-        ignore_relationships: bool = False,
+        ignore_relationships: bool = True,
+        attributes: Optional[Union[List[str], List[AtlanField]]] = None,
+        relationships_attributes: Optional[Union[List[str], List[AtlanField]]] = None,
     ) -> A:
         """
         Retrieves an asset by its qualified_name.
@@ -366,15 +368,48 @@ def get_by_qualified_name(
         :param asset_type: type of asset to be retrieved ( must be the actual asset type not a super type)
         :param min_ext_info: whether to minimize extra info (True) or not (False)
         :param ignore_relationships: whether to include relationships (False) or exclude them (True)
+        :param attributes: a specific list of attributes to retrieve for the asset
+        :param relationships_attributes: a specific list of relationships attributes to retrieve for the asset
         :returns: the requested asset
         :raises NotFoundError: if the asset does not exist
         :raises AtlanError: on any API communication issue
         """
+        from pyatlan.client.atlan import AtlanClient
+        from pyatlan.model.fluent_search import FluentSearch
+
         query_params = {
             "attr:qualifiedName": qualified_name,
             "minExtInfo": min_ext_info,
             "ignoreRelationships": ignore_relationships,
         }
+        attributes = attributes or []
+        relationships_attributes = relationships_attributes or []
+
+        if (attributes and len(attributes)) or (
+            relationships_attributes and len(relationships_attributes)
+        ):
+            client = AtlanClient.get_default_client()
+            search = (
+                FluentSearch().select().where(Asset.QUALIFIED_NAME.eq(qualified_name))
+            )
+            for attribute in attributes:
+                search = search.include_on_results(attribute)
+            for relation_attribute in relationships_attributes:
+                search = search.include_on_relations(relation_attribute)
+            results = search.execute(client=client)
+            if results and results.current_page():
+                first_result = results.current_page()[0]
+                if isinstance(first_result, asset_type):
+                    return first_result
+                else:
+                    raise ErrorCode.ASSET_NOT_FOUND_BY_NAME.exception_with_parameters(
+                        asset_type.__name__, qualified_name
+                    )
+            else:
+                raise ErrorCode.ASSET_NOT_FOUND_BY_QN.exception_with_parameters(
+                    qualified_name, asset_type.__name__
+                )
+
         raw_json = self._client._call_api(
             GET_ENTITY_BY_UNIQUE_ATTRIBUTE.format_path_with_params(asset_type.__name__),
             query_params,
@@ -390,13 +425,15 @@ def get_by_qualified_name(
             )
         return asset
 
-    @validate_arguments
+    @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def get_by_guid(
         self,
         guid: str,
         asset_type: Type[A] = Asset,  # type: ignore[assignment]
         min_ext_info: bool = False,
-        ignore_relationships: bool = False,
+        ignore_relationships: bool = True,
+        attributes: Optional[Union[List[str], List[AtlanField]]] = None,
+        relationships_attributes: Optional[Union[List[str], List[AtlanField]]] = None,
     ) -> A:
         """
         Retrieves an asset by its GUID.
@@ -405,14 +442,42 @@ def get_by_guid(
         :param asset_type: type of asset to be retrieved, defaults to `Asset`
         :param min_ext_info: whether to minimize extra info (True) or not (False)
         :param ignore_relationships: whether to include relationships (False) or exclude them (True)
+        :param attributes: a specific list of attributes to retrieve for the asset
+        :param relationships_attributes: a specific list of relationships attributes to retrieve for the asset
         :returns: the requested asset
         :raises NotFoundError: if the asset does not exist, or is not of the type requested
         :raises AtlanError: on any API communication issue
         """
+        from pyatlan.client.atlan import AtlanClient
+        from pyatlan.model.fluent_search import FluentSearch
+
         query_params = {
             "minExtInfo": min_ext_info,
             "ignoreRelationships": ignore_relationships,
         }
+        attributes = attributes or []
+        relationships_attributes = relationships_attributes or []
+
+        if (attributes and len(attributes)) or (
+            relationships_attributes and len(relationships_attributes)
+        ):
+            client = AtlanClient.get_default_client()
+            search = FluentSearch().select().where(Asset.GUID.eq(guid))
+            for attribute in attributes:
+                search = search.include_on_results(attribute)
+            for relation_attribute in relationships_attributes:
+                search = search.include_on_relations(relation_attribute)
+            results = search.execute(client=client)
+            if results and results.current_page():
+                first_result = results.current_page()[0]
+                if isinstance(first_result, asset_type):
+                    return first_result
+                else:
+                    raise ErrorCode.ASSET_NOT_TYPE_REQUESTED.exception_with_parameters(
+                        guid, asset_type.__name__
+                    )
+            else:
+                raise ErrorCode.ASSET_NOT_FOUND_BY_GUID.exception_with_parameters(guid)
 
         raw_json = self._client._call_api(
             GET_ENTITY_BY_GUID.format_path_with_params(guid),
@@ -743,7 +808,9 @@ def _restore(self, asset_type: Type[A], qualified_name: str, retries: int) -> bo
         if not asset_type.can_be_archived():
             return False
         existing = self.get_by_qualified_name(
-            asset_type=asset_type, qualified_name=qualified_name
+            asset_type=asset_type,
+            qualified_name=qualified_name,
+            ignore_relationships=False,
         )
         if not existing:
             # Nothing to restore, so cannot be restored
@@ -1308,10 +1375,14 @@ def replace_terms(
         if guid:
             if qualified_name:
                 raise ErrorCode.QN_OR_GUID_NOT_BOTH.exception_with_parameters()
-            asset = self.get_by_guid(guid=guid, asset_type=asset_type)
+            asset = self.get_by_guid(
+                guid=guid, asset_type=asset_type, ignore_relationships=False
+            )
         elif qualified_name:
             asset = self.get_by_qualified_name(
-                qualified_name=qualified_name, asset_type=asset_type
+                qualified_name=qualified_name,
+                asset_type=asset_type,
+                ignore_relationships=False,
             )
         else:
             raise ErrorCode.QN_OR_GUID.exception_with_parameters()
@@ -1346,10 +1417,14 @@ def remove_terms(
         if guid:
             if qualified_name:
                 raise ErrorCode.QN_OR_GUID_NOT_BOTH.exception_with_parameters()
-            asset = self.get_by_guid(guid=guid, asset_type=asset_type)
+            asset = self.get_by_guid(
+                guid=guid, asset_type=asset_type, ignore_relationships=False
+            )
         elif qualified_name:
             asset = self.get_by_qualified_name(
-                qualified_name=qualified_name, asset_type=asset_type
+                qualified_name=qualified_name,
+                asset_type=asset_type,
+                ignore_relationships=False,
             )
         else:
             raise ErrorCode.QN_OR_GUID.exception_with_parameters()

pyatlan/samples/events/lambda_scorer.py

@@ -209,7 +209,9 @@ def calculate_changes(self, asset: Asset) -> List[Asset]:
             s_readme = 0
             readme = asset.readme
             if readme and readme.guid:
-                readme = client.asset.get_by_guid(readme.guid, asset_type=Readme)
+                readme = client.asset.get_by_guid(
+                    readme.guid, asset_type=Readme, ignore_relationships=False
+                )
                 if description := readme.description:
                     if len(description) > 1000:
                         s_readme = 20

pyatlan/test_utils/__init__.py

@@ -112,7 +112,9 @@ def create_connection(
     )
     response = client.asset.save(to_create)
     result = response.assets_created(asset_type=Connection)[0]
-    return client.asset.get_by_guid(result.guid, asset_type=Connection)
+    return client.asset.get_by_guid(
+        result.guid, asset_type=Connection, ignore_relationships=False
+    )
 
 
 def create_group(client: AtlanClient, name: str) -> CreateGroupResponse:

tests/integration/adls_asset_test.py

@@ -249,7 +249,9 @@ def test_retrieve_adls_object(
     adls_container: ADLSContainer,
     adls_object: ADLSObject,
 ):
-    b = client.asset.get_by_guid(adls_object.guid, asset_type=ADLSObject)
+    b = client.asset.get_by_guid(
+        adls_object.guid, asset_type=ADLSObject, ignore_relationships=False
+    )
     assert b
     assert not b.is_incomplete
     assert b.guid == adls_object.guid
@@ -319,7 +321,9 @@ def test_read_deleted_adls_object(
     adls_container: ADLSContainer,
     adls_object: ADLSObject,
 ):
-    deleted = client.asset.get_by_guid(adls_object.guid, asset_type=ADLSObject)
+    deleted = client.asset.get_by_guid(
+        adls_object.guid, asset_type=ADLSObject, ignore_relationships=False
+    )
     assert deleted
     assert deleted.guid == adls_object.guid
     assert deleted.qualified_name == adls_object.qualified_name
@@ -339,7 +343,9 @@ def test_restore_object(
     )
     assert adls_object.qualified_name
     restored = client.asset.get_by_qualified_name(
-        asset_type=ADLSObject, qualified_name=adls_object.qualified_name
+        asset_type=ADLSObject,
+        qualified_name=adls_object.qualified_name,
+        ignore_relationships=False,
     )
     assert restored
     assert restored.guid == adls_object.guid

tests/integration/airflow_asset_test.py

@@ -170,7 +170,9 @@ def test_update_airflow_assets(
 
 
 def _retrieve_airflow_assets(client, asset, asset_type):
-    retrieved = client.asset.get_by_guid(asset.guid, asset_type=asset_type)
+    retrieved = client.asset.get_by_guid(
+        asset.guid, asset_type=asset_type, ignore_relationships=False
+    )
     assert retrieved
     assert not retrieved.is_incomplete
     assert retrieved.guid == asset.guid
@@ -215,7 +217,9 @@ def test_read_deleted_airflow_task(
     client: AtlanClient,
     airflow_task: AirflowTask,
 ):
-    deleted = client.asset.get_by_guid(airflow_task.guid, asset_type=AirflowTask)
+    deleted = client.asset.get_by_guid(
+        airflow_task.guid, asset_type=AirflowTask, ignore_relationships=False
+    )
     assert deleted
     assert deleted.status == EntityStatus.DELETED
     assert deleted.guid == airflow_task.guid
@@ -233,7 +237,9 @@ def test_restore_airflow_task(
     )
     assert airflow_task.qualified_name
     restored = client.asset.get_by_qualified_name(
-        asset_type=AirflowTask, qualified_name=airflow_task.qualified_name
+        asset_type=AirflowTask,
+        qualified_name=airflow_task.qualified_name,
+        ignore_relationships=False,
     )
     assert restored
     assert restored.guid == airflow_task.guid

tests/integration/api_asset_test.py

@@ -187,7 +187,9 @@ def test_update_api_path(
 def test_retrieve_api_path(
     client: AtlanClient, connection: Connection, api_spec: APISpec, api_path: APIPath
 ):
-    b = client.asset.get_by_guid(api_path.guid, asset_type=APIPath)
+    b = client.asset.get_by_guid(
+        api_path.guid, asset_type=APIPath, ignore_relationships=False
+    )
     assert b
     assert not b.is_incomplete
     assert b.guid == api_path.guid
@@ -250,7 +252,9 @@ def test_delete_api_path(
 def test_read_deleted_api_path(
     client: AtlanClient, connection: Connection, api_spec: APISpec, api_path: APIPath
 ):
-    deleted = client.asset.get_by_guid(api_path.guid, asset_type=APIPath)
+    deleted = client.asset.get_by_guid(
+        api_path.guid, asset_type=APIPath, ignore_relationships=False
+    )
     assert deleted
     assert deleted.guid == api_path.guid
     assert deleted.qualified_name == api_path.qualified_name
@@ -267,7 +271,9 @@ def test_restore_path(
     )
     assert api_path.qualified_name
     restored = client.asset.get_by_qualified_name(
-        asset_type=APIPath, qualified_name=api_path.qualified_name
+        asset_type=APIPath,
+        qualified_name=api_path.qualified_name,
+        ignore_relationships=False,
     )
     assert restored
     assert restored.guid == api_path.guid
@@ -365,7 +371,9 @@ def test_update_api_object(
 def test_retrieve_api_object(
     client: AtlanClient, connection: Connection, api_object_overload: APIObject
 ):
-    b = client.asset.get_by_guid(api_object_overload.guid, asset_type=APIObject)
+    b = client.asset.get_by_guid(
+        api_object_overload.guid, asset_type=APIObject, ignore_relationships=False
+    )
     assert b
     assert not b.is_incomplete
     assert b.guid == api_object_overload.guid
@@ -398,7 +406,9 @@ def test_delete_api_object(
 def test_read_deleted_api_object(
     client: AtlanClient, connection: Connection, api_object_overload: APIObject
 ):
-    deleted = client.asset.get_by_guid(api_object_overload.guid, asset_type=APIObject)
+    deleted = client.asset.get_by_guid(
+        api_object_overload.guid, asset_type=APIObject, ignore_relationships=False
+    )
     assert deleted
     assert deleted.guid == api_object_overload.guid
     assert deleted.qualified_name == api_object_overload.qualified_name
@@ -415,7 +425,9 @@ def test_restore_object(
     )
     assert api_object_overload.qualified_name
     restored = client.asset.get_by_qualified_name(
-        asset_type=APIObject, qualified_name=api_object_overload.qualified_name
+        asset_type=APIObject,
+        qualified_name=api_object_overload.qualified_name,
+        ignore_relationships=False,
     )
     assert restored
     assert restored.guid == api_object_overload.guid
@@ -597,7 +609,9 @@ def test_update_api_query(
 def test_retrieve_api_query(
     client: AtlanClient, connection: Connection, api_query_overload_3: APIQuery
 ):
-    b = client.asset.get_by_guid(api_query_overload_3.guid, asset_type=APIQuery)
+    b = client.asset.get_by_guid(
+        api_query_overload_3.guid, asset_type=APIQuery, ignore_relationships=False
+    )
     assert b
     assert not b.is_incomplete
     assert b.guid == api_query_overload_3.guid
@@ -630,7 +644,9 @@ def test_delete_api_query(
 def test_read_deleted_api_query(
     client: AtlanClient, connection: Connection, api_query_overload_3: APIQuery
 ):
-    deleted = client.asset.get_by_guid(api_query_overload_3.guid, asset_type=APIQuery)
+    deleted = client.asset.get_by_guid(
+        api_query_overload_3.guid, asset_type=APIQuery, ignore_relationships=False
+    )
     assert deleted
     assert deleted.guid == api_query_overload_3.guid
     assert deleted.qualified_name == api_query_overload_3.qualified_name
@@ -647,7 +663,9 @@ def test_restore_query(
     )
     assert api_query_overload_3.qualified_name
     restored = client.asset.get_by_qualified_name(
-        asset_type=APIQuery, qualified_name=api_query_overload_3.qualified_name
+        asset_type=APIQuery,
+        qualified_name=api_query_overload_3.qualified_name,
+        ignore_relationships=False,
     )
     assert restored
     assert restored.guid == api_query_overload_3.guid
@@ -989,7 +1007,9 @@ def test_retrieve_api_field(
     api_field_parent_query_overload: APIField,
 ):
     b = client.asset.get_by_guid(
-        api_field_parent_query_overload.guid, asset_type=APIField
+        api_field_parent_query_overload.guid,
+        asset_type=APIField,
+        ignore_relationships=False,
     )
     assert b
     assert not b.is_incomplete
@@ -1031,7 +1051,9 @@ def test_read_deleted_api_field(
     api_field_parent_query_overload: APIField,
 ):
     deleted = client.asset.get_by_guid(
-        api_field_parent_query_overload.guid, asset_type=APIField
+        api_field_parent_query_overload.guid,
+        asset_type=APIField,
+        ignore_relationships=False,
     )
     assert deleted
     assert deleted.guid == api_field_parent_query_overload.guid
@@ -1054,6 +1076,7 @@ def test_restore_field(
     restored = client.asset.get_by_qualified_name(
         asset_type=APIField,
         qualified_name=api_field_parent_query_overload.qualified_name,
+        ignore_relationships=False,
     )
     assert restored
     assert restored.guid == api_field_parent_query_overload.guid