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