Feature: Spark add output logs flag (#468)

* add output flag to cluster submit

* add output flag to cluster app-logs

* add output flag to job get-app-logs

* sort imports

* make spinner context

aztk_cli/spark/endpoints/cluster/cluster_add_user.py

@@ -1,8 +1,8 @@
 import argparse
 import typing
+
 import aztk.spark
-from aztk_cli import log
-from aztk_cli import utils, config
+from aztk_cli import config, log, utils
 
 
 def setup_parser(parser: argparse.ArgumentParser):

aztk_cli/spark/endpoints/cluster/cluster_app_logs.py

@@ -1,8 +1,10 @@
 import argparse
+import os
 import typing
+
 import aztk
+from aztk_cli import config, utils
 
-from aztk_cli import utils, config
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id',
@@ -14,7 +16,12 @@ def setup_parser(parser: argparse.ArgumentParser):
                         required=True,
                         help='The unique id of your job name')
 
-    parser.add_argument('--tail', dest='tail', action='store_true')
+    output_group = parser.add_mutually_exclusive_group()
+
+    output_group.add_argument('--output',
+                              help='Path to the file you wish to output to. If not \
+                                    specified, output is printed to stdout')
+    output_group.add_argument('--tail', dest='tail', action='store_true')
 
 
 def execute(args: typing.NamedTuple):
@@ -23,5 +30,10 @@ def execute(args: typing.NamedTuple):
     if args.tail:
         utils.stream_logs(client=spark_client, cluster_id=args.cluster_id, application_name=args.app_name)
     else:
-        app_logs = spark_client.get_application_log(cluster_id=args.cluster_id, application_name=args.app_name)
-        print(app_logs.log)
+        app_log = spark_client.get_application_log(cluster_id=args.cluster_id, application_name=args.app_name)
+        if args.output:
+            with utils.Spinner():
+                with open(os.path.abspath(os.path.expanduser(args.output)), "w", encoding="UTF-8") as f:
+                    f.write(app_log.log)
+        else:
+            print(app_log.log)

aztk_cli/spark/endpoints/cluster/cluster_copy.py

@@ -1,5 +1,6 @@
 import argparse
 import typing
+
 import aztk.spark
 from aztk_cli import config
 

aztk_cli/spark/endpoints/cluster/cluster_create.py

@@ -1,11 +1,11 @@
-import os
 import argparse
+import os
 import typing
+
 import aztk.spark
 from aztk.spark.models import ClusterConfiguration, UserConfiguration
-from aztk_cli import log
+from aztk_cli import config, log, utils
 from aztk_cli.config import load_aztk_spark_config
-from aztk_cli import utils, config
 
 
 def setup_parser(parser: argparse.ArgumentParser):
@@ -70,19 +70,14 @@ def execute(args: typing.NamedTuple):
         cluster_conf.user_configuration = None
 
     utils.print_cluster_conf(cluster_conf, wait)
-    spinner = utils.Spinner()
-    spinner.start()
-
-    # create spark cluster
-    cluster = spark_client.create_cluster(
-        cluster_conf,
-        wait=wait
-    )
-
-    spinner.stop()
+    with utils.Spinner():
+        # create spark cluster
+        cluster = spark_client.create_cluster(
+            cluster_conf,
+            wait=wait
+        )
 
     if wait:
         log.info("Cluster %s created successfully.", cluster.id)
     else:
         log.info("Cluster %s is being provisioned.", cluster.id)
-

aztk_cli/spark/endpoints/cluster/cluster_delete.py

@@ -1,7 +1,8 @@
 import argparse
 import typing
+
 import aztk
-from aztk_cli import log, config
+from aztk_cli import config, log
 
 
 def setup_parser(parser: argparse.ArgumentParser):

aztk_cli/spark/endpoints/cluster/cluster_get.py

@@ -1,8 +1,8 @@
 import argparse
 import typing
+
 import aztk
-from aztk_cli import log
-from aztk_cli import utils, config
+from aztk_cli import config, log, utils
 
 
 def setup_parser(parser: argparse.ArgumentParser):

aztk_cli/spark/endpoints/cluster/cluster_list.py

@@ -1,7 +1,8 @@
 import argparse
 import typing
+
 import aztk
-from aztk_cli import utils, config
+from aztk_cli import config, utils
 
 
 def setup_parser(_: argparse.ArgumentParser):

aztk_cli/spark/endpoints/cluster/cluster_run.py

@@ -1,7 +1,9 @@
 import argparse
 import typing
+
 import aztk.spark
-from aztk_cli import utils, config
+from aztk_cli import config, utils
+
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id',

aztk_cli/spark/endpoints/cluster/cluster_ssh.py

@@ -1,11 +1,12 @@
 import argparse
 import typing
-from aztk_cli import log
-from aztk_cli import utils, config
-from aztk_cli.config import SshConfig
-import aztk
+
 import azure.batch.models.batch_error as batch_error
+
+import aztk
 from aztk.models import ClusterConfiguration
+from aztk_cli import config, log, utils
+from aztk_cli.config import SshConfig
 
 
 def setup_parser(parser: argparse.ArgumentParser):
@@ -106,4 +107,3 @@ def print_plugin_ports(cluster_config: ClusterConfiguration):
 
                         url = "{0}{1}".format(http_prefix, port.public_port)
                         utils.log_property(label, url)
-

aztk_cli/spark/endpoints/cluster/cluster_submit.py

@@ -1,8 +1,10 @@
 import argparse
+import os
 import sys
 import typing
-from aztk_cli import utils, config, log
+
 import aztk.spark
+from aztk_cli import config, log, utils
 
 
 def setup_parser(parser: argparse.ArgumentParser):
@@ -64,6 +66,10 @@ def setup_parser(parser: argparse.ArgumentParser):
                         help='Number of times the Spark job may be retried \
                               if there is a failure')
 
+    parser.add_argument('--output',
+                        help='Path to the file you wish to output to. If not \
+                              specified, output is printed to stdout')
+
     parser.add_argument('app',
                         help='App jar OR python file to execute. Use absolute \
                               path to reference file.')
@@ -73,6 +79,9 @@ def setup_parser(parser: argparse.ArgumentParser):
 
 
 def execute(args: typing.NamedTuple):
+    if not args.wait and args.output:
+        raise aztk.error.AztkError("--output flag requires --wait flag")
+
     spark_client = aztk.spark.Client(config.load_aztk_secrets())
     jars = []
     py_files = []
@@ -141,6 +150,14 @@ def execute(args: typing.NamedTuple):
     )
 
     if args.wait:
-        exit_code = utils.stream_logs(client=spark_client, cluster_id=args.cluster_id, application_name=args.name)
-        sys.exit(exit_code)
+        if not args.output:
+            exit_code = utils.stream_logs(client=spark_client, cluster_id=args.cluster_id, application_name=args.name)
+        else:
+            with utils.Spinner():
+                spark_client.wait_until_application_done(cluster_id=args.cluster_id, task_id=args.name)
+                application_log = spark_client.get_application_log(cluster_id=args.cluster_id, application_name=args.name)
+                with open(os.path.abspath(os.path.expanduser(args.output)), "w", encoding="UTF-8") as f:
+                    f.write(application_log.log)
+                exit_code = application_log.exit_code
 
+        sys.exit(exit_code)

aztk_cli/spark/endpoints/init.py

@@ -1,7 +1,8 @@
-import os
 import argparse
+import os
 import typing
 from distutils.dir_util import copy_tree
+
 import aztk.utils.constants as constants
 
 

aztk_cli/spark/endpoints/job/delete.py

@@ -1,7 +1,8 @@
 import argparse
 import typing
+
 import aztk.spark
-from aztk_cli import log, config
+from aztk_cli import config, log
 
 
 def setup_parser(parser: argparse.ArgumentParser):

aztk_cli/spark/endpoints/job/get.py

@@ -1,9 +1,10 @@
 import argparse
-import typing
 import time
+import typing
+
 import aztk.spark
-from aztk_cli import config
-from aztk_cli import utils
+from aztk_cli import config, utils
+
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id',

aztk_cli/spark/endpoints/job/get_app.py

@@ -1,9 +1,10 @@
 import argparse
-import typing
 import time
+import typing
+
 import aztk.spark
-from aztk_cli import config
-from aztk_cli import utils
+from aztk_cli import config, utils
+
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id',

aztk_cli/spark/endpoints/job/get_app_logs.py

@@ -1,7 +1,10 @@
 import argparse
+import os
 import typing
+
 import aztk.spark
-from aztk_cli import utils, config
+from aztk_cli import config, utils
+
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id',
@@ -12,9 +15,17 @@ def setup_parser(parser: argparse.ArgumentParser):
                         dest='app_name',
                         required=True,
                         help='The unique id of your job name')
+    parser.add_argument('--output',
+                        help='Path to the file you wish to output to. If not \
+                              specified, output is printed to stdout')
 
 
 def execute(args: typing.NamedTuple):
     spark_client = aztk.spark.Client(config.load_aztk_secrets())
-    app_logs = spark_client.get_job_application_log(args.job_id, args.app_name)
-    print(app_logs.log)
+    app_log = spark_client.get_job_application_log(args.job_id, args.app_name)
+    if args.output:
+        with utils.Spinner():
+            with open(os.path.abspath(os.path.expanduser(args.output)), "w", encoding="UTF-8") as f:
+                f.write(app_log.log)
+    else:
+        print(app_log.log)

aztk_cli/spark/endpoints/job/list.py

@@ -1,9 +1,10 @@
 import argparse
-import typing
 import time
+import typing
+
 import aztk.spark
-from aztk_cli import config
-from aztk_cli import utils
+from aztk_cli import config, utils
+
 
 def setup_parser(_: argparse.ArgumentParser):
     # No arguments for list yet

aztk_cli/spark/endpoints/job/list_apps.py

@@ -1,8 +1,9 @@
 import argparse
 import typing
+
 import aztk.spark
-from aztk_cli import config
-from aztk_cli import utils
+from aztk_cli import config, utils
+
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id',

aztk_cli/spark/endpoints/job/stop.py

@@ -1,9 +1,10 @@
 import argparse
-import typing
 import time
+import typing
+
 import aztk.spark
-from aztk_cli import config
-from aztk_cli import utils
+from aztk_cli import config, utils
+
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id',

aztk_cli/spark/endpoints/job/stop_app.py

@@ -1,10 +1,9 @@
 import argparse
-import typing
 import time
+import typing
+
 import aztk.spark
-from aztk_cli import config
-from aztk_cli import utils
-from aztk_cli import log
+from aztk_cli import config, log, utils
 
 
 def setup_parser(parser: argparse.ArgumentParser):

aztk_cli/spark/endpoints/job/submit.py

@@ -1,8 +1,9 @@
 import argparse
-import typing
 import time
+import typing
+
 import aztk.spark
-from aztk_cli import config, utils, log
+from aztk_cli import config, log, utils
 from aztk_cli.config import JobConfig, load_aztk_spark_config
 
 

aztk_cli/spark/endpoints/spark.py

@@ -1,9 +1,10 @@
 import argparse
 import typing
 
+from . import init
 from .cluster import cluster
 from .job import job
-from . import init
+
 
 def setup_parser(parser: argparse.ArgumentParser):
     subparsers = parser.add_subparsers(

aztk_cli/utils.py

@@ -379,6 +379,12 @@ def __init__(self, delay=None):
         self.spinner_generator = self.spinning_cursor()
         if delay and float(delay): self.delay = delay
 
+    def __enter__(self):
+        return self.start()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self.stop()
+
     def spinner_task(self):
         while self.busy:
             sys.stdout.write(next(self.spinner_generator))