Skip to content

Commit

Permalink
Merge pull request #23 from mogproject/topic-upload-with-#20
Browse files Browse the repository at this point in the history 
implement --upload-with option closes #20
  • Loading branch information
@mogproject
mogproject committed Dec 30, 2015
2 parents 8815610 + f10a2b1 commit eb3da0c
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 17 deletions.
50 changes: 33 additions & 17 deletions src/color_ssh/color_ssh.py
View file
Expand Up @@ -21,6 +21,9 @@ class Setting(object):
' %prog [options...] -H "[user@]hostname [[user@]hostname]...]" command'
])
DEFAULT_PARALLELISM = 32
CMD_SSH = str('ssh')
CMD_UPLOAD = [str('rsync'), str('-a')]
CMD_MKDIR = [str('mkdir'), str('-p')]

def __init__(self, parallelism=None, tasks=None):
self.parallelism = parallelism
Expand All @@ -39,7 +42,7 @@ def parse_args(self, argv, stdout=io2bytes(sys.stdout)):
help='label name'
)
parser.add_option(
'--ssh', dest='ssh', default=str('ssh'), type='string', metavar='SSH',
'--ssh', dest='ssh', default=self.CMD_SSH, type='string', metavar='SSH',
help='override ssh command line string'
)
parser.add_option(
Expand All @@ -62,11 +65,16 @@ def parse_args(self, argv, stdout=io2bytes(sys.stdout)):
'--upload', dest='upload', default=False, action='store_true',
help='upload files before executing a command (all args are regarded as paths)'
)
parser.add_option(
'--upload-with', dest='upload_with', default=None, type='string', metavar='PATH',
help='file paths to be uploaded before executing a command'
)

option, args = parser.parse_args(argv[1:])
hosts = self._load_hosts(option.host_file) + (option.host_string.split() if option.host_string else [])

if len(args) < (1 if hosts else 2):
print(option.__dict__)
stdout.write(arg2bytes(parser.format_help().encode('utf-8')))
parser.exit(2)

Expand All @@ -77,31 +85,26 @@ def parse_args(self, argv, stdout=io2bytes(sys.stdout)):
# parse hosts
parsed_hosts = [self._parse_host(h) for h in hosts]

# parse upload-with option
upload_with = [] if option.upload_with is None else shlex.split(option.upload_with)

tasks = []
if option.distribute:
# distribute args
dist_prefix = shlex.split(option.distribute)
d = distribute(len(hosts), args)
for i, (user, host, port) in enumerate(parsed_hosts):
if d[i]:
setup_commands = []
if option.upload:
# create directories
dirs = list(set(x for x in [os.path.dirname(arg) for arg in d[i]] if x != '' and x != '.'))
if dirs:
setup_commands.append(
self._ssh_args(option.ssh, user, host, port) + [str('mkdir'), str('-p')] + dirs
)

# upload files before executing main commands
setup_commands.extend([self._scp_args(str('rsync -a'), user, host, port, arg) for arg in d[i]])

upload_paths = upload_with + d[i] if option.upload else []
label = option.label or host
ssh_args = self._ssh_args(option.ssh, user, host, port)
tasks.append((label, ssh_args + dist_prefix + d[i], setup_commands))
tasks.append((label, ssh_args + dist_prefix + d[i],
self._build_upload_commands(user, host, port, option.ssh, upload_paths)))
else:
for user, host, port in parsed_hosts:
tasks.append((option.label or host, self._ssh_args(option.ssh, user, host, port) + args, []))
tasks.append((option.label or host,
self._ssh_args(option.ssh, user, host, port) + args,
self._build_upload_commands(user, host, port, option.ssh, upload_with)))

self.parallelism = option.parallelism
self.tasks = tasks
Expand Down Expand Up @@ -140,12 +143,25 @@ def _ssh_args(ssh_cmd, user, host, port):
[] if port is None else [str('-p'), port]) + [Setting._build_host_string(user, host)]

@staticmethod
def _scp_args(scp_cmd, user, host, port, path):
return shlex.split(scp_cmd) + ([] if port is None else [str('-P'), port]) + [
def _upload_args(user, host, port, path):
return Setting.CMD_UPLOAD + ([] if port is None else [str('-P'), port]) + [
path,
Setting._build_host_string(user, host) + str(':') + path
]

@staticmethod
def _build_upload_commands(user, host, port, ssh_cmd, paths):
# create directories
dirs = list(set(x for x in [os.path.dirname(path) for path in paths] if x != '' and x != '.'))

ret = []
if dirs:
ret.append(Setting._ssh_args(ssh_cmd, user, host, port) + Setting.CMD_MKDIR + sorted(dirs))

# upload files
ret.extend([Setting._upload_args(user, host, port, path) for path in paths])
return ret


def run_task(args):
label, command, setup_commands = args
Expand Down
29 changes: 29 additions & 0 deletions tests/color_ssh/test_color_ssh.py
View file
Expand Up @@ -115,6 +115,35 @@ def test_parse_args(self):
[['rsync', '-a', 'z', 'root@server-12:z']]),
])

# upload-with
self._check(self._parse(['--upload-with=dir1/x', 'server-1', 'pwd']),
[('server-1', ['ssh', 'server-1', 'pwd'], [
['ssh', 'server-1', 'mkdir', '-p', 'dir1'],
['rsync', '-a', 'dir1/x', 'server-1:dir1/x'],
])])

self._check(
self._parse([
'--upload-with', 'dir2/c dir2/d dir3/e',
'-H', 'server-11 root@server-12', '--distribute', 'echo "foo bar"', '--upload', 'dir1/x', 'dir1/y', 'z'
]), [
('server-11', ['ssh', 'server-11', 'echo', 'foo bar', 'dir1/x', 'dir1/y'], [
['ssh', 'server-11', 'mkdir', '-p', 'dir1', 'dir2', 'dir3'],
['rsync', '-a', 'dir2/c', 'server-11:dir2/c'],
['rsync', '-a', 'dir2/d', 'server-11:dir2/d'],
['rsync', '-a', 'dir3/e', 'server-11:dir3/e'],
['rsync', '-a', 'dir1/x', 'server-11:dir1/x'],
['rsync', '-a', 'dir1/y', 'server-11:dir1/y']
]),
('server-12', ['ssh', 'root@server-12', 'echo', 'foo bar', 'z'], [
['ssh', 'root@server-12', 'mkdir', '-p', 'dir2', 'dir3'],
['rsync', '-a', 'dir2/c', 'root@server-12:dir2/c'],
['rsync', '-a', 'dir2/d', 'root@server-12:dir2/d'],
['rsync', '-a', 'dir3/e', 'root@server-12:dir3/e'],
['rsync', '-a', 'z', 'root@server-12:z']
]),
])

def test_parse_args_error(self):
with self.withBytesOutput() as (out, err):
self.assertSystemExit(2, Setting().parse_args, ['color-ssh'], out)
Expand Down

0 comments on commit eb3da0c

Please sign in to comment.