Add new method create_hierarchy() to AssetsAPI (#494)

* Add create_hierarchy method to AssetsAPI and update tests. * Update documentation. * Update CHANGELOG.md * Remove redundant code, add check that all assets have external_id, and update tests. * Update documentation. * Update integration test after merge. * Update docstrings. * Update CHANGELOG.md * Update assets.py * Update documentation.
cognitedata · Aug 15, 2019 · 3313610 · 3313610
1 parent 5ec1833
commit 3313610
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 27 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,10 +21,14 @@ Changes are grouped as follows
 
 ## [Unreleased]
 ### Added
+- New method create_hierarchy() added to assets API.
 - SequencesAPI.list now accepts an asset_ids parameter for searching by asset
 - SequencesDataAPI.insert now accepts a SequenceData object for easier copying
 - DatapointsAPI.insert now accepts a Datapoints object for easier copying
 
+### Changed
+- assets.create() no longer validates asset hierarchy and sorts assets before posting. This functionality has been moved to assets.create_hierarchy().
+
 ## [1.0.5] - 2019-08-15
 ### Added
 - files.create() method to enable creating a file without uploading content.

diff --git a/cognite/client/_api/assets.py b/cognite/client/_api/assets.py
@@ -194,7 +194,7 @@ def list(
         return self._list(method="POST", limit=limit, filter=filter)
 
     def create(self, asset: Union[Asset, List[Asset]]) -> Union[Asset, AssetList]:
-        """Create one or more assets.
+        """Create one or more assets. You can create an arbitrary number of assets, and the SDK will split the request into multiple requests.
 
         Args:
             asset (Union[Asset, List[Asset]]): Asset or list of assets to create.
@@ -213,9 +213,32 @@ def create(self, asset: Union[Asset, List[Asset]]) -> Union[Asset, AssetList]:
                 >>> res = c.assets.create(assets)
         """
         utils._auxiliary.assert_type(asset, "asset", [Asset, list])
-        if isinstance(asset, Asset) or len(asset) <= self._CREATE_LIMIT:
-            return self._create_multiple(asset)
-        return _AssetPoster(asset, client=self).post()
+        return self._create_multiple(asset)
+
+    def create_hierarchy(self, assets: List[Asset]) -> AssetList:
+        """Create asset hierarchy. Like the create() method, when posting a large number of assets, the IDE will split the request into smaller requests.
+        However, create_hierarchy() will additionally make sure that the assets are posted in correct order. The ordering is determined from the
+        external_id and parent_external_id properties of the assets, and the external_id is therefore required for all assets. Before posting, it is
+        checked that all assets have a unique external_id and that there are no circular dependencies.
+
+        Args:
+            assets (List[Asset]]): List of assets to create. Requires each asset to have a unique external id.
+
+        Returns:
+            AssetList: Created asset hierarchy
+
+        Examples:
+
+            Create asset hierarchy::
+
+                >>> from cognite.client import CogniteClient
+                >>> from cognite.client.data_classes import Asset
+                >>> c = CogniteClient()
+                >>> assets = [Asset(external_id="root"), Asset(external_id="child1", parent_external_id="root"), Asset(external_id="child2", parent_external_id="root")]
+                >>> res = c.assets.create_hierarchy(assets)
+        """
+        utils._auxiliary.assert_type(assets, "assets", [list])
+        return _AssetPoster(assets, client=self).post()
 
     def delete(
         self, id: Union[int, List[int]] = None, external_id: Union[str, List[str]] = None, recursive: bool = False
@@ -374,14 +397,9 @@ def __init__(self, assets: List[Asset], client: AssetsAPI):
         self.external_id_to_asset = {}
 
         for asset in assets:
-            asset_copy = Asset(**asset.dump())
-            external_id = asset.external_id
-            if external_id is None:
-                external_id = utils._auxiliary.random_string()
-                asset_copy.external_id = external_id
-            self.remaining_external_ids[external_id] = None
-            self.remaining_external_ids_set.add(external_id)
-            self.external_id_to_asset[external_id] = asset_copy
+            self.remaining_external_ids[asset.external_id] = None
+            self.remaining_external_ids_set.add(asset.external_id)
+            self.external_id_to_asset[asset.external_id] = asset
 
         self.client = client
 
@@ -409,10 +427,11 @@ def __init__(self, assets: List[Asset], client: AssetsAPI):
     def _validate_asset_hierarchy(assets) -> None:
         external_ids_seen = set()
         for asset in assets:
-            if asset.external_id:
-                if asset.external_id in external_ids_seen:
-                    raise AssertionError("Duplicate external_id '{}' found".format(asset.external_id))
-                external_ids_seen.add(asset.external_id)
+            if asset.external_id is None:
+                raise AssertionError("An asset does not have external_id.")
+            if asset.external_id in external_ids_seen:
+                raise AssertionError("Duplicate external_id '{}' found".format(asset.external_id))
+            external_ids_seen.add(asset.external_id)
 
             parent_ref = asset.parent_external_id
             if parent_ref:

diff --git a/docs/source/cognite.rst b/docs/source/cognite.rst
@@ -102,8 +102,10 @@ To make an asset a child of an existing asset, you must specify a parent ID.
 
 To post an entire asset hierarchy, you can describe the relations within your asset hierarchy
 using the :code:`external_id` and :code:`parent_external_id` attributes on the :code:`Asset` object. You can post
-an arbitrary number of assets, and the SDK will split the request into multiple requests and create the assets
-in the correct order
+an arbitrary number of assets, and the SDK will split the request into multiple requests. To make sure that the
+assets are posted in the correct order, you can use the .create_hierarchy() function, which takes care of the
+sorting before splitting the request into smaller chunks. However, note that the .create_hierarchy() function requires the
+external_id property to be set for all assets.
 
 This example shows how to post a three levels deep asset hierarchy consisting of three assets.
 
@@ -115,9 +117,9 @@ This example shows how to post a three levels deep asset hierarchy consisting of
     >>> root = Asset(name="root", external_id="1")
     >>> child = Asset(name="child", external_id="2", parent_external_id="1")
     >>> descendant = Asset(name="descendant", external_id="3", parent_external_id="2")
-    >>> c.assets.create([root, child, descendant])
+    >>> c.assets.create_hierarchy([root, child, descendant])
 
-Wrap the .create() call in a try-except to get information if posting the assets fails:
+Wrap the .create_hierarchy() call in a try-except to get information if posting the assets fails:
 
 - Which assets were posted. (The request yielded a 201.)
 - Which assets may have been posted. (The request yielded 5xx.)
@@ -127,7 +129,7 @@ Wrap the .create() call in a try-except to get information if posting the assets
 
     >>> from cognite.client.exceptions import CogniteAPIError
     >>> try:
-    ...     c.create([root, child, descendant])
+    ...     c.create_hierarchy([root, child, descendant])
     >>> except CogniteAPIError as e:
     ...     assets_posted = e.successful
     ...     assets_may_have_been_posted = e.unknown
@@ -277,6 +279,10 @@ Create assets
 ^^^^^^^^^^^^^
 .. automethod:: cognite.client._api.assets.AssetsAPI.create
 
+Create asset hierarchy
+^^^^^^^^^^^^^^^^^^^^^^
+.. automethod:: cognite.client._api.assets.AssetsAPI.create_hierarchy
+
 Delete assets
 ^^^^^^^^^^^^^
 .. automethod:: cognite.client._api.assets.AssetsAPI.delete

diff --git a/tests/tests_integration/test_api/test_assets.py b/tests/tests_integration/test_api/test_assets.py
@@ -47,7 +47,7 @@ def new_asset_hierarchy(post_spy):
     assets = generate_asset_tree(random_prefix + "0", depth=5, children_per_node=5)
 
     with set_request_limit(COGNITE_CLIENT.assets, 50):
-        COGNITE_CLIENT.assets.create(assets)
+        COGNITE_CLIENT.assets.create_hierarchy(assets)
 
     assert 20 < COGNITE_CLIENT.assets._post.call_count < 30
 
@@ -116,7 +116,7 @@ def test_create_asset_hierarchy_parent_external_id_not_in_request(self, new_root
             root_external_id=root.external_id, depth=5, children_per_node=10, current_depth=2
         )
 
-        COGNITE_CLIENT.assets.create(children)
+        COGNITE_CLIENT.assets.create_hierarchy(children)
 
         external_ids = [asset.external_id for asset in children] + [root.external_id]
         posted_assets = COGNITE_CLIENT.assets.retrieve_multiple(external_ids=external_ids)

diff --git a/tests/tests_unit/test_api/test_assets.py b/tests/tests_unit/test_api/test_assets.py
@@ -253,7 +253,7 @@ def test_validate_asset_hierarchy_duplicate_ref_ids(self):
 
     def test_validate_asset_hierarchy__more_than_limit_only_resolved_assets(self):
         with set_request_limit(ASSETS_API, 1):
-            _AssetPoster([Asset(parent_id=1), Asset(parent_id=2)], ASSETS_API)
+            _AssetPoster([Asset(external_id="a1", parent_id=1), Asset(external_id="a2", parent_id=2)], ASSETS_API)
 
     def test_validate_asset_hierarchy_circular_dependencies(self):
         assets = [
@@ -358,7 +358,7 @@ def request_callback(request):
     def test_post_hierarchy(self, limit, depth, children_per_node, expected_num_calls, mock_post_asset_hierarchy):
         assets = generate_asset_tree(root_external_id="0", depth=depth, children_per_node=children_per_node)
         with set_request_limit(ASSETS_API, limit):
-            created_assets = ASSETS_API.create(assets)
+            created_assets = ASSETS_API.create_hierarchy(assets)
 
         assert len(assets) == len(created_assets)
         assert expected_num_calls - 1 <= len(mock_post_asset_hierarchy.calls) <= expected_num_calls + 1
@@ -370,7 +370,9 @@ def test_post_hierarchy(self, limit, depth, children_per_node, expected_num_call
 
     def test_post_assets_over_limit_only_resolved(self, mock_post_asset_hierarchy):
         with set_request_limit(ASSETS_API, 1):
-            _AssetPoster([Asset(parent_id=1), Asset(parent_id=2)], ASSETS_API).post()
+            _AssetPoster(
+                [Asset(external_id="a1", parent_id=1), Asset(external_id="a2", parent_id=2)], ASSETS_API
+            ).post()
         assert 2 == len(mock_post_asset_hierarchy.calls)
 
     @pytest.fixture
@@ -422,7 +424,7 @@ def test_post_with_failures(self, mock_post_assets_failures):
             Asset(name="200", parent_external_id="03", external_id="031"),
         ]
         with pytest.raises(CogniteAPIError) as e:
-            ASSETS_API.create(assets)
+            ASSETS_API.create_hierarchy(assets)
 
         assert {a.external_id for a in e.value.unknown} == {"03"}
         assert {a.external_id for a in e.value.failed} == {"02", "021", "0211", "031"}