def _upload_zip(session, name, path, create=False, archive_name=None):
"""Upload zip to project in Azkaban.
:param session: Remote Azkaban session.
:param name: Project name
:param path: Path to zip file.
:param create: Create project if it doesn't exist.
:param archive_name: Optional zip file name (used by Azkaban).
"""
def _callback(cur_bytes, tot_bytes, file_index, _stdout=sys.stdout):
"""Callback for streaming upload.
:param cur_bytes: Total bytes uploaded so far.
:param tot_bytes: Total bytes to be uploaded.
:param file_index: (0-based) index of the file currently uploaded.
:param _stdout: Performance caching.
"""
if cur_bytes != tot_bytes:
_stdout.write('Uploading project: %.1f%%\r' %
(100. * cur_bytes / tot_bytes, ))
else:
_stdout.write('Validating project... \r')
_stdout.flush()
while True:
try:
res = session.upload_project(
name=name,
path=path,
archive_name=archive_name,
callback=_callback,
)
except AzkabanError as err:
if create and str(err).endswith("doesn't exist."):
session.create_project(name, name)
else:
raise err
except HTTPError as err:
# See https://github.com/mtth/azkaban/pull/34 for more context on why the
# logic below is necessary.
code = err.response.status_code
if code == 400:
raise AzkabanError(
"Failed to upload project (%s HTTP error). Check that the project "
"is not locked, exists, and that your user has write permissions."
% (code, ))
elif code == 401:
raise AzkabanError(
"Not authorized to upload project (%s HTTP error). Check that your "
"user has write permissions." % (code, ))
elif code == 410:
session.create_project(name, name)
else:
raise err
else:
return res
def view_log(_execution, _job, _url, _alias):
"""View workflow or job execution logs."""
session = _get_session(_url, _alias)
exc = Execution(session, _execution)
logs = exc.job_logs(_job[0]) if _job else exc.logs()
try:
for line in logs:
sys.stdout.write('%s\n' % (line.encode('utf-8'), ))
except HTTPError:
# Azkaban responds with 500 if the execution or job isn't found
if _job:
raise AzkabanError('Execution %s and/or job %s not found.',
_execution, _job)
else:
raise AzkabanError('Execution %s not found.', _execution)
def _load_project(_project):
"""Resolve project from CLI argument.
:param _project: `--project` argument.
"""
try:
name, project = _parse_project(_project, require_project=True)
except ImportError:
raise AzkabanError(
'This command requires a project configuration module which was not '
'found.\nYou can specify another location using the `--project` option.'
)
else:
return project
def _parse_option(_option):
"""Parse `--option` argument.
:param _option: `--option` argument.
Returns a dictionary.
"""
paths = (opt for opt in _option if not '=' in opt)
opts = read_properties(*paths)
try:
opts.update(dict(s.split('=', 1) for s in _option if '=' in s))
except ValueError:
raise AzkabanError('Invalid `--option` flag.')
return opts
def main(argv=None):
"""Entry point."""
# enable general logging
logger = lg.getLogger()
logger.setLevel(lg.DEBUG)
handler = Config().get_file_handler('azkaban')
if handler:
logger.addHandler(handler)
# capture pesky unverified requests warnings
suppress_urllib_warnings()
# parse arguments
argv = argv or sys.argv[1:]
_logger.debug('Running command %r from %r.', ' '.join(argv), os.getcwd())
args = docopt(__doc__, version=__version__)
CLI_ARGS.update(args)
# do things
if args['--log']:
if handler:
sys.stdout.write('%s\n' % (handler.baseFilename, ))
else:
raise AzkabanError('No log file active.')
elif args['build']:
build_project(
_load_project(args['--project']),
**_forward(
args,
['ZIP', '--url', '--alias', '--replace', '--create', '--option']
)
)
elif args['log']:
view_log(
**_forward(args, ['EXECUTION', 'JOB', '--url', '--alias'])
)
elif args['info']:
view_info(
_load_project(args['--project']),
**_forward(args, ['--files', '--option', 'JOB', '--include-properties'])
)
elif args['run']:
run_workflow(
_get_project_name(args['--project']),
**_forward(
args,
[
'FLOW', 'JOB', '--bounce', '--url', '--alias', '--kill', '--email',
'--option',
]
)
)
elif args['schedule']:
schedule_workflow(
_get_project_name(args['--project']),
**_forward(
args,
[
'FLOW', 'JOB', '--bounce', '--url', '--alias', '--kill',
'--email', '--option', '--date', '--time', '--span'
]
)
)
elif args['upload']:
upload_project(
_get_project_name(args['--project']),
**_forward(args, ['ZIP', '--create', '--url', '--alias'])
)
def _parse_project(_project, require_project=False):
"""Parse `--project` argument into `(name, project)`.
:param _project: `--project` argument.
:param require_project: Fail if we fail to load the project.
Note that `name` is guaranteed to be non-`None` (this function will throw an
exception otherwise) but `project` can be.
The rules are as follows:
+ If at least one `':'` is found in `_project` then the rightmost one is
interpreted as delimitor between the path to the module and the project
name.
+ Else:
+ We first try to interpret `_project` as a module path and find a unique
project inside.
+ If the above attempt raises an `ImportError`, we interpret it as a name.
"""
default_module = Config().get_option('azkaban', 'default.project', 'jobs')
projects = {}
_project = _project or default_module
if ':' in _project:
# unambiguous case
path, name = _project.rsplit(':', 1)
try:
projects = Project.load(path or default_module)
# adding the default here lets options like `-p :name` work as intended
except ImportError:
pass
else:
# the option could be a name or module
try:
# try first as a module
projects = Project.load(_project)
except ImportError:
# if that fails, try as a name: load the default module and look there
name = _project
try:
projects = Project.load(default_module)
except ImportError:
pass
else:
name = None
if name:
if name in projects:
return name, projects[name]
elif projects:
# harder consistency requirement
raise AzkabanError(
'Project %r not found. Available projects: %s\n'
'You can also specify another location using the `--project` option.'
% (name, ', '.join(projects))
)
elif require_project:
raise AzkabanError(
'This command requires a project configuration module.\n'
'You can specify another location using the `--project` option.'
)
else:
return name, None
else:
if not projects:
raise AzkabanError(
'No registered project found in %r.\n'
'You can also specify another location using the `--project` option.'
% (_project, )
)
elif len(projects) > 1:
raise AzkabanError(
'Multiple registered projects found: %s\n'
'You can use the `--project` option to disambiguate.'
% (', '.join(projects), )
)
else:
return projects.popitem()
def _parse_project(_project, require_project=False):
"""Parse `--project` argument into `(name, project)`.
:param _project: `--project` argument.
:param require_project: Fail if we fail to load the project.
Note that `name` is guaranteed to be non-`None` (this function will throw an
exception otherwise) but `project` can be.
The rules are as follows:
+ If at least one `':'` is found in `_project` then the rightmost one is
interpreted as delimitor between the path to the module and the project
name.
+ Else:
+ We first try to interpret `_project` as a module path and find a unique
project inside.
+ If the above attempt raises an `ImportError`, we interpret it as a name.
"""
default_project = Config().get_option('azkaban', 'default.project', 'jobs')
exceptions = {}
projects = {}
def try_load(path):
try:
projects.update(Project.load(path))
return True
except Exception:
exceptions[path] = format_exc()
return False
_project = _project or default_project
if ':' in _project:
# unambiguous case
path, name = _project.rsplit(':', 1)
if ':' in default_project:
try_load(Project.load(path))
else:
# adding the default here lets options like `-p :name` work as intended
try_load(path or default_project)
else:
# the option could be a name or module
if not try_load(_project): # try first as a module
# if that fails, try as a name
name = _project
if not ':' in default_project:
path = default_project
else:
path = default_project.rsplit(':', 1)[0]
# if the default project could be a mdule, try loading it
try_load(path)
else:
name = None
path = _project
if exceptions:
footer = '\nErrors occurred while loading the following modules:\n'
for t in exceptions.items():
footer += '\n> %r\n\n%s' % t
else:
footer = ''
if name:
if name in projects:
return name, projects[name]
elif projects:
# harder consistency requirement
raise AzkabanError(
'Project %r not found. Available projects: %s\n'
'You can also specify another location using the `--project` option.'
'%s'
% (name, ', '.join(projects), footer)
)
elif require_project:
raise AzkabanError(
'This command requires a project configuration module.\n'
'You can specify another location using the `--project` option.'
'%s'
% (footer, )
)
else:
return name, None
else:
if not projects:
raise AzkabanError(
'No registered project found in %r.\n'
'You can specify another location using the `--project` option.'
'%s'
% (path, footer)
)
elif len(projects) > 1:
raise AzkabanError(
'Multiple registered projects found: %s\n'
'You can use the `--project` option to disambiguate.'
'%s'
% (', '.join(projects), footer)
)
else:
return projects.popitem()