diff --git a/bin/qds.py b/bin/qds.py index 074081db..08301349 100755 --- a/bin/qds.py +++ b/bin/qds.py @@ -15,6 +15,7 @@ from qds_sdk.app import AppCmdLine from qds_sdk.nezha import NezhaCmdLine from qds_sdk.user import UserCmdLine +from qds_sdk.template import TemplateCmdLine import os import sys @@ -78,6 +79,8 @@ " action --help\n" "\nScheduler subcommand:\n" " scheduler --help\n" + "\nTemplate subcommand:\n" + " template --help\n" "\nAccount subcommand:\n" " account --help\n" "\nNezha subcommand:\n" @@ -474,8 +477,12 @@ def nezhamain(args): result = NezhaCmdLine.run(args) print(result) -def main(): +def templatemain(args): + result = TemplateCmdLine.run(args) + print(result) + +def main(): optparser = OptionParser(usage=usage_str) optparser.add_option("--token", dest="api_token", default=os.getenv('QDS_API_TOKEN'), @@ -508,7 +515,7 @@ def main(): optparser.disable_interspersed_args() (options, args) = optparser.parse_args() - + if options.chatty: logging.basicConfig(level=logging.DEBUG) elif options.verbose: @@ -581,11 +588,13 @@ def main(): if a0 == "user": return usermain(args) + if a0 == "template": + return templatemain(args) cmdset = set(CommandClasses.keys()) sys.stderr.write("First command must be one of <%s>\n" % "|".join(cmdset.union(["cluster", "action", "scheduler", "report", - "dbtap", "role", "group", "app", "account", "nezha", "user"]))) + "dbtap", "role", "group", "app", "account", "nezha", "user", "template"]))) usage(optparser) diff --git a/qds_sdk/commands.py b/qds_sdk/commands.py index e3928633..d09e9fc7 100644 --- a/qds_sdk/commands.py +++ b/qds_sdk/commands.py @@ -243,6 +243,7 @@ class HiveCommand(Command): optparser.add_option("--print-logs", action="store_true", dest="print_logs", default=False, help="Fetch logs and print them to stderr.") + optparser.add_option("--retry", dest="retry", default=0, choices=[1,2,3], help="Number of retries for a job") @classmethod def parse(cls, args): @@ -418,6 +419,8 @@ class SparkCommand(Command): optparser.add_option("--print-logs", action="store_true", dest="print_logs", default=False, help="Fetch logs and print them to stderr.") + optparser.add_option("--retry", dest="retry", default=0, help="Number of retries") + @classmethod def validate_program(cls, options): bool_program = options.program is not None @@ -573,6 +576,7 @@ class PrestoCommand(Command): optparser.add_option("--print-logs", action="store_true", dest="print_logs", default=False, help="Fetch logs and print them to stderr.") + optparser.add_option("--retry", dest="retry", default=0, choices=[1,2,3], help="Number of retries for a job") @classmethod def parse(cls, args): @@ -646,6 +650,7 @@ class HadoopCommand(Command): optparser.add_option("--print-logs", action="store_true", dest="print_logs", default=False, help="Fetch logs and print them to stderr.") + optparser.add_option("--retry", dest="retry", default=0, choices=[1,2,3], help="Number of retries for a job") optparser.disable_interspersed_args() @@ -689,7 +694,7 @@ def parse(cls, args): "|".join(cls.subcmdlist)) parsed["sub_command"] = subcmd - parsed["sub_command_args"] = " ".join("'" + a + "'" for a in args) + parsed["sub_command_args"] = " ".join("'" + str(a) + "'" for a in args) return parsed @@ -816,6 +821,7 @@ class PigCommand(Command): optparser.add_option("--print-logs", action="store_true", dest="print_logs", default=False, help="Fetch logs and print them to stderr.") + optparser.add_option("--retry", dest="retry", choices=[1,2,3], default=0, help="Number of retries for a job") @classmethod def parse(cls, args): @@ -931,6 +937,7 @@ class DbExportCommand(Command): optparser.add_option("--print-logs", action="store_true", dest="print_logs", default=False, help="Fetch logs and print them to stderr.") + optparser.add_option("--retry", dest="retry", default=0, choices=[1,2,3], help="Number of retries for a job") @classmethod def parse(cls, args): @@ -1031,6 +1038,7 @@ class DbImportCommand(Command): optparser.add_option("--print-logs", action="store_true", dest="print_logs", default=False, help="Fetch logs and print them to stderr.") + optparser.add_option("--retry", dest="retry", default=0, choices=[1,2,3], help="Number of retries for a job") @classmethod def parse(cls, args): diff --git a/qds_sdk/scheduler.py b/qds_sdk/scheduler.py index f49e97fb..b4d414ab 100644 --- a/qds_sdk/scheduler.py +++ b/qds_sdk/scheduler.py @@ -265,7 +265,7 @@ def list_instances(self, page=None, per_page=None): if per_page is not None: page_attr.append("per_page=%s" % per_page) if page_attr: - url_path = "%s/instances?%s" % (self.element_path(args.id), "&".join(page_attr)) + url_path = "%s/instances?%s" % (self.element_path(self.id), "&".join(page_attr)) #Todo Page numbers are thrown away right now cmdjson = conn.get(url_path) cmdlist = [] @@ -277,5 +277,5 @@ def list_instances(self, page=None, per_page=None): def rerun(self, instance_id): conn = Qubole.agent() - url_path = self.element_path(id) + "/instances/" + instance_id + "/rerun" + url_path = self.element_path(self.id) + "/instances/%s/rerun" % instance_id return conn.post(url_path)['status'] diff --git a/qds_sdk/template.py b/qds_sdk/template.py new file mode 100644 index 00000000..8eddd58f --- /dev/null +++ b/qds_sdk/template.py @@ -0,0 +1,274 @@ +""" +The Template Module contains the base definition for Executing Templates +""" +import json +import logging +import os + +from argparse import ArgumentParser +from qds_sdk.qubole import Qubole +from qds_sdk.resource import Resource + +log = logging.getLogger("qds_template") + +from qds_sdk.commands import * + +class TemplateCmdLine: + """ + qds_sdk.TemplateCmdLine is the interface used for template related operation in qds.py + """ + + @staticmethod + def parsers(): + argparser = ArgumentParser(prog="qds.py template", description="Template Client for Qubole Data Service.") + subparsers = argparser.add_subparsers() + + #create + create = subparsers.add_parser("create", help="To Create a new Template") + create.add_argument("--data", dest="data",required=True, help="Path to JSON file with template details") + create.set_defaults(func=TemplateCmdLine.create) + + #edit + edit = subparsers.add_parser("edit", help="To Edit an existing Template") + edit.add_argument("--data", dest="data", required=True, help="Path to JSON file with template details") + edit.add_argument("--id", dest="id", required=True, help="Id for the Template") + edit.set_defaults(func=TemplateCmdLine.edit) + + #clone + clone = subparsers.add_parser("clone", help="To Clone an existing Template") + clone.add_argument("--id", dest="id",required=True, help="Id for the Template to be Cloned") + clone.add_argument("--data", dest="data", required=True, help="Path to JSON file with template details to override") + clone.set_defaults(func=TemplateCmdLine.clone) + + #view + view = subparsers.add_parser("view", help="To View an existing Template") + view.add_argument("--id", dest="id", required=True, help="Id for the Template") + view.set_defaults(func=TemplateCmdLine.view) + + #list + list = subparsers.add_parser("list", help="To List existing Templates") + list.add_argument("--per-page", dest="per_page", help="Number of items per page") + list.add_argument("--page", dest="page", help="Page Number") + list.set_defaults(func=TemplateCmdLine.list) + + #run + run = subparsers.add_parser("run", help="To Run Template and wait to print Result") + run.add_argument("--id", dest="id", required=True, help="Id of the template to run") + run.add_argument("--j", dest="data", required=True, help="Path to JSON file or json string with input field details") + run.set_defaults(func=TemplateCmdLine.execute) + + #submit + submit = subparsers.add_parser("submit", help="To Submit Template and get the command Id") + submit.add_argument("--id", dest="id", required=True, help="Id of the template to Submit") + submit.add_argument("--j", dest="data", required=True, help="Path to JSON file or json string with input field details") + submit.set_defaults(func=TemplateCmdLine.submit) + + return argparser + + @staticmethod + def run(args): + parser = TemplateCmdLine.parsers() + parsed = parser.parse_args(args) + return parsed.func(parsed) + + @staticmethod + def create(args): + with open(args.data) as f: + spec = json.load(f) + return Template.createTemplate(spec) + + @staticmethod + def edit(args): + with open(args.data) as f: + spec = json.load(f) + return Template.editTemplate(args.id, spec) + + @staticmethod + def clone(args): + with open(args.data) as f: + spec = json.load(f) + id = args.id + return Template.cloneTemplate(id, spec) + + @staticmethod + def submit(args): + spec = getSpec(args) + res = Template.submitTemplate(args.id, spec) + log.info("Submitted Template with Id: %s, Command Id: %s, CommandType: %s" % (args.id, res['id'], res['command_type'])) + return res + + @staticmethod + def execute(args): + spec = getSpec(args) + return Template.runTemplate(args.id, spec) + + @staticmethod + def view(args): + id = args.id + return Template.viewTemplate(id) + + @staticmethod + def list(args): + return Template.listTemplates(args) + +def getSpec(args): + if args.data is not None: + if os.path.isfile(args.data): + with open(args.data) as f: + spec = json.load(f) + else: + spec = json.loads(args.data) + if 'input_vars' in spec: + inputs = formatData(spec['input_vars']) + spec["input_vars"] = inputs + else: + spec = {} + return spec + + +def formatData(inputs): + res = [] + if len(inputs) != 0: + for obj in inputs: + o = {} + for key in obj: + o[key] = "'" + obj[key] + "'" + res.append(o) + return res + +class Template(Resource): + """ + qds_sdk.Template is the base Qubole Template class. + + it uses /command_templates endpoint + """ + + rest_entity_path = "command_templates" + + @staticmethod + def createTemplate(data): + """ + Create a new template. + + Args: + `data`: json data required for creating a template + Returns: + Dictionary containing the details of the template with its ID. + """ + conn = Qubole.agent() + return conn.post(Template.rest_entity_path, data) + + @staticmethod + def editTemplate(id, data): + """ + Edit an existing template. + + Args: + `id`: ID of the template to edit + `data`: json data to be updated + Returns: + Dictionary containing the updated details of the template. + """ + conn = Qubole.agent() + return conn.put(Template.element_path(id), data) + + @staticmethod + def cloneTemplate(id, data): + """ + Clone an existing template. + + Args: + `id`: ID of the template to be cloned + `data`: json data to override + Returns: + Dictionary containing the updated details of the template. + """ + conn = Qubole.agent() + path = str(id) + "/duplicate" + return conn.post(Template.element_path(path), data) + + @staticmethod + def viewTemplate(id): + """ + View an existing Template details. + + Args: + `id`: ID of the template to fetch + + Returns: + Dictionary containing the details of the template. + """ + conn = Qubole.agent() + return conn.get(Template.element_path(id)) + + @staticmethod + def submitTemplate(id, data): + """ + Submit an existing Template. + + Args: + `id`: ID of the template to submit + `data`: json data containing the input_vars + Returns: + Dictionary containing Command Object details. + """ + conn = Qubole.agent() + path = str(id) + "/run" + return conn.post(Template.element_path(path), data) + + @staticmethod + def runTemplate(id, data): + """ + Run an existing Template and waits for the Result. + Prints result to stdout. + + Args: + `id`: ID of the template to run + `data`: json data containing the input_vars + + Returns: + An integer as status (0: success, 1: failure) + """ + conn = Qubole.agent() + path = str(id) + "/run" + res = conn.post(Template.element_path(path), data) + cmdType = res['command_type'] + cmdId = res['id'] + cmdClass = eval(cmdType) + cmd = cmdClass.find(cmdId) + while not Command.is_done(cmd.status): + time.sleep(Qubole.poll_interval) + cmd = cmdClass.find(cmd.id) + return Template.getResult(cmdClass, cmd) + + @staticmethod + def getResult(cmdClass, cmd): + if Command.is_success(cmd.status): + log.info("Fetching results for %s, Id: %s" % (cmdClass.__name__, cmd.id)) + cmd.get_results(sys.stdout, delim='\t') + return 0 + else: + log.error("Cannot fetch results - command Id: %s failed with status: %s" % (cmd.id, cmd.status)) + return 1 + + @staticmethod + def listTemplates(args): + """ + Fetch existing Templates details. + + Args: + `args`: dictionary containing the value of page number and per-page value + Returns: + Dictionary containing paging_info and command_templates details + """ + conn = Qubole.agent() + url_path = Template.rest_entity_path + page_attr = [] + if args.page is not None: + page_attr.append("page=%s" % args.page) + if args.per_page is not None: + page_attr.append("per_page=%s" % args.per_page) + if page_attr: + url_path = "%s?%s" % (url_path, "&".join(page_attr)) + + return conn.get(url_path) \ No newline at end of file diff --git a/setup.py b/setup.py index 34bdac5a..c4a29efe 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,7 @@ def read(fname): setup( name="qds_sdk", - version="1.9.1", + version="1.9.2", author="Qubole", author_email="dev@qubole.com", description=("Python SDK for coding to the Qubole Data Service API"), diff --git a/tests/input_clone_template.json b/tests/input_clone_template.json new file mode 100644 index 00000000..d903555e --- /dev/null +++ b/tests/input_clone_template.json @@ -0,0 +1,12 @@ +{ + "name" : "template_name_clone", + "input_vars" : [ + { + "name":"col", + "default_value" : "'name'"}, + { + "name":"table", + "default_value" : "'account'" + } + ] +} diff --git a/tests/input_create_template.json b/tests/input_create_template.json new file mode 100644 index 00000000..e27b7783 --- /dev/null +++ b/tests/input_create_template.json @@ -0,0 +1,26 @@ +{ + "name" : "template_name", + "command_type" : "HiveCommand", + "command" : { + "name": null, + "tags": null, + "label": null, + "macros": null, + "hive_version": null, + "query": "select $col$ from $table$", + + "command_type": "HiveCommand", + "can_notify": false, + "script_location": null, + "sample_size": null + }, + "input_vars" : [ + { + "name":"col", + "default_value" : "'age'"}, + { + "name":"table", + "default_value" : "'user'" + } + ] +} diff --git a/tests/input_edit_template.json b/tests/input_edit_template.json new file mode 100644 index 00000000..2f7544a5 --- /dev/null +++ b/tests/input_edit_template.json @@ -0,0 +1,26 @@ +{ + "name" : "tempplate_new_name", + "command_type" : "HiveCommand", + "command" : { + "name": null, + "tags": null, + "label": null, + "macros": null, + "hive_version": null, + "query": "select $col$ from $table$", + + "command_type": "HiveCommand", + "can_notify": false, + "script_location": null, + "sample_size": null + }, + "input_vars" : [ + { + "name":"col", + "default_value" : "'age'"}, + { + "name":"table", + "default_value" : "'user'" + } + ] +} diff --git a/tests/input_run_template.json b/tests/input_run_template.json new file mode 100644 index 00000000..40cbdb2c --- /dev/null +++ b/tests/input_run_template.json @@ -0,0 +1,7 @@ +{ + "input_vars" : [ + { + "table" : "'accounts'" + } + ] +} diff --git a/tests/test_command.py b/tests/test_command.py index d2d3af1b..82c60de2 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -181,7 +181,7 @@ def test_done(self): class TestHiveCommand(QdsCliTestCase): def test_submit_query(self): - sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables'] + sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables', '--retry', 2] print_command() Connection._api_call = Mock(return_value={'id': 1234}) qds.main() @@ -195,7 +195,8 @@ def test_submit_query(self): 'query': 'show tables', 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 2}) def test_submit_query_with_hive_version(self): sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables', '--hive-version', '0.13'] @@ -212,7 +213,8 @@ def test_submit_query_with_hive_version(self): 'query': 'show tables', 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_script_location(self): sys.argv = ['qds.py', 'hivecmd', 'submit', '--script_location', 's3://bucket/path-to-script'] @@ -229,7 +231,8 @@ def test_submit_script_location(self): 'query': None, 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': 's3://bucket/path-to-script'}) + 'script_location': 's3://bucket/path-to-script', + 'retry': 0}) def test_submit_none(self): sys.argv = ['qds.py', 'hivecmd', 'submit'] @@ -260,7 +263,8 @@ def test_submit_macros(self): 'query': None, 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': 's3://bucket/path-to-script'}) + 'script_location': 's3://bucket/path-to-script', + 'retry': 0}) def test_submit_tags(self): sys.argv = ['qds.py', 'hivecmd', 'submit', '--script_location', 's3://bucket/path-to-script', @@ -278,7 +282,8 @@ def test_submit_tags(self): 'query': None, 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': 's3://bucket/path-to-script'}) + 'script_location': 's3://bucket/path-to-script', + 'retry': 0}) def test_submit_cluster_label(self): sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables', @@ -296,7 +301,8 @@ def test_submit_cluster_label(self): 'query': 'show tables', 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_name(self): sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables', @@ -314,7 +320,8 @@ def test_submit_name(self): 'query': 'show tables', 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_notify(self): sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables', @@ -332,7 +339,8 @@ def test_submit_notify(self): 'query': 'show tables', 'command_type': 'HiveCommand', 'can_notify': True, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_sample_size(self): sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables', @@ -350,7 +358,15 @@ def test_submit_sample_size(self): 'query': 'show tables', 'command_type': 'HiveCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) + + def test_retry_out_of_range(self): + sys.argv = ['qds.py', 'hivecmd', 'submit', '--query', 'show tables', + '--retry', 4] + print_command() + with self.assertRaises(qds_sdk.exception.ParseError): + qds.main() class TestSparkCommand(QdsCliTestCase): @@ -373,7 +389,8 @@ def test_submit_query(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_script_location_aws(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--script_location', 's3://bucket/path-to-script'] @@ -403,7 +420,8 @@ def test_submit_script_location_local_py(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_script_location_local_scala(self): with NamedTemporaryFile(suffix=".scala") as tmp: @@ -427,7 +445,8 @@ def test_submit_script_location_local_scala(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_script_location_local_java(self): with NamedTemporaryFile(suffix=".java") as tmp: @@ -460,13 +479,14 @@ def test_submit_script_location_local_R(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_script_location_local_sql(self): with NamedTemporaryFile(suffix=".sql") as tmp: tmp.write('show tables'.encode("utf8")) tmp.seek(0) - sys.argv = ['qds.py', 'sparkcmd', 'submit', '--script_location', tmp.name] + sys.argv = ['qds.py', 'sparkcmd', 'submit', '--script_location', tmp.name ] print_command() Connection._api_call = Mock(return_value={'id': 1234}) qds.main() @@ -484,7 +504,8 @@ def test_submit_script_location_local_sql(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_sql(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--sql', 'show dummy'] @@ -505,7 +526,8 @@ def test_submit_sql(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_sql_with_language(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--language','python', '--sql', 'show dummy'] @@ -566,7 +588,8 @@ def test_submit_macros(self): 'command_type': 'SparkCommand', 'cmdline': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_tags(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--language','scala','--program',"println(\"hello, world!\")", @@ -588,7 +611,8 @@ def test_submit_tags(self): 'user_program_arguments': None, 'cmdline': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_cluster_label(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--cmdline', '/usr/lib/spark/bin/spark-submit --class Test Test.jar', @@ -610,7 +634,8 @@ def test_submit_cluster_label(self): 'user_program_arguments': None, 'command_type': 'SparkCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_name(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--cmdline', '/usr/lib/spark/bin/spark-submit --class Test Test.jar', @@ -632,7 +657,8 @@ def test_submit_name(self): 'app_id': None, 'command_type': 'SparkCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_notify(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--cmdline', '/usr/lib/spark/bin/spark-submit --class Test Test.jar', @@ -654,7 +680,8 @@ def test_submit_notify(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': True, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_python_program(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--language','python','--program', 'print "hello, world!"'] @@ -675,7 +702,8 @@ def test_submit_python_program(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_user_program_arguments(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--language','scala','--program', @@ -699,7 +727,8 @@ def test_submit_user_program_arguments(self): 'arguments': '--class HelloWorld', 'user_program_arguments': 'world', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_scala_program(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--language','scala','--program', 'println("hello, world!")'] @@ -720,7 +749,8 @@ def test_submit_scala_program(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_R_program(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--language','R','--program', 'cat("hello, world!")'] @@ -741,7 +771,8 @@ def test_submit_R_program(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_program_to_app(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--language', 'scala', @@ -763,7 +794,8 @@ def test_submit_program_to_app(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_sql_to_app(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--sql', 'show tables', @@ -785,7 +817,8 @@ def test_submit_sql_to_app(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_script_location_local_py_to_app(self): with NamedTemporaryFile(suffix=".py") as tmp: @@ -810,7 +843,8 @@ def test_submit_script_location_local_py_to_app(self): 'arguments': None, 'user_program_arguments': None, 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_cmdline_to_app(self): sys.argv = ['qds.py', 'sparkcmd', 'submit', '--cmdline', @@ -824,7 +858,7 @@ def test_submit_cmdline_to_app(self): class TestPrestoCommand(QdsCliTestCase): def test_submit_query(self): - sys.argv = ['qds.py', 'prestocmd', 'submit', '--query', 'show tables'] + sys.argv = ['qds.py', 'prestocmd', 'submit', '--query', 'show tables', '--retry', 1] print_command() Connection._api_call = Mock(return_value={'id': 1234}) qds.main() @@ -836,7 +870,8 @@ def test_submit_query(self): 'query': 'show tables', 'command_type': 'PrestoCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 1}) def test_submit_script_location(self): sys.argv = ['qds.py', 'prestocmd', 'submit', '--script_location', 's3://bucket/path-to-script'] @@ -851,7 +886,8 @@ def test_submit_script_location(self): 'query': None, 'command_type': 'PrestoCommand', 'can_notify': False, - 'script_location': 's3://bucket/path-to-script'}) + 'script_location': 's3://bucket/path-to-script', + 'retry': 0}) def test_submit_none(self): sys.argv = ['qds.py', 'prestocmd', 'submit'] @@ -880,7 +916,8 @@ def test_submit_macros(self): 'query': None, 'command_type': 'PrestoCommand', 'can_notify': False, - 'script_location': 's3://bucket/path-to-script'}) + 'script_location': 's3://bucket/path-to-script', + 'retry': 0}) def test_submit_tags(self): sys.argv = ['qds.py', 'prestocmd', 'submit', '--script_location', 's3://bucket/path-to-script', @@ -896,7 +933,8 @@ def test_submit_tags(self): 'query': None, 'command_type': 'PrestoCommand', 'can_notify': False, - 'script_location': 's3://bucket/path-to-script'}) + 'script_location': 's3://bucket/path-to-script', + 'retry': 0}) def test_submit_cluster_label(self): sys.argv = ['qds.py', 'prestocmd', 'submit', '--query', 'show tables', @@ -912,7 +950,8 @@ def test_submit_cluster_label(self): 'query': 'show tables', 'command_type': 'PrestoCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_name(self): sys.argv = ['qds.py', 'prestocmd', 'submit', '--query', 'show tables', @@ -928,7 +967,8 @@ def test_submit_name(self): 'query': 'show tables', 'command_type': 'PrestoCommand', 'can_notify': False, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) def test_submit_notify(self): sys.argv = ['qds.py', 'prestocmd', 'submit', '--query', 'show tables', @@ -944,8 +984,15 @@ def test_submit_notify(self): 'query': 'show tables', 'command_type': 'PrestoCommand', 'can_notify': True, - 'script_location': None}) + 'script_location': None, + 'retry': 0}) + def test_retry_out_of_range(self): + sys.argv = ['qds.py', 'prestocmd', 'submit', '--query', 'show tables', + '--retry', 5] + print_command() + with self.assertRaises(qds_sdk.exception.ParseError): + qds.main() class TestHadoopCommand(QdsCliTestCase): @@ -1087,7 +1134,7 @@ class TestDbExportCommand(QdsCliTestCase): def test_submit_command(self): sys.argv = ['qds.py', 'dbexportcmd', 'submit', '--mode', '1', '--dbtap_id', '1', - '--db_table', 'mydbtable', '--hive_table', 'myhivetable'] + '--db_table', 'mydbtable', '--hive_table', 'myhivetable', '--retry', 3] print_command() Connection._api_call = Mock(return_value={'id': 1234}) qds.main() @@ -1104,7 +1151,8 @@ def test_submit_command(self): 'command_type': 'DbExportCommand', 'dbtap_id': '1', 'can_notify': False, - 'db_update_mode': None}) + 'db_update_mode': None, + 'retry': 3}) def test_submit_fail_with_no_parameters(self): sys.argv = ['qds.py', 'dbexportcmd', 'submit'] @@ -1131,7 +1179,8 @@ def test_submit_with_notify(self): 'command_type': 'DbExportCommand', 'dbtap_id': '1', 'can_notify': True, - 'db_update_mode': None}) + 'db_update_mode': None, + 'retry': 0}) def test_submit_with_name(self): sys.argv = ['qds.py', 'dbexportcmd', 'submit', '--mode', '1', '--dbtap_id', '1', @@ -1152,7 +1201,8 @@ def test_submit_with_name(self): 'command_type': 'DbExportCommand', 'dbtap_id': '1', 'can_notify': False, - 'db_update_mode': None}) + 'db_update_mode': None, + 'retry': 0}) def test_submit_with_update_mode_and_keys(self): sys.argv = ['qds.py', 'dbexportcmd', 'submit', '--mode', '1', '--dbtap_id', '1', @@ -1174,7 +1224,8 @@ def test_submit_with_update_mode_and_keys(self): 'command_type': 'DbExportCommand', 'dbtap_id': '1', 'can_notify': False, - 'db_update_mode': 'updateonly'}) + 'db_update_mode': 'updateonly', + 'retry': 0}) def test_submit_with_mode_2(self): sys.argv = ['qds.py', 'dbexportcmd', 'submit', '--mode', '2', '--dbtap_id', '1', @@ -1196,7 +1247,15 @@ def test_submit_with_mode_2(self): 'command_type': 'DbExportCommand', 'dbtap_id': '1', 'can_notify': False, - 'db_update_mode': None}) + 'db_update_mode': None, + 'retry': 0}) + + def test_retry_out_of_range(self): + sys.argv = ['qds.py', 'dbexportcmd', 'submit', '--mode', '1', '--dbtap_id', '1', + '--db_table', 'mydbtable', '--hive_table', 'myhivetable', '--retry', 5] + print_command() + with self.assertRaises(qds_sdk.exception.ParseError): + qds.main() class TestDbImportCommand(QdsCliTestCase): @@ -1205,7 +1264,7 @@ class TestDbImportCommand(QdsCliTestCase): # The test cases might give false positivies def test_submit_command(self): sys.argv = ['qds.py', 'dbimportcmd', 'submit', '--mode', '1', '--dbtap_id', '1', - '--db_table', 'mydbtable', '--hive_table', 'myhivetable'] + '--db_table', 'mydbtable', '--hive_table', 'myhivetable', '--retry', 2] print_command() Connection._api_call = Mock(return_value={'id': 1234}) qds.main() @@ -1222,8 +1281,15 @@ def test_submit_command(self): 'can_notify': False, 'hive_table': 'myhivetable', 'db_table': 'mydbtable', - 'db_extract_query': None}) + 'db_extract_query': None, + 'retry': 2}) + def test_retry_out_of_range(self): + sys.argv = ['qds.py', 'dbimportcmd', 'submit', '--mode', '1', '--dbtap_id', '1', + '--db_table', 'mydbtable', '--hive_table', 'myhivetable', '--retry', 6] + print_command() + with self.assertRaises(qds_sdk.exception.ParseError): + qds.main() class TestDbTapQueryCommand(QdsCliTestCase): diff --git a/tests/test_template.py b/tests/test_template.py new file mode 100644 index 00000000..cdb53da1 --- /dev/null +++ b/tests/test_template.py @@ -0,0 +1,156 @@ +from __future__ import print_function + +import sys +import os + +if sys.version_info > (2,7,0): + import unittest +else: + import unittest2 as unittest + +from mock import * + +sys.path.append(os.path.join(os.path.dirname(__file__), '../bin')) +import qds + +from qds_sdk.connection import Connection +from qds_sdk.template import * +from test_base import print_command +from test_base import QdsCliTestCase + +class TestTemplateCheck(QdsCliTestCase): + def test_create_template(self): + file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'input_create_template.json') + sys.argv = ['qds.py', 'template', 'create', '--data', file_path] + print_command() + Connection._api_call = Mock() + qds.main() + with open(file_path) as f: + data = json.load(f) + Connection._api_call.assert_called_with("POST", "command_templates", data) + + def test_create_template_without_data(self): + sys.argv = ['qds.py', 'template', 'create'] + print_command() + with self.assertRaises(SystemExit): + qds.main() + + def test_edit_template(self): + file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'input_edit_template.json') + sys.argv = ['qds.py', 'template', 'edit', '--id', '12', '--data', file_path] + print_command() + Connection._api_call = Mock() + qds.main() + with open(file_path) as f: + data = json.load(f) + Connection._api_call.assert_called_with("PUT", "command_templates/12", data) + + def test_edit_template_without_data(self): + sys.argv = ['qds.py', 'template', 'edit', '--id', '12'] + print_command() + with self.assertRaises(SystemExit): + qds.main() + + def test_view_template(self): + sys.argv = ['qds.py', 'template', 'view', '--id', '12'] + print_command() + Connection._api_call = Mock() + qds.main() + Connection._api_call.assert_called_with("GET", "command_templates/12", params=None) + + def test_view_template_without_id(self): + sys.argv = ['qds.py', 'template', 'view'] + print_command() + with self.assertRaises(SystemExit): + qds.main() + + def test_clone_template(self): + file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'input_clone_template.json') + sys.argv = ['qds.py', 'template', 'clone', '--id', '12', '--data', file_path] + print_command() + Connection._api_call = Mock() + qds.main() + with open(file_path) as f: + data = json.load(f) + Connection._api_call.assert_called_with("POST", "command_templates/12/duplicate", data) + + def test_clone_template_without_id_data(self): + sys.argv = ['qds.py', 'template', 'clone'] + print_command() + with self.assertRaises(SystemExit): + qds.main() + + def test_list_template(self): + sys.argv = ['qds.py', 'template', 'list', '--page', '1', '--per-page','10'] + print_command() + Connection._api_call = Mock() + qds.main() + Connection._api_call.assert_called_with("GET", "command_templates?page=1&per_page=10", params=None) + + def test_submit_template(self): + file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'input_run_template.json') + sys.argv = ['qds.py', 'template', 'submit', '--id', '14', '--j', file_path] + print_command() + Connection._api_call = Mock() + Connection._api_call.side_effect = submit_actions_side_effect + qds.main() + with open(file_path) as f: + data = json.load(f) + Connection._api_call.assert_called_with("POST", "command_templates/14/run", data) + + def test_submit_template_without_id_data(self): + sys.argv = ['qds.py', 'template', 'submit'] + print_command() + with self.assertRaises(SystemExit): + qds.main() + + def test_submit_template_with_inline_json(self): + sys.argv = ['qds.py', 'template', 'submit', '--id', '14', '--j', '{"input_vars" : [{"table" : "accounts"}]}'] + print_command() + Connection._api_call = Mock() + Connection._api_call.side_effect = submit_actions_side_effect + qds.main() + data = {'input_vars': [{'table': "'accounts'"}]} + Connection._api_call.assert_called_with("POST", "command_templates/14/run", data) + + def test_run_template(self): + file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'input_run_template.json') + sys.argv = ['qds.py', 'template', 'run','--id', '14','--j',file_path] + print_command() + Connection._api_call = Mock() + Connection._api_call.side_effect = submit_actions_side_effect + HiveCommand.find = Mock() + HiveCommand.find.side_effect = find_command_side_effect + HiveCommand.get_results = Mock() + qds.main() + with open(file_path) as f: + data = json.load(f) + Connection._api_call.assert_called_with("POST", "command_templates/14/run", data) + + def test_run_template_without_id_data(self): + sys.argv = ['qds.py', 'template', 'run'] + print_command() + with self.assertRaises(SystemExit): + qds.main() + + def test_invalid_template_operation(self): + sys.argv = ['qds.py', 'template', 'load', '--id','12'] + print_command() + with self.assertRaises(SystemExit): + qds.main() + +def submit_actions_side_effect(*args, **kwargs): + res = { + "id" : 122, + "command_type" : 'HiveCommand' + } + return res + +def find_command_side_effect(id): + cmd = HiveCommand() + cmd.status = 'done' + cmd.id = 122 + return cmd + +if __name__ == '__main__': + unittest.main()