From 6424daf2f96db0caa230bac435a601e825daa4fd Mon Sep 17 00:00:00 2001
From: Troy Dai <troy.dai@outlook.com>
Date: Thu, 26 Apr 2018 16:27:32 -0700
Subject: [PATCH] Async command

---
 src/a01/cli.py                 | 6 +++++-
 src/a01/commands/get_run.py    | 9 +++++----
 src/a01/operations/__init__.py | 4 ++--
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/a01/cli.py b/src/a01/cli.py
index 49cbd4b..92b02b8 100644
--- a/src/a01/cli.py
+++ b/src/a01/cli.py
@@ -1,4 +1,5 @@
 import argparse
+import asyncio
 import inspect
 from typing import Callable, Collection
 
@@ -145,7 +146,10 @@ class CommandDefinition(object):
         logger.info(f'execute [{self.name}] -> {self.func}')
 
         kwargs = {parameter: getattr(args, parameter) for parameter in self.signature.parameters}
-        return self.func(**kwargs)
+        if inspect.iscoroutinefunction(self.func):
+            return asyncio.get_event_loop().run_until_complete(self.func(**kwargs))
+        else:
+            return self.func(**kwargs)
 
     def setup(self, parser: argparse.ArgumentParser) -> None:
         parser.description = self.description
diff --git a/src/a01/commands/get_run.py b/src/a01/commands/get_run.py
index d98d3df..d39c0a8 100644
--- a/src/a01/commands/get_run.py
+++ b/src/a01/commands/get_run.py
@@ -8,7 +8,7 @@ import colorama
 from a01.cli import cmd, arg
 from a01.output import output_in_table
 from a01.models import TaskCollection
-from a01.operations import query_tasks_by_run, query_run
+from a01.operations import query_tasks_by_run_async, query_run_async
 
 
 @cmd('get run', desc='Retrieve a run')
@@ -23,12 +23,13 @@ from a01.operations import query_tasks_by_run, query_run
      help='When download the recording files the files are arranged in directory structure mimic Azure CLI '
           'source code.')
 @arg('raw', help='For debug.')
-def get_run(run_id: str, log: bool = False, recording: bool = False, recording_az_mode: bool = False,
+async def get_run(run_id: str, log: bool = False, recording: bool = False, recording_az_mode: bool = False,
             show_all: bool = False, raw: bool = False) -> None:
     logger = logging.getLogger(__name__)
 
     try:
-        tasks = TaskCollection(query_tasks_by_run(run_id), run_id)
+        tasks = await query_tasks_by_run_async(run_id)
+        tasks = TaskCollection(tasks, run_id)
         output_in_table(tasks.get_table_view(failed=not show_all), headers=tasks.get_table_header())
         output_in_table(tasks.get_summary(), tablefmt='plain')
 
@@ -46,7 +47,7 @@ def get_run(run_id: str, log: bool = False, recording: bool = False, recording_a
                 task.download_recording(recording_az_mode)
 
         if raw:
-            run = query_run(run_id=run_id)
+            run = await query_run_async(run_id=run_id)
             print(json.dumps(run.to_dict(), indent=2))
     except ValueError as err:
         logger.error(err)
diff --git a/src/a01/operations/__init__.py b/src/a01/operations/__init__.py
index ad5d932..7624df6 100644
--- a/src/a01/operations/__init__.py
+++ b/src/a01/operations/__init__.py
@@ -1,3 +1,3 @@
 # pylint: disable=unused-import
-from .query_tasks import query_tasks, query_tasks_by_run
-from .query_runs import query_run, query_runs
+from .query_tasks import query_tasks, query_tasks_by_run, query_tasks_by_run_async, query_tasks_async
+from .query_runs import query_run, query_runs, query_run_async, query_runs_async