diff --git a/nipyapi/canvas.py b/nipyapi/canvas.py index 09e6c2a8..f37226f1 100644 --- a/nipyapi/canvas.py +++ b/nipyapi/canvas.py @@ -719,13 +719,16 @@ def _starting_schedule_processor(processor_): return False -def update_process_group(pg, update): +def update_process_group(pg, update, refresh=True): """ Updates a given Process Group. Args: - pg (ProcessGroupEntity): The Processor to target for update + pg (ProcessGroupEntity): The Process Group to + target for update update (dict): key:value pairs to update + refresh (bool): Whether to refresh the Process Group before + applying the update Returns: (ProcessGroupEntity): The updated ProcessorEntity @@ -733,6 +736,8 @@ def update_process_group(pg, update): """ assert isinstance(pg, nipyapi.nifi.ProcessGroupEntity) with nipyapi.utils.rest_exceptions(): + if refresh: + pg = get_process_group(pg.id, 'id') return nipyapi.nifi.ProcessGroupsApi().update_process_group( id=pg.id, body=nipyapi.nifi.ProcessGroupEntity( @@ -1282,6 +1287,7 @@ def get_controller(identifier, identifier_type='name', bool_response=False): assert isinstance(identifier, six.string_types) assert identifier_type in ['name', 'id'] handle = nipyapi.nifi.ControllerServicesApi() + out = None try: if identifier_type == 'id': out = handle.get_controller_service(identifier) diff --git a/nipyapi/config.py b/nipyapi/config.py index 8179cd0e..7839cad3 100644 --- a/nipyapi/config.py +++ b/nipyapi/config.py @@ -25,7 +25,7 @@ # convenience function for this in nipyapi.utils.set_endpoint # Set Default Host for NiFi -default_host = '34.244.104.232' # Default to localhost for release +default_host = 'localhost' # Default to localhost for release # nifi_config.host = os.getenv( 'NIFI_API_ENDPOINT', diff --git a/nipyapi/parameters.py b/nipyapi/parameters.py index a3faa017..5fefa9b4 100644 --- a/nipyapi/parameters.py +++ b/nipyapi/parameters.py @@ -18,7 +18,9 @@ "list_all_parameter_contexts", "create_parameter_context", "delete_parameter_context", "get_parameter_context", "update_parameter_context", "prepare_parameter", - "delete_parameter_from_context" + "delete_parameter_from_context", "upsert_parameter_to_context", + "assign_context_to_process_group", + "remove_context_from_process_group" ] @@ -217,3 +219,45 @@ def upsert_parameter_to_context(context, parameter): enforce_min_ver('1.10.0') context.component.parameters = [parameter] return update_parameter_context(context=context) + + +def assign_context_to_process_group(pg, context_id): + """ + Assigns a given Parameter Context to a specific Process Group + + Args: + pg (ProcessGroupEntity): The Process Group to target + context_id (str): The ID of the Parameter Context + + Returns: + (ProcessGroupEntity) The updated Process Group + """ + assert isinstance(context_id, str) + return nipyapi.canvas.update_process_group( + pg=pg, + update={ + 'parameter_context': { + 'id': context_id + } + } + ) + + +def remove_context_from_process_group(pg): + """ + Clears any Parameter Context from the given Process Group + + Args: + pg (ProcessGroupEntity): The Process Group to target + + Returns: + (ProcessGroupEntity) The updated Process Group + """ + return nipyapi.canvas.update_process_group( + pg=pg, + update={ + 'parameter_context': { + 'id': None + } + } + ) diff --git a/nipyapi/utils.py b/nipyapi/utils.py index 5851ca1b..28cc022e 100644 --- a/nipyapi/utils.py +++ b/nipyapi/utils.py @@ -573,7 +573,6 @@ def wrapper(*args, **kwargs): nipyapi.registry.rest.ApiException) as e: if status_code is not None and e.status == int(status_code): return response - else: - _raise(ValueError(e.body), e) + _raise(ValueError(e.body), e) return wrapper return func_wrapper diff --git a/tests/test_parameters.py b/tests/test_parameters.py index 1fdc4503..701b8dd3 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -111,3 +111,23 @@ def test_upsert_parameter_to_context(regress_nifi, fix_context): parameters.prepare_parameter('Black', 'Lives', 'Matter', True) ) assert r1.component.parameters[0].parameter.description == 'Matter' + + +def test_assign_context_to_process_group(regress_nifi, fix_pg, fix_context): + if check_version('1.10.0') > 0: + pytest.skip("NiFi not 1.10+") + c1 = fix_context.generate() + pg1 = fix_pg.generate() + r1 = parameters.assign_context_to_process_group(pg1, c1.id) + assert r1.component.parameter_context.id == c1.id + + +def test_remove_context_from_process_group(regress_nifi, fix_pg, fix_context): + if check_version('1.10.0') > 0: + pytest.skip("NiFi not 1.10+") + c1 = fix_context.generate() + pg1 = fix_pg.generate() + r1 = parameters.assign_context_to_process_group(pg1, c1.id) + assert r1.component.parameter_context.id == c1.id + r2 = parameters.remove_context_from_process_group(pg1) + assert r2.component.parameter_context is None