def _get_children(self, assets: AssetList) -> AssetList:
ids = [a.id for a in assets]
tasks = []
chunk_size = 100
for i in range(0, len(ids), chunk_size):
tasks.append({"parent_ids": ids[i:i + chunk_size], "limit": -1})
tasks_summary = utils._concurrency.execute_tasks_concurrently(
self.list, tasks=tasks, max_workers=self._config.max_workers)
tasks_summary.raise_compound_exception_if_failed_tasks()
res_list = tasks_summary.results
children = AssetList([])
for res in res_list:
children.extend(res)
return children
def test_to_pandas_nullable_int(self):
import pandas as pd
for camel_case in [False, True]:
assert (pd.Int64Dtype() == AssetList([
Asset(parent_id=123),
Asset(parent_id=None)
]).to_pandas(camel_case=camel_case).dtypes[0])
def test_get_related_resources_should_not_return_duplicates(self, resource_class, resource_list_class, method):
r1 = resource_class(id=1)
r2 = resource_class(id=2)
r3 = resource_class(id=3)
resources_a1 = resource_list_class([r1])
resources_a2 = resource_list_class([r2, r3])
resources_a3 = resource_list_class([r2, r3])
mock_cognite_client = mock.MagicMock()
mock_method = getattr(mock_cognite_client, method)
mock_method.list.side_effect = [resources_a1, resources_a2, resources_a3]
mock_method._config = mock.Mock(max_workers=3)
assets = AssetList([Asset(id=1), Asset(id=2), Asset(id=3)], cognite_client=mock_cognite_client)
assets._retrieve_chunk_size = 1
resources = getattr(assets, method)()
expected = [r1, r2, r3]
assert expected == resources
def run(self):
unblocked_assets_lists = self._get_unblocked_assets()
for unblocked_assets in unblocked_assets_lists:
self.request_queue.put(list(unblocked_assets))
while self.assets_remaining():
res = self.response_queue.get()
if isinstance(res, _AssetsFailedToPost):
if isinstance(res.exc, CogniteAPIError):
self.exception = res.exc
for asset in res.assets:
if res.exc.code >= 500:
self.may_have_been_posted_assets.add(asset)
elif res.exc.code >= 400:
self.not_posted_assets.add(asset)
for descendant in self._get_descendants(asset):
self.not_posted_assets.add(descendant)
else:
raise res.exc
else:
for asset in res:
self.posted_assets.add(asset)
self.successfully_posted_external_ids.add(
asset.external_id)
unblocked_assets_lists = self._get_unblocked_assets()
for unblocked_assets in unblocked_assets_lists:
self.request_queue.put(list(unblocked_assets))
if len(self.may_have_been_posted_assets) > 0 or len(
self.not_posted_assets) > 0:
if isinstance(self.exception, CogniteAPIError):
raise CogniteAPIError(
message=self.exception.message,
code=self.exception.code,
x_request_id=self.exception.x_request_id,
missing=self.exception.missing,
duplicated=self.exception.duplicated,
successful=AssetList(list(self.posted_assets)),
unknown=AssetList(list(self.may_have_been_posted_assets)),
failed=AssetList(list(self.not_posted_assets)),
unwrap_fn=lambda a: a.external_id,
)
raise self.exception
def post(self):
workers = []
for _ in range(self.client._config.max_workers):
worker = _AssetPosterWorker(self.client, self.request_queue, self.response_queue)
workers.append(worker)
worker.start()
self.run()
for worker in workers:
worker.stop = True
return AssetList(self.posted_assets)
def retrieve_subtree(self, id: int = None, external_id: str = None, depth: int = None) -> AssetList:
"""Retrieve the subtree for this asset up to a specified depth.
Args:
id (int): Id of the root asset in the subtree.
external_id (str): External id of the root asset in the subtree.
depth (int): Retrieve assets up to this depth below the root asset in the subtree. Omit to get the entire
subtree.
Returns:
AssetList: The requested assets.
"""
utils._auxiliary.assert_exactly_one_of_id_or_external_id(id, external_id)
asset = self.retrieve(id=id, external_id=external_id)
subtree = self._get_asset_subtree(AssetList([asset]), current_depth=0, depth=depth)
subtree._cognite_client = self._cognite_client
return subtree
def test_get_files(self):
c.files.list = mock.MagicMock()
a = AssetList(resources=[Asset(id=1)], cognite_client=c)
a.files()
assert c.files.list.call_args == call(asset_ids=[1], limit=-1)
assert c.files.list.call_count == 1
