diff --git a/riocli/device/execute.py b/riocli/device/execute.py index 0d4505a2..ae42ef97 100644 --- a/riocli/device/execute.py +++ b/riocli/device/execute.py @@ -11,13 +11,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import functools import typing +from concurrent.futures import ThreadPoolExecutor +from queue import Queue import click from click_help_colors import HelpColorsCommand +from riocli.config import new_client from riocli.constants import Colors -from riocli.device.util import name_to_guid +from riocli.device.util import fetch_devices from riocli.utils.execute import run_on_device @@ -29,29 +33,53 @@ ) @click.option('--user', default='root') @click.option('--shell', default='/bin/bash') -@click.argument('device-name', type=str) +@click.option('--workers', '-w', + help="Number of parallel workers for executing command. Defaults to 10.", type=int, default=10) +@click.argument('device-name-or-regex', type=str) @click.argument('command', nargs=-1) -@name_to_guid def execute_command( - device_name: str, - device_guid: str, + device_name_or_regex: str, user: str, shell: str, - command: typing.List[str] + command: typing.List[str], + workers: int = 10, ) -> None: """ Execute commands on the Device """ + client = new_client() + try: - response = run_on_device( - device_guid=device_guid, - user=user, - shell=shell, - command=command, - background=False, - ) - - click.secho(response) + devices = fetch_devices(client, device_name_or_regex, include_all=False, online_devices=True) except Exception as e: click.secho(str(e), fg=Colors.RED) raise SystemExit(1) from e + + if not devices: + click.secho('No device(s) found', fg=Colors.RED) + raise SystemExit(1) + + device_guids = [d.uuid for d in devices] + + try: + result = Queue() + func = functools.partial(_run_on_device, user=user, shell=shell, + command=command, background=False, result=result) + with ThreadPoolExecutor(max_workers=workers) as executor: + executor.map(func, device_guids) + + for device_guid, success, response in result.queue: + click.echo('>>> {}: {} -> {}'.format(device_guid, 'Success' if success else 'Failed', response)) + except Exception as e: + click.secho(str(e), fg=Colors.RED) + raise SystemExit(1) from e + + +def _run_on_device(device_guid, user, shell, command, background, result): + """Wrapper on top of run_on_device to capture the output in a queue""" + try: + response = run_on_device(device_guid=device_guid, command=command, + user=user, shell=shell, background=background) + result.put((device_guid, True, response)) + except Exception as e: + result.put((device_guid, False, str(e)))