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

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:

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('/'),