diff options
Diffstat (limited to 'python/refcat/cli.py')
| -rw-r--r-- | python/refcat/cli.py | 251 |
1 files changed, 251 insertions, 0 deletions
diff --git a/python/refcat/cli.py b/python/refcat/cli.py new file mode 100644 index 0000000..ca1cab7 --- /dev/null +++ b/python/refcat/cli.py @@ -0,0 +1,251 @@ +""" + ____ __ + ________ / __/________ _/ /_ + / ___/ _ \/ /_/ ___/ __ `/ __/ + / / / __/ __/ /__/ /_/ / /_ +/_/ \___/_/ \___/\__,_/\__/ + +Command line entry point for running various data tasks. + + $ refcat.pyz [COMMAND | TASK] [OPTIONS] + +Commands: ls, ll, deps, tasks, files, config, cat, completion + +To install completion run: + + $ source <(refcat.pyz completion) +""" + +import logging +import subprocess +import sys +import tempfile + +import gluish +import luigi +from luigi.parameter import MissingParameterException +from luigi.task import Register +from luigi.task_register import TaskClassNotFoundException + +from refcat import __version__ +from refcat.settings import LOGGING_CONF_FILE, settings +from refcat.deps import dump_deps +from refcat.tasks import * +from refcat.utils import columnize + +# XXX: get rid of gluish dep, include them in refcat +suppress_task_names = [ + "Available", + "BaseTask", + "Config", + "Executable", + "ExternalTask", + "FillSolrIndex", + "GitCloneRepository", + "GitUpdateRepository", + "MockTask", + "RangeBase", + "RangeByMinutes", + "RangeByMinutesBase", + "RangeDaily", + "RangeDailyBase", + "RangeHourly", + "RangeHourlyBase", + "RangeMonthly", + "Task", + "TestNotificationsTask", + "WrapperTask", +] + + +def tasks(): + """ + Print task name. + """ + names = [name for name in sorted(Register.task_names()) if name not in suppress_task_names and not name.islower()] + for name in names: + print(name) + + +def files(): + """ + Print task name and file. + """ + names = [name for name in sorted(Register.task_names()) if name not in suppress_task_names and not name.islower()] + for name in names: + klass = Register.get_task_cls(name) + try: + print("{:30s}{}".format(name, klass().output().path)) + except AttributeError: + print("{:30s}".format(name)) + + +def cat(*args): + """ + Inspect file. + """ + if len(args) == 0: + raise ValueError("ls failed: task name required") + task = find_task(args[0]) + try: + filename = task().output().path + if filename.endswith(".zst"): + subprocess.run(["zstdcat", "-T0", filename]) + elif filename.endswith(".gz"): + subprocess.run(["zcat", filename]) + else: + subprocess.run(["cat", filename]) + except FileNotFoundError: + print("file not found: {}".format(filename), file=sys.stderr) + except AttributeError: + print("{} is most likely not a task object".format(name), file=sys.stderr) + + +def ls(*args): + """ + Print output filename for task. + """ + if len(args) == 0: + raise ValueError("ls failed: task name required") + task = find_task(args[0]) + print(task().output().path) + + +def ll(*args): + """ + Long output. + """ + if len(args) == 0: + raise ValueError("ls failed: task name required") + task = find_task(args[0]) + try: + filename = task().output().path + subprocess.run(["ls", "-lah", filename]) + except FileNotFoundError: + print("file not found: {}".format(filename), file=sys.stderr) + except AttributeError: + print("{} is most likely not a task object".format(name), file=sys.stderr) + + +def deps(*args): + """ + Render task dependencies. + """ + if len(args) == 0: + raise ValueError("deps failed: task name required") + task = find_task(args[0]) + dump_deps(task()) + + +def config(): + """ + Dump config to stdout. + """ + with open(settings.settings_file) as f: + print(f.read()) + + +def completion(): + snippet = """ +_refcat_completion() +{ + local cur prev + COMPREPLY=() + cur=${COMP_WORDS[COMP_CWORD]} + prev=${COMP_WORDS[COMP_CWORD-1]} + + if [ $COMP_CWORD -eq 1 ]; then + COMPREPLY=( $(compgen -W "cat completion config deps files ll ls tasks" -- $cur) ) + elif [ $COMP_CWORD -eq 2 ]; then + case "$prev" in + "cat") + COMPREPLY=( $(compgen -W "$(refcat.pyz tasks)" -- $cur) ) + ;; + "ll") + COMPREPLY=( $(compgen -W "$(refcat.pyz tasks)" -- $cur) ) + ;; + "ls") + COMPREPLY=( $(compgen -W "$(refcat.pyz tasks)" -- $cur) ) + ;; + *) + ;; + esac + fi + return 0 +} + +complete -F _refcat_completion "refcat.pyz" +complete -F _refcat_completion "refcat" + """ + print(snippet) + + +def find_task(name): + """ + Note: return task class, not instance. + """ + task = globals().get(name) + if task is None: + raise RuntimeError("no such task: {}".format(name)) + if not isinstance(task(), luigi.Task): + raise RuntimeError("not a task: {}".format(name)) + return task + + +def main(): + """ + The main command line entry point. + """ + if os.path.exists(LOGGING_CONF_FILE): + logging.config.fileConfig(LOGGING_CONF_FILE) + + if settings.TMPDIR: + # This might not be passed on to subprocesses. + tempfile.tempdir = settings.TMPDIR + + # Explicitly name valid subcommands. + sub_commands = [ + "cat", + "completion", + "config", + "deps", + "files", + "ll", + "ls", + "tasks", + ] + + if len(sys.argv) >= 2 and sys.argv[1] in sub_commands: + try: + func = globals()[sys.argv[1]] + if callable(func): + func(*sys.argv[2:]) + else: + print("not implemented: {}".format(sys.argv[1])) + except KeyError: + print("sub-command not implemented: {}".format(sys.argv[1]), file=sys.stderr) + except (AttributeError, ValueError, RuntimeError) as exc: + print(exc, file=sys.stderr) + sys.exit(1) + else: + sys.exit(0) + elif len(sys.argv) == 1: + print(__doc__) + print("VERSION {}".format(__version__)) + print("SETTINGS {}".format(settings.settings_file)) + print("BASE {}".format(settings.BASE)) + print("TMPDIR {}".format(settings.TMPDIR)) + print() + names = [name for name in sorted(Register.task_names()) if name not in suppress_task_names and not name.islower()] + print(columnize(names)) + sys.exit(0) + + # If we found not subcommand, assume task name. + try: + luigi.run(local_scheduler=True) # XXX: add logging_conf_file + except MissingParameterException as err: + print('missing parameter: %s' % err, file=sys.stderr) + sys.exit(1) + except TaskClassNotFoundException as err: + print(err, file=sys.stderr) + sys.exit(1) |