feat(device): execute command on multiple devices

riocli/device/execute.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import typing
-
+import functools
 import click
 from click_help_colors import HelpColorsCommand
-
+from concurrent.futures import ThreadPoolExecutor
+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
+from queue import Queue
 
 
 @click.command(
@@ -37,19 +39,25 @@
     default=False,
     help="Run the command asynchronously.",
 )
-@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,
     timeout: int,
     shell: str,
     run_async: bool,
     command: typing.List[str],
+    workers: int = 10,
 ) -> None:
-    """Execute commands on a device.
+    """Execute commands on one or more devices.
 
     You can specify the user, shell, run-async, and timeout options to customize
     the command execution. To specify the user, use the --user flag.
@@ -63,19 +71,69 @@ def execute_command(
 
     Usage Examples:
 
-        $ rio device execute DEVICE_NAME "ls -l"
+        $ rio device execute DEVICE_NAME_OR_REGEX "ls -l"
+
+        To execute the command asynchronously:
+
+            $ rio device execute DEVICE_NAME_OR_REGEX "ls -l" --async
+
+        To run the command on all devices:
+
+            $ rio device execute ".*" "ls -l"
     """
+
+    client = new_client()
+
     try:
-        response = run_on_device(
-            device_guid=device_guid,
+        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=run_async,
             timeout=timeout,
+            result=result,
         )
+        with ThreadPoolExecutor(max_workers=workers) as executor:
+            executor.map(func, device_guids)
 
-        click.secho(response)
+        for device_guid, success, response in result.queue:
+            click.echo(
+                ">>> {}: {}\n{}".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, timeout, 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,
+            timeout=timeout,
+        )
+        result.put((device_guid, True, response))
+    except Exception as e:
+        result.put((device_guid, False, str(e)))