diff --git a/tdp/cli/commands/plan/dag.py b/tdp/cli/commands/plan/dag.py index 93897ad1..3a83f910 100644 --- a/tdp/cli/commands/plan/dag.py +++ b/tdp/cli/commands/plan/dag.py @@ -12,6 +12,7 @@ database_dsn_option, force_option, hosts_option, + no_hosts_limit_option, preview_option, rolling_interval_option, ) @@ -60,6 +61,9 @@ is_flag=True, help="Replace 'start' operations by 'stop' operations. This option should be used with `--reverse`.", ) +@no_hosts_limit_option( + help="Works with --host and does not limit the operation to the specified hosts." +) @hosts_option(help="Hosts where operations are launched. Can be used multiple times.") @rolling_interval_option @preview_option @@ -80,6 +84,7 @@ def dag( is_regex: bool = False, rolling_interval: Optional[int] = None, hosts: Optional[tuple[str]] = None, + no_host_limit: Optional[list[str]] = None, ): """Deploy from the DAG.""" @@ -118,6 +123,9 @@ def dag( else: click.echo("Creating a deployment plan for the whole DAG.") + if no_host_limit and not hosts: + click.BadOptionUsage("Cannot use `--no-host-limit` without --host argument.") + deployment = DeploymentModel.from_dag( dag, sources=sources, @@ -128,7 +136,8 @@ def dag( reverse=reverse, stop=stop, rolling_interval=rolling_interval, - host_names=hosts + host_names=hosts, + no_host_limit_operation=no_host_limit, ) if preview: print_deployment(deployment) diff --git a/tdp/cli/commands/plan/ops.py b/tdp/cli/commands/plan/ops.py index d298f8be..2fce2327 100644 --- a/tdp/cli/commands/plan/ops.py +++ b/tdp/cli/commands/plan/ops.py @@ -12,6 +12,7 @@ database_dsn_option, force_option, hosts_option, + no_hosts_limit_option, preview_option, rolling_interval_option, ) @@ -37,6 +38,9 @@ @database_dsn_option @preview_option @force_option +@no_hosts_limit_option( + help="Works with --host and does not limit the operation to the specified hosts." +) @rolling_interval_option def ops( operation_names: tuple[str], @@ -47,6 +51,7 @@ def ops( preview: bool, force: bool, rolling_interval: Optional[int] = None, + no_host_limit: Optional[tuple[str]] = None, ): """Run a list of operations.""" @@ -57,8 +62,10 @@ def ops( click.echo( f"Creating a deployment plan to run {len(operation_names)} operation(s)." ) + if no_host_limit and not hosts: + click.BadOptionUsage("Cannot use `--no-host-limit` without --host argument.") deployment = DeploymentModel.from_operations( - collections, operation_names, hosts, extra_vars, rolling_interval + collections, operation_names, hosts, no_host_limit, extra_vars, rolling_interval ) if preview: print_deployment(deployment) diff --git a/tdp/cli/params.py b/tdp/cli/params.py index 358d567e..da84e311 100644 --- a/tdp/cli/params.py +++ b/tdp/cli/params.py @@ -157,6 +157,34 @@ def decorator(fn: FC) -> FC: return decorator(func) +def no_hosts_limit_option( + func: Optional[FC] = None, *, help: str +) -> Callable[[FC], FC]: + """Add the `--no-host-limit` option to a Click command. + + Takes multiple operations and transforms them into a tuple of strings. Available as + "no_host_limit" in the command context. + + Args: + help: The help text for the option. + """ + + def decorator(fn: FC) -> FC: + return click.option( + "--no-host-limit", + "no_host_limit", + type=str, + multiple=True, + help=help, + )(fn) + + # Checks if the decorator was used without parentheses. + if func is None: + return decorator + else: + return decorator(func) + + def conf_option(func: FC) -> FC: """Add the `--conf` option to a Click command.""" return click.option( diff --git a/tdp/core/models/deployment_model.py b/tdp/core/models/deployment_model.py index 5bba4297..5c73e8d5 100644 --- a/tdp/core/models/deployment_model.py +++ b/tdp/core/models/deployment_model.py @@ -119,6 +119,7 @@ def from_dag( stop: bool = False, rolling_interval: Optional[int] = None, host_names: Optional[Iterable[str]] = None, + no_host_limit_operation: Optional[Iterable[str]] = None, ) -> DeploymentModel: """Generate a deployment plan from a DAG. @@ -174,6 +175,10 @@ def from_dag( }, state=DeploymentStateEnum.PLANNED, ) + if no_host_limit_operation: + NO_HOST_LIMIT_OPERATION = no_host_limit_operation + else: + NO_HOST_LIMIT_OPERATION = [None] operation_order = 1 for operation in operations: can_perform_rolling_restart = ( @@ -187,12 +192,21 @@ def from_dag( if host is None or ( isinstance(operation, PlaybookOperation) and host in operation.playbook.hosts + and ( + operation.name.name + not in [op.operation for op in deployment.operations] + or operation.name.name not in NO_HOST_LIMIT_OPERATION + ) ): deployment.operations.append( OperationModel( operation=operation.name.name, operation_order=operation_order, - host=host, + host=( + host + if not operation.name.name in NO_HOST_LIMIT_OPERATION + else None + ), extra_vars=None, state=OperationStateEnum.PLANNED, ) @@ -218,6 +232,7 @@ def from_operations( collections: Collections, operation_names: list[str], host_names: Optional[Iterable[str]] = None, + no_host_limit_operation: Optional[Iterable[str]] = None, extra_vars: Optional[Iterable[str]] = None, rolling_interval: Optional[int] = None, ) -> DeploymentModel: @@ -255,6 +270,10 @@ def from_operations( }, state=DeploymentStateEnum.PLANNED, ) + if no_host_limit_operation: + NO_HOST_LIMIT_OPERATION = no_host_limit_operation + else: + NO_HOST_LIMIT_OPERATION = [None] operation_order = 1 for operation in operations: can_perform_rolling_restart = ( @@ -270,29 +289,38 @@ def from_operations( if can_perform_rolling_restart else [None] ): - deployment.operations.append( - OperationModel( - operation=operation.name.name, - operation_order=operation_order, - host=host_name, - extra_vars=list(extra_vars) if extra_vars else None, - state=OperationStateEnum.PLANNED, - ) - ) - if can_perform_rolling_restart: - operation_order += 1 + if ( + operation.name.name + not in [op.operation for op in deployment.operations] + or operation.name.name not in NO_HOST_LIMIT_OPERATION + ): deployment.operations.append( OperationModel( - operation=OPERATION_SLEEP_NAME, + operation=operation.name.name, operation_order=operation_order, - host=None, - extra_vars=[ - f"{OPERATION_SLEEP_VARIABLE}={rolling_interval}" - ], + host=( + host_name + if not operation.name.name in NO_HOST_LIMIT_OPERATION + else None + ), + extra_vars=list(extra_vars) if extra_vars else None, state=OperationStateEnum.PLANNED, ) ) - operation_order += 1 + if can_perform_rolling_restart: + operation_order += 1 + deployment.operations.append( + OperationModel( + operation=OPERATION_SLEEP_NAME, + operation_order=operation_order, + host=None, + extra_vars=[ + f"{OPERATION_SLEEP_VARIABLE}={rolling_interval}" + ], + state=OperationStateEnum.PLANNED, + ) + ) + operation_order += 1 return deployment @staticmethod