[AKS] az aks command invoke: Add support for --no-wait

src/azure-cli/azure/cli/command_modules/acs/custom.py

@@ -2655,7 +2655,7 @@ def aks_get_versions(cmd, client, location):
     return client.list_orchestrators(location, resource_type='managedClusters')
 
 
-def aks_runcommand(cmd, client, resource_group_name, name, command_string="", command_files=None):
+def aks_runcommand(cmd, client, resource_group_name, name, command_string="", command_files=None, no_wait=False):
     colorama.init()
 
     mc = client.get(resource_group_name, name)
@@ -2674,21 +2674,49 @@ def aks_runcommand(cmd, client, resource_group_name, name, command_string="", co
         request_payload.cluster_token = _get_dataplane_aad_token(
             cmd.cli_ctx, "6dae42f8-4368-4678-94ff-3960e28e3630")
 
-    commandResultFuture = client.begin_run_command(
-        resource_group_name, name, request_payload, polling_interval=5, retry_total=0)
-
-    return _print_command_result(cmd.cli_ctx, commandResultFuture.result(300))
+    command_result_poller = sdk_no_wait(
+        no_wait, client.begin_run_command, resource_group_name, name, request_payload, polling_interval=5, retry_total=0
+    )
+    if no_wait:
+        # pylint: disable=protected-access
+        command_result_polling_url = command_result_poller.polling_method()._initial_response.http_response.headers[
+            "location"
+        ]
+        command_id_regex = re.compile(r"commandResults\/(\w*)\?")
+        command_id = command_id_regex.findall(command_result_polling_url)[0]
+        _aks_command_result_in_progess_helper(client, resource_group_name, name, command_id)
+        return
+    return _print_command_result(cmd.cli_ctx, command_result_poller.result(300))
 
 
 def aks_command_result(cmd, client, resource_group_name, name, command_id=""):
     if not command_id:
         raise ValidationError('CommandID cannot be empty.')
 
-    commandResult = client.get_command_result(
-        resource_group_name, name, command_id)
+    commandResult = client.get_command_result(resource_group_name, name, command_id)
+    if commandResult is None:
+        _aks_command_result_in_progess_helper(client, resource_group_name, name, command_id)
+        return
     return _print_command_result(cmd.cli_ctx, commandResult)
 
 
+def _aks_command_result_in_progess_helper(client, resource_group_name, name, command_id):
+    # pylint: disable=unused-argument
+    def command_result_direct_response_handler(pipeline_response, *args, **kwargs):
+        deserialized_data = pipeline_response.context.get("deserialized_data", {})
+        if deserialized_data:
+            provisioning_state = deserialized_data.get("properties", {}).get("provisioningState", None)
+            started_at = deserialized_data.get("properties", {}).get("startedAt", None)
+            print(f"command id: {command_id}, started at: {started_at}, status: {provisioning_state}")
+            print(
+                f"Please use command \"az aks command result -g {resource_group_name} -n {name} -i {command_id}\" "
+                "to get the future execution result"
+            )
+        else:
+            print(f"failed to fetch command result for command id: {command_id}")
+    client.get_command_result(resource_group_name, name, command_id, cls=command_result_direct_response_handler)
+
+
 def _print_command_result(cli_ctx, commandResult):
     # cli_ctx.data['safe_params'] contains list of parameter name user typed in, without value.
     # cli core also use this calculate ParameterSetName header for all http request from cli.
@@ -2715,8 +2743,8 @@ def _print_command_result(cli_ctx, commandResult):
         return
 
     # *-ing state
-    print(f"{colorama.Fore.BLUE}command is in : {commandResult.provisioning_state} state{colorama.Style.RESET_ALL}")
-    return None
+    print(f"{colorama.Fore.BLUE}command is in {commandResult.provisioning_state} state{colorama.Style.RESET_ALL}")
+    return
 
 
 def _get_command_context(command_files):