Allow to specify storage policy when uploading objects

Client already supports -H/--header option when creating container
or uploading objects. This patch extends this option to support
Storage Policy.

e.g.,
    swift post con -H 'X-Storage-Policy:p1'
This creates one container 'con' with storage policy 'p1'.

    swift upload con obj -H 'X-Storage-Policy:p2'
This creates container 'con' with storage policy 'p2' and uploads
object 'obj' into it.

Also fixes segmented uploading to non-default storage policy container

When uploading large objects with segmentation to container with
non-default storage policy, there will be another 'xxx_segments'
container created, but with the default storage policy. This
results all the segments to be stored with the wrong policy.

This patch is for the Storage Policy feature, and also
compatible with old versions w/o Storage Policy support.

Change-Id: I5c19e90604a0bcf2c85e1732b8a0b97ae6801994

2 files changed, 55 insertions, 8 deletions

diff --git a/swiftclient/shell.py b/swiftclient/shell.py
index d10fc70..ef0ee45 100755
--- a/swiftclient/shell.py
+++ b/swiftclient/shell.py
@@ -45,6 +45,7 @@ from swiftclient import __version__ as client_version
 
 
 BASENAME = 'swift'
+POLICY = 'X-Storage-Policy'
 
 
 def get_conn(options):
@@ -1158,13 +1159,48 @@ def st_upload(parser, args, thread_manager):
     # fails, it might just be because the user doesn't have container PUT
     # permissions, so we'll ignore any error. If there's really a problem,
     # it'll surface on the first object PUT.
+    container_name = args[0]
     try:
-        conn.put_container(args[0])
+        policy_header = {}
+        _header = split_headers(options.header)
+        if POLICY in _header:
+            policy_header[POLICY] = \
+                _header[POLICY]
+        try:
+            conn.put_container(args[0], policy_header)
+        except ClientException as err:
+            if err.http_status != 409:
+                raise
+            if POLICY in _header:
+                thread_manager.error('Error trying to create %s with '
+                                     'Storage Policy %s', args[0],
+                                     _header[POLICY].strip())
         if options.segment_size is not None:
-            seg_container = args[0] + '_segments'
+            container_name = seg_container = args[0] + '_segments'
             if options.segment_container:
-                seg_container = options.segment_container
-            conn.put_container(seg_container)
+                container_name = seg_container = options.segment_container
+            seg_headers = {}
+            if POLICY in _header:
+                seg_headers[POLICY] = \
+                    _header[POLICY]
+            else:
+                # Since no storage policy was specified on the command line,
+                # rather than just letting swift pick the default storage
+                # policy, we'll try to create the segments container with the
+                # same as the upload container
+                _meta = conn.head_container(args[0])
+                if 'x-storage-policy' in _meta:
+                    seg_headers[POLICY] = \
+                        _meta.get('x-storage-policy')
+            try:
+                conn.put_container(seg_container, seg_headers)
+            except ClientException as err:
+                if err.http_status != 409:
+                    raise
+                if POLICY in seg_headers:
+                    thread_manager.error('Error trying to create %s with '
+                                         'Storage Policy %s', seg_container,
+                                         seg_headers[POLICY].strip())
     except ClientException as err:
         msg = ' '.join(str(x) for x in (err.http_status, err.http_reason))
         if err.http_response_content:
@@ -1172,11 +1208,11 @@ def st_upload(parser, args, thread_manager):
                 msg += ': '
             msg += err.http_response_content[:60]
         thread_manager.error(
-            'Error trying to create container %r: %s', args[0],
+            'Error trying to create container %r: %s', container_name,
             msg)
     except Exception as err:
         thread_manager.error(
-            'Error trying to create container %r: %s', args[0],
+            'Error trying to create container %r: %s', container_name,
             err)
 
     if options.object_name is not None:
diff --git a/tests/unit/test_shell.py b/tests/unit/test_shell.py
index 0c28297..3def6f8 100644
--- a/tests/unit/test_shell.py
+++ b/tests/unit/test_shell.py
@@ -208,14 +208,20 @@ class TestShell(unittest.TestCase):
         connection.return_value.head_object.return_value = {
             'content-length': '0'}
         connection.return_value.attempts = 0
-        argv = ["", "upload", "container", self.tmpfile]
+        argv = ["", "upload", "container", self.tmpfile,
+                "-H", "X-Storage-Policy:one"]
         swiftclient.shell.main(argv)
+        connection.return_value.put_container.assert_called_with(
+            'container',
+            {'X-Storage-Policy': mock.ANY})
+
         connection.return_value.put_object.assert_called_with(
             'container',
             self.tmpfile.lstrip('/'),
             mock.ANY,
             content_length=0,
-            headers={'x-object-meta-mtime': mock.ANY})
+            headers={'x-object-meta-mtime': mock.ANY,
+                     'X-Storage-Policy': 'one'})
 
         # Upload whole directory
         argv = ["", "upload", "container", "/tmp"]
@@ -229,10 +235,15 @@ class TestShell(unittest.TestCase):
             headers={'x-object-meta-mtime': mock.ANY})
 
         # Upload in segments
+        connection.return_value.head_container.return_value = {
+            'x-storage-policy': 'one'}
         argv = ["", "upload", "container", self.tmpfile, "-S", "10"]
         with open(self.tmpfile, "wb") as fh:
             fh.write(b'12345678901234567890')
         swiftclient.shell.main(argv)
+        connection.return_value.put_container.assert_called_with(
+            'container_segments',
+            {'X-Storage-Policy': mock.ANY})
         connection.return_value.put_object.assert_called_with(
             'container',
             self.tmpfile.lstrip('/'),