Merge lp:~abentley/workspace-runner/s3-artifacts into lp:workspace-runner

[Aaron Bentley]

Upload artifacts to s3 from workspace runner.

This branch applies the upload_artifacts script.

It updates workspace_run to accept artifact_prefix and --s3-config. It uses the access-key and secret key from the s3config file, plus the artifact prefix, plus the 'bucket' and 'artifacts' config values to generate a config file. That file is then used to remotely run upload_artifacts.

[lp:~abentley/workspace-runner/s3-artifacts]

37. By Aaron Bentley on 2015-06-30

Merged trunk into s3-artifacts.

[Curtis Hovey (sinzui)]

Thank you. I have a question/suggestion inline.

review: Approve (code)

[Aaron Bentley (abentley)]

AFAICT, the main advantage of SafeConfigParser over RawConfigParser is that SafeConfigParser supports interpolation. I didn't think interpolation was useful for this case so, I went with RawConfigParser.

=== added file 'upload.yaml'
--- upload.yaml	1970-01-01 00:00:00 +0000
+++ upload.yaml	2015-06-30 18:41:09 +0000
@@ -0,0 +1,6 @@
+command: [ls > foo]
+install:
+  docs: [README]
+artifacts:
+  ls_output: [foo]
+bucket: ws-runner-test
 
=== modified file 'workspace_runner/__init__.py'
--- workspace_runner/__init__.py	2015-06-25 17:05:49 +0000
+++ workspace_runner/__init__.py	2015-06-30 18:41:09 +0000
@@ -1,15 +1,20 @@
 from argparse import ArgumentParser
+from ConfigParser import RawConfigParser
 from contextlib import contextmanager
 from itertools import chain
+import json
 import logging
 from pipes import quote
 import os
+from shutil import rmtree
 import subprocess
 import sys
+from tempfile import mkdtemp
 from textwrap import dedent
 
-from yaml import safe_load
-
+from yaml import (
+    safe_load,
+    )
 __metaclass__ = type
 
 
@@ -18,9 +23,13 @@
     parser = ArgumentParser()
     parser.add_argument('config', help='Config file to use.')
     parser.add_argument('host', help='Machine to run the command on.')
+    parser.add_argument('artifact_prefix', nargs='?',
+                        help='Prefix to use storing artifacts.')
     parser.add_argument('--private-key', '-i', help='Private SSH key to use.')
     parser.add_argument('--verbose', '-v', help='Verbose output.',
                         action='store_true')
+    parser.add_argument('--s3-config',
+                        help='s3cmd config file for credentials.')
     return parser.parse_args(argv)
 
 
@@ -185,7 +194,49 @@
         primitives.destroy()
 
 
-def workspace_run(argv=None, primitives_factory=SSHPrimitives):
+@contextmanager
+def temp_dir():
+    temp_dir = mkdtemp()
+    try:
+        yield temp_dir
+    finally:
+        rmtree(temp_dir)
+
+
+@contextmanager
+def temp_config(config):
+    with temp_dir() as config_dir:
+        config_filename = os.path.join(config_dir, 'upload.json')
+        with open(config_filename, 'w') as config_file:
+            json.dump(config, config_file)
+            config_file.flush()
+        yield config_filename
+
+
+def run_from_config(primitives, config, s3_config, artifact_prefix, output):
+    """Install the files and run the command."""
+    for target, sources in config['install'].items():
+        primitives.install(sources, target)
+    if 'artifacts' in config:
+        import upload_artifacts
+        sources = [upload_artifacts.__file__.replace('.pyc', '.py')]
+        upload_config = {
+            'bucket': config['bucket'],
+            'files': config['artifacts'],
+            'prefix': artifact_prefix,
+            }
+        upload_config.update(s3_config)
+        with temp_config(upload_config) as upload_config_filename:
+            sources.append(upload_config_filename)
+            primitives.install(sources, '.wsr')
+    primitives.run(config['command'], output)
+    if 'artifacts' in config:
+        primitives.run(['python', '.wsr/upload_artifacts.py',
+                        '.wsr/upload.json', primitives.workspace], output)
+
+
+def workspace_run(argv=None, primitives_factory=SSHPrimitives,
+                  output=sys.stdout):
     """Run an operation in a workspace."""
     args = parse_args(argv)
     if args.verbose:
@@ -195,8 +246,15 @@
     logging.basicConfig(level=level)
     with open(args.config) as config_file:
         config = safe_load(config_file)
+    s3_config = None
+    if args.s3_config is not None:
+        config_parser = RawConfigParser()
+        config_parser.read(args.s3_config)
+        s3_config = {
+            'access_key': config_parser.get('default', 'access_key'),
+            'secret_key': config_parser.get('default', 'secret_key'),
+        }
     with workspace_context(args.host, args.private_key,
                            primitives_factory) as runner:
-        for target, sources in config['install'].items():
-            runner.install(sources, target)
-        runner.run(config['command'], sys.stdout)
+        run_from_config(runner, config, s3_config, args.artifact_prefix,
+                        output)
 
=== modified file 'workspace_runner/tests/__init__.py'
--- workspace_runner/tests/__init__.py	2015-06-25 17:05:49 +0000
+++ workspace_runner/tests/__init__.py	2015-06-30 18:41:09 +0000
@@ -1,5 +1,6 @@
 from argparse import Namespace
 from contextlib import contextmanager
+import json
 import os
 import logging
 from mock import patch
@@ -9,6 +10,7 @@
     rmtree,
     )
 from StringIO import StringIO
+import sys
 from tempfile import (
     mkdtemp,
     NamedTemporaryFile,
@@ -16,7 +18,6 @@
 from textwrap import dedent
 from unittest import TestCase
 import subprocess
-
 from yaml import safe_dump
 
 from workspace_runner import (
@@ -25,26 +26,21 @@
     retry_ssh,
     SSHConnection,
     SSHPrimitives,
+    temp_config,
+    temp_dir,
+    run_from_config,
     workspace_context,
     workspace_run,
     )
 
 
-@contextmanager
-def temp_dir():
-    temp_dir = mkdtemp()
-    try:
-        yield temp_dir
-    finally:
-        rmtree(temp_dir)
-
-
 class TestParseArgs(TestCase):
 
     def test_minimal(self):
         args = parse_args(['foo', 'bar'])
-        self.assertEqual(args, Namespace(config='foo', host='bar',
-                         private_key=None, verbose=False))
+        self.assertEqual(args, Namespace(
+            config='foo', host='bar', private_key=None, verbose=False,
+            s3_config=None, artifact_prefix=None))
 
     def test_private_key(self):
         args = parse_args(['foo', 'bar', '--private-key', 'key'])
@@ -58,6 +54,10 @@
         args = parse_args(['foo', 'bar', '-v'])
         self.assertEqual(args.verbose, True)
 
+    def test_s3_config(self):
+        args = parse_args(['foo', 'bar', '--s3-config', 'foobar'])
+        self.assertEqual(args.s3_config, 'foobar')
+
 
 class FakePrimitives(Primitives):
 
@@ -364,26 +364,75 @@
         self.assertFalse(os.path.exists(primitives.workspace))
 
 
+class TestRunFromConfig(TestCase):
+
+    def test_minimal(self):
+        config = {
+            'command': ['run', 'this'],
+            'install': {},
+            }
+        with workspace_context('foo', None, FakePrimitives) as primitives:
+            run_from_config(primitives, config, None, None, StringIO())
+        self.assertEqual(primitives.run_calls, [['run', 'this']])
+        self.assertEqual(primitives.walk_results,
+                         [(primitives.workspace, [], [])])
+
+    def test_s3_upload(self):
+        config = {
+            'command': ['run', 'this'],
+            'install': {},
+            'artifacts': {'foo': ['bar']},
+            'bucket': 'bucket1',
+            }
+        s3_config = {'access_key': 'access1', 'secret_key': 'secret1'}
+        with workspace_context('foo', None, FakePrimitives) as primitives:
+            run_from_config(primitives, config, s3_config, 'prefix/midfix',
+                            StringIO())
+            upload_json_path = os.path.join(primitives.workspace, '.wsr',
+                                            'upload.json')
+            with open(upload_json_path) as remote_config_file:
+                remote_config = json.load(remote_config_file)
+        self.assertEqual(remote_config, {
+            'access_key': 'access1',
+            'secret_key': 'secret1',
+            'bucket': 'bucket1',
+            'files': {'foo': ['bar']},
+            'prefix': 'prefix/midfix',
+            })
+        self.assertEqual(primitives.run_calls, [
+            ['run', 'this'],
+            ['python', '.wsr/upload_artifacts.py', '.wsr/upload.json',
+             primitives.workspace],
+            ])
+        self.assertEqual(primitives.walk_results, [
+            (primitives.workspace, ['.wsr'], []),
+            (primitives.workspace + '/.wsr', [], [
+                'upload.json', 'upload_artifacts.py'])
+            ])
+
+
 class TestWorkspaceRun(TestCase):
 
     @contextmanager
     def config_file(self):
-        with NamedTemporaryFile() as config_file:
-            safe_dump({
-                'command': ['run', 'this'],
-                'install': {},
-                }, config_file)
-            yield config_file
+        config = {
+            'command': ['run', 'this'],
+            'install': {},
+            }
+        with temp_config(config) as config_file_name:
+            yield config_file_name
 
-    def run_primitives(self, args):
+    def run_primitives(self, args, output=None):
         fp_factory = FakePrimitivesFactory()
-        workspace_run(args, fp_factory)
+        if output is None:
+            output = StringIO()
+        workspace_run(args, fp_factory, output=output)
         return fp_factory.last_instance
 
     def test_minimal(self):
-        with self.config_file() as config_file:
+        with self.config_file() as config_file_name:
             with patch('logging.root.handlers', []):
-                primitives = self.run_primitives([config_file.name, 'bar'])
+                primitives = self.run_primitives([config_file_name, 'bar'])
                 self.assertEqual(logging.getLogger().getEffectiveLevel(),
                                  logging.WARNING)
         self.assertEqual(primitives.run_calls, [['run', 'this']])
@@ -392,9 +441,9 @@
                          [(primitives.workspace, [], [])])
 
     def test_private_key(self):
-        with self.config_file() as config_file:
+        with self.config_file() as config_file_name:
             primitives = self.run_primitives(
-                [config_file.name, 'bar', '-i', 'qux'])
+                [config_file_name, 'bar', '-i', 'qux'])
         self.assertEqual(primitives.ssh_connection.private_key, 'qux')
 
     def test_install(self):
@@ -407,7 +456,7 @@
                     'install': {'bin-dir': [install_file.name]},
                     }, config_file)
                 workspace_run([config_file.name, 'bar'],
-                              fp_factory)
+                              fp_factory, StringIO())
         primitives = fp_factory.last_instance
         bin_path = os.path.join(primitives.workspace, 'bin-dir')
         install_base = os.path.basename(install_file.name)
@@ -417,8 +466,27 @@
             ])
 
     def test_verbose(self):
-        with self.config_file() as config_file:
+        with self.config_file() as config_file_name:
             with patch('logging.root.handlers', []):
-                self.run_primitives([config_file.name, 'bar', '-v'])
+                self.run_primitives([config_file_name, 'bar', '-v'])
                 self.assertEqual(logging.getLogger().getEffectiveLevel(),
                                  logging.INFO)
+
+    def test_s3config(self):
+        config = {'command': ['run', 'this'], 'install': {}}
+        credentials = {'access_key': 'foobar', 'secret_key': 'barfoo'}
+        with NamedTemporaryFile() as s3_config_file:
+            s3_config_file.write(dedent("""\
+                [default]
+                access_key = foobar
+                secret_key = barfoo
+            """))
+            s3_config_file.flush()
+            with self.config_file() as config_file_name:
+                with patch('workspace_runner.run_from_config',
+                           autospec=True) as rfc_mock:
+                    argv = [config_file_name, 'bar',
+                            '--s3-config', s3_config_file.name]
+                    primitives = self.run_primitives(argv, output=sys.stdout)
+                    rfc_mock.assert_called_once_with(
+                        primitives, config, credentials, None, sys.stdout)
 
=== modified file 'workspace_runner/upload_artifacts.py'
--- workspace_runner/upload_artifacts.py	2015-06-30 18:41:09 +0000
+++ workspace_runner/upload_artifacts.py	2015-06-30 18:41:09 +0000
@@ -11,8 +11,8 @@
 
 def parse_args(argv=None):
     parser = ArgumentParser()
-    parser.add_argument('artifacts_file')
-    parser.add_argument('root')
+    parser.add_argument('artifacts_file', help='Configuration file.')
+    parser.add_argument('root', help='The root directory to upload from.')
     return parser.parse_args(argv)
 
 
@@ -40,7 +40,7 @@
     args = parse_args(argv)
     # Use JSON rather than YAML because a program will emit it and no external
     # libs required.
-    with file(args.artifacts_file) as settings_file:
+    with open(args.artifacts_file) as settings_file:
         settings = json.load(settings_file)
     upload_artifacts(args.root, settings)