# coding=utf-8 # -------------------------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- import json from knack.log import get_logger from azure.cli.core.azclierror import ForbiddenError, RequiredArgumentMissingError, InvalidArgumentValueError from azext_iot.common.utility import process_json_arg, handle_service_exception, scantree from azext_iot.digitaltwins.common import ADTModelCreateFailurePolicy from azext_iot.digitaltwins.providers.base import DigitalTwinsProvider from azext_iot.sdk.digitaltwins.dataplane.models import ErrorResponseException from tqdm import tqdm logger = get_logger(__name__) MAX_MODELS_API_LIMIT = 250 def get_model_dependencies(model, model_id_to_model_map=None): """Return a list of dependency DTMIs for a given model""" dependencies = [] # Add everything that would have dependency DTMIs, worry about flattening later if "contents" in model: components = [item["schema"] for item in model["contents"] if item["@type"] == "Component"] dependencies.extend(components) if "extends" in model: dependencies.append(model['extends']) # Go through gathered items, get the DTMI references, and flatten if needed no_dup = set() for item in dependencies: # Models defined in a DTDL can implement extensions of up to two interfaces. # These interfaces can be in the form of a DTMI reference, or a nested model. if isinstance(item, str): # If its just a string, thats a single DTMI reference, so just add that to our set no_dup.add(item) # Calculate recursive dependencies if model id to model map is passed if model_id_to_model_map is not None: dep_model = model_id_to_model_map[item] no_dup.update(set(get_model_dependencies(dep_model, model_id_to_model_map))) elif isinstance(item, dict): # If its a single nested model, get its dtmi reference, dependencies and add them no_dup.update(set(get_model_dependencies(item, model_id_to_model_map))) elif isinstance(item, list): # If its a list, could have DTMIs or nested models for sub_item in item: if isinstance(sub_item, str): # If there are strings in the list, that's a DTMI reference, so add it no_dup.add(sub_item) # Calculate recursive dependencies if model id to model map is passed if model_id_to_model_map is not None: sub_dep_model = model_id_to_model_map[sub_item] no_dup.update(set(get_model_dependencies(sub_dep_model, model_id_to_model_map))) elif isinstance(sub_item, dict): # This is a nested model. Now go get its dependencies and add them no_dup.update(set(get_model_dependencies(sub_item, model_id_to_model_map))) return list(no_dup) class ModelProvider(DigitalTwinsProvider): def __init__(self, cmd, name, rg=None): super(ModelProvider, self).__init__( cmd=cmd, name=name, rg=rg, ) self.model_sdk = self.get_sdk().digital_twin_models def add(self, max_models_per_batch: int, models=None, from_directory=None, failure_policy=ADTModelCreateFailurePolicy.ROLLBACK.value): if not any([models, from_directory]): raise RequiredArgumentMissingError("Provide either --models or --from-directory.") # If both arguments are provided. --models wins. payload = [] models_per_batch = max_models_per_batch if models: models_result = process_json_arg(content=models, argument_name="models") if isinstance(models_result, list): payload.extend(models_result) elif isinstance(models_result, dict): payload.append(models_result) elif from_directory: payload = self._process_directory(from_directory=from_directory) logger.info("Models payload %s", json.dumps(payload)) models_created = [] try: # Process models in batches if models to process exceed the API limit if len(payload) > MAX_MODELS_API_LIMIT: model_id_to_model_map = {} for model_def in payload: model_id_to_model_map[model_def['@id']] = model_def # Create a dictionary to categorize models by their number of dependencies dep_count_to_models_map = {} for model in payload: num_dependencies = len(get_model_dependencies(model, model_id_to_model_map)) if num_dependencies not in dep_count_to_models_map: dep_count_to_models_map[num_dependencies] = [] dep_count_to_models_map[num_dependencies].append(model) # Sort by dependency count dep_count_to_models_tuples = sorted(dep_count_to_models_map.items()) models_batch = [] response = [] pbar = tqdm(total=len(payload), desc='Creating models...', ascii=' #') # The tuples being iterated are sorted by dependency count, hence models with 0 dependencies go first, # followed by models with 1 dependency, then 2 dependencies and so on... This ensures that all dependencies # of each model being added were either already added in a previous iteration or are in the current payload. for _, models_list in dep_count_to_models_tuples: while len(models_batch) + len(models_list) > models_per_batch: num_models_to_add = models_per_batch - len(models_batch) models_batch.extend(models_list[0:num_models_to_add]) response.extend(self.model_sdk.add(models_batch, raw=True).response.json()) models_created.extend([model['@id'] for model in models_batch]) pbar.update(len(models_batch)) # Remove the model ids which have been processed models_list = models_list[num_models_to_add:] models_batch = [] models_batch.extend(models_list) # Process the last set of model ids if len(models_batch) > 0: pbar.update(len(models_batch)) response.extend(self.model_sdk.add(models_batch, raw=True).response.json()) pbar.close() return response return self.model_sdk.add(payload, raw=True).response.json() except ErrorResponseException as e: if len(models_created) > 0: pbar.close() # Delete all models created by this operation when the failure policy is set to 'Rollback' if failure_policy == ADTModelCreateFailurePolicy.ROLLBACK.value: logger.error( "Error creating models. Deleting {} models created by this operation...".format(len(models_created)) ) # Models will be deleted in the reverse order they were created. # Hence, ensuring each model's dependencies are deleted after deleting the model. models_created.reverse() for model_id in models_created: self.delete(model_id) # Models created by this operation are not deleted when the failure policy is set to 'None' elif failure_policy == ADTModelCreateFailurePolicy.NONE.value: logger.error( "Error creating current model batch. Successfully created {} models.".format(len(models_created)) ) else: raise InvalidArgumentValueError( "Invalid failure policy: {}. Supported values are: '{}' and '{}'".format( failure_policy, ADTModelCreateFailurePolicy.ROLLBACK.value, ADTModelCreateFailurePolicy.NONE.value ) ) # @vilit - hack to customize 403's to have more specific error messages if e.response.status_code == 403: error_text = "Current principal access is forbidden. Please validate rbac role assignments." raise ForbiddenError(error_text) handle_service_exception(e) def _process_directory(self, from_directory): logger.debug( "Documents contained in directory: {}, processing...".format(from_directory) ) payload = [] for entry in scantree(from_directory): if all( [not entry.name.endswith(".json"), not entry.name.endswith(".dtdl")] ): logger.debug( "Skipping {} - model file must end with .json or .dtdl".format( entry.path ) ) continue entry_json = process_json_arg(content=entry.path, argument_name=entry.name) payload.append(entry_json) return payload def get(self, id, get_definition=False): try: return self.model_sdk.get_by_id( id=id, include_model_definition=get_definition, raw=True ).response.json() except ErrorResponseException as e: handle_service_exception(e) def list( self, get_definition=False, dependencies_for=None, top=None ): # top is guarded for int() in arg def from azext_iot.sdk.digitaltwins.dataplane.models import DigitalTwinModelsListOptions list_options = DigitalTwinModelsListOptions(max_items_per_page=top) return self.model_sdk.list( dependencies_for=dependencies_for, include_model_definition=get_definition, digital_twin_models_list_options=list_options, ) def update(self, id, decommission: bool): patched_model = [ {"op": "replace", "path": "/decommissioned", "value": decommission} ] # Does not return model object upon updating try: self.model_sdk.update(id=id, update_model=patched_model) except ErrorResponseException as e: handle_service_exception(e) return self.get(id=id) def delete(self, id: str): try: self.model_sdk.delete(id=id) except ErrorResponseException as e: handle_service_exception(e) def delete_all(self): # Get all models incoming_pager = self.list(get_definition=True) incoming_result = [] try: while True: incoming_result.extend(incoming_pager.advance_page()) except StopIteration: pass except ErrorResponseException as e: handle_service_exception(e) # Build dict of model_id : set of parent_ids parsed_models = {model.id: set() for model in incoming_result} for model in incoming_result: # Parse dependents, add current model as parent of dependents dependencies = get_model_dependencies(model.model) for d_id in dependencies: parsed_models[d_id].add(model.id) def delete_parents(model_id, model_dict): # Check if current model has been deleted already if model_id not in model_dict: return # Delete parents first for parent_id in model_dict[model_id]: if parent_id in model_dict: delete_parents(parent_id, model_dict) # Delete current model and remove references del model_dict[model_id] try: self.delete(model_id) except Exception as e: logger.warning(f"Could not delete model {model_id}; error is {e}") while len(parsed_models) > 0: model_id = next(iter(parsed_models)) delete_parents(model_id, parsed_models)