Merge "Add additional headers for HEAD/GET/DELETE requests."

5 files changed, 119 insertions, 47 deletions

diff --git a/swiftclient/client.py b/swiftclient/client.py
index 6eefec4..21cbe27 100644
--- a/swiftclient/client.py
+++ b/swiftclient/client.py
@@ -703,7 +703,7 @@ def store_response(resp, response_dict):
 
 def get_account(url, token, marker=None, limit=None, prefix=None,
                 end_marker=None, http_conn=None, full_listing=False,
-                service_token=None):
+                service_token=None, headers=None):
     """
     Get a listing of containers for the account.
 
@@ -718,20 +718,28 @@ def get_account(url, token, marker=None, limit=None, prefix=None,
     :param full_listing: if True, return a full listing, else returns a max
                          of 10000 listings
     :param service_token: service auth token
+    :param headers: additional headers to include in the request
     :returns: a tuple of (response headers, a list of containers) The response
               headers will be a dict and all header names will be lowercase.
     :raises ClientException: HTTP GET request failed
     """
+    req_headers = {'X-Auth-Token': token, 'Accept-Encoding': 'gzip'}
+    if service_token:
+        req_headers['X-Service-Token'] = service_token
+    if headers:
+        req_headers.update(headers)
+
     if not http_conn:
         http_conn = http_connection(url)
     if full_listing:
         rv = get_account(url, token, marker, limit, prefix,
-                         end_marker, http_conn)
+                         end_marker, http_conn, headers=req_headers)
         listing = rv[1]
         while listing:
             marker = listing[-1]['name']
             listing = get_account(url, token, marker, limit, prefix,
-                                  end_marker, http_conn)[1]
+                                  end_marker, http_conn,
+                                  headers=req_headers)[1]
             if listing:
                 rv[1].extend(listing)
         return rv
@@ -746,14 +754,12 @@ def get_account(url, token, marker=None, limit=None, prefix=None,
     if end_marker:
         qs += '&end_marker=%s' % quote(end_marker)
     full_path = '%s?%s' % (parsed.path, qs)
-    headers = {'X-Auth-Token': token, 'Accept-Encoding': 'gzip'}
-    if service_token:
-        headers['X-Service-Token'] = service_token
     method = 'GET'
-    conn.request(method, full_path, '', headers)
+    conn.request(method, full_path, '', req_headers)
     resp = conn.getresponse()
     body = resp.read()
-    http_log(("%s?%s" % (url, qs), method,), {'headers': headers}, resp, body)
+    http_log(("%s?%s" % (url, qs), method,), {'headers': req_headers},
+             resp, body)
 
     resp_headers = resp_header_dict(resp)
     if resp.status < 200 or resp.status >= 300:
@@ -763,7 +769,8 @@ def get_account(url, token, marker=None, limit=None, prefix=None,
     return resp_headers, parse_api_response(resp_headers, body)
 
 
-def head_account(url, token, http_conn=None, service_token=None):
+def head_account(url, token, http_conn=None, headers=None,
+                 service_token=None):
     """
     Get account stats.
 
@@ -771,6 +778,7 @@ def head_account(url, token, http_conn=None, service_token=None):
     :param token: auth token
     :param http_conn: a tuple of (parsed url, HTTPConnection object),
                       (If None, it will create the conn object)
+    :param headers: additional headers to include in the request
     :param service_token: service auth token
     :returns: a dict containing the response's headers (all header names will
               be lowercase)
@@ -781,13 +789,16 @@ def head_account(url, token, http_conn=None, service_token=None):
     else:
         parsed, conn = http_connection(url)
     method = "HEAD"
-    headers = {'X-Auth-Token': token}
+    req_headers = {'X-Auth-Token': token}
     if service_token:
-        headers['X-Service-Token'] = service_token
-    conn.request(method, parsed.path, '', headers)
+        req_headers['X-Service-Token'] = service_token
+    if headers:
+        req_headers.update(headers)
+
+    conn.request(method, parsed.path, '', req_headers)
     resp = conn.getresponse()
     body = resp.read()
-    http_log((url, method,), {'headers': headers}, resp, body)
+    http_log((url, method,), {'headers': req_headers}, resp, body)
     if resp.status < 200 or resp.status >= 300:
         raise ClientException.from_response(resp, 'Account HEAD failed', body)
     resp_headers = resp_header_dict(resp)
@@ -1054,7 +1065,7 @@ def post_container(url, token, container, headers, http_conn=None,
 
 def delete_container(url, token, container, http_conn=None,
                      response_dict=None, service_token=None,
-                     query_string=None):
+                     query_string=None, headers=None):
     """
     Delete a container
 
@@ -1067,6 +1078,7 @@ def delete_container(url, token, container, http_conn=None,
                      the response - status, reason and headers
     :param service_token: service auth token
     :param query_string: if set will be appended with '?' to generated path
+    :param headers: additional headers to include in the request
     :raises ClientException: HTTP DELETE request failed
     """
     if http_conn:
@@ -1074,7 +1086,12 @@ def delete_container(url, token, container, http_conn=None,
     else:
         parsed, conn = http_connection(url)
     path = '%s/%s' % (parsed.path, quote(container))
-    headers = {'X-Auth-Token': token}
+    if headers:
+        headers = dict(headers)
+    else:
+        headers = {}
+
+    headers['X-Auth-Token'] = token
     if service_token:
         headers['X-Service-Token'] = service_token
     if query_string:
@@ -1692,19 +1709,19 @@ class Connection(object):
             if reset_func:
                 reset_func(func, *args, **kwargs)
 
-    def head_account(self):
+    def head_account(self, headers=None):
         """Wrapper for :func:`head_account`"""
-        return self._retry(None, head_account)
+        return self._retry(None, head_account, headers=headers)
 
     def get_account(self, marker=None, limit=None, prefix=None,
-                    end_marker=None, full_listing=False):
+                    end_marker=None, full_listing=False, headers=None):
         """Wrapper for :func:`get_account`"""
         # TODO(unknown): With full_listing=True this will restart the entire
         # listing with each retry. Need to make a better version that just
         # retries where it left off.
         return self._retry(None, get_account, marker=marker, limit=limit,
                            prefix=prefix, end_marker=end_marker,
-                           full_listing=full_listing)
+                           full_listing=full_listing, headers=headers)
 
     def post_account(self, headers, response_dict=None,
                      query_string=None, data=None):
@@ -1743,11 +1760,12 @@ class Connection(object):
                            response_dict=response_dict)
 
     def delete_container(self, container, response_dict=None,
-                         query_string=None):
+                         query_string=None, headers={}):
         """Wrapper for :func:`delete_container`"""
         return self._retry(None, delete_container, container,
                            response_dict=response_dict,
-                           query_string=query_string)
+                           query_string=query_string,
+                           headers=headers)
 
     def head_object(self, container, obj, headers=None):
         """Wrapper for :func:`head_object`"""
@@ -1818,11 +1836,12 @@ class Connection(object):
                            response_dict=response_dict)
 
     def delete_object(self, container, obj, query_string=None,
-                      response_dict=None):
+                      response_dict=None, headers=None):
         """Wrapper for :func:`delete_object`"""
         return self._retry(None, delete_object, container, obj,
                            query_string=query_string,
-                           response_dict=response_dict)
+                           response_dict=response_dict,
+                           headers=headers)
 
     def get_capabilities(self, url=None):
         url = url or self.url
diff --git a/swiftclient/command_helpers.py b/swiftclient/command_helpers.py
index 0822699..49ccad1 100644
--- a/swiftclient/command_helpers.py
+++ b/swiftclient/command_helpers.py
@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from swiftclient.utils import prt_bytes
+from swiftclient.utils import prt_bytes, split_request_headers
 
 
 POLICY_HEADER_PREFIX = 'x-account-storage-policy-'
@@ -19,8 +19,9 @@ POLICY_HEADER_PREFIX = 'x-account-storage-policy-'
 
 def stat_account(conn, options):
     items = []
+    req_headers = split_request_headers(options.get('header', []))
 
-    headers = conn.head_account()
+    headers = conn.head_account(headers=req_headers)
     if options['verbose'] > 1:
         items.extend([
             ('StorageURL', conn.url),
@@ -91,8 +92,11 @@ def print_account_stats(items, headers, output_manager):
 
 
 def stat_container(conn, options, container):
-    headers = conn.head_container(container)
+    req_headers = split_request_headers(options.get('header', []))
+
+    headers = conn.head_container(container, headers=req_headers)
     items = []
+
     if options['verbose'] > 1:
         path = '%s/%s' % (conn.url, container)
         items.extend([
@@ -137,7 +141,9 @@ def print_container_stats(items, headers, output_manager):
 
 
 def stat_object(conn, options, container, obj):
-    headers = conn.head_object(container, obj)
+    req_headers = split_request_headers(options.get('header', []))
+
+    headers = conn.head_object(container, obj, headers=req_headers)
     items = []
     if options['verbose'] > 1:
         path = '%s/%s/%s' % (conn.url, container, obj)
diff --git a/swiftclient/service.py b/swiftclient/service.py
index f204895..de30d50 100644
--- a/swiftclient/service.py
+++ b/swiftclient/service.py
@@ -44,7 +44,7 @@ from swiftclient.command_helpers import (
 )
 from swiftclient.utils import (
     config_true_value, ReadableToIterable, LengthWrapper, EMPTY_ETAG,
-    parse_api_response, report_traceback, n_groups
+    parse_api_response, report_traceback, n_groups, split_request_headers
 )
 from swiftclient.exceptions import ClientException
 from swiftclient.multithreading import MultiThreadingManager
@@ -279,15 +279,11 @@ def split_headers(options, prefix=''):
         reporting.
     """
     headers = {}
-    for item in options:
-        split_item = item.split(':', 1)
-        if len(split_item) == 2:
-            headers[(prefix + split_item[0]).title()] = split_item[1].strip()
-        else:
-            raise SwiftError(
-                "Metadata parameter %s must contain a ':'.\n%s"
-                % (item, "Example: 'Color:Blue' or 'Size:Large'")
-            )
+    try:
+        headers = split_request_headers(options, prefix)
+    except ValueError as e:
+        raise SwiftError(e)
+
     return headers
 
 
@@ -467,7 +463,8 @@ class SwiftService(object):
                         performed by this call::
 
                             {
-                                'human': False
+                                'human': False,
+                                'header': []
                             }
 
         :returns: Either a single dictionary containing stats about an account
@@ -849,6 +846,7 @@ class SwiftService(object):
                                 'long': False,
                                 'prefix': None,
                                 'delimiter': None,
+                                'header': []
                             }
 
         :returns: A generator for returning the results of the list operation
@@ -884,10 +882,12 @@ class SwiftService(object):
     def _list_account_job(conn, options, result_queue):
         marker = ''
         error = None
+        req_headers = split_headers(options.get('header', []))
         try:
             while True:
                 _, items = conn.get_account(
-                    marker=marker, prefix=options['prefix']
+                    marker=marker, prefix=options['prefix'],
+                    headers=req_headers
                 )
 
                 if not items:
@@ -943,11 +943,12 @@ class SwiftService(object):
     def _list_container_job(conn, container, options, result_queue):
         marker = options.get('marker', '')
         error = None
+        req_headers = split_headers(options.get('header', []))
         try:
             while True:
                 _, items = conn.get_container(
                     container, marker=marker, prefix=options['prefix'],
-                    delimiter=options['delimiter']
+                    delimiter=options['delimiter'], headers=req_headers
                 )
 
                 if not items:
@@ -2112,6 +2113,7 @@ class SwiftService(object):
                                 'yes_all': False,
                                 'leave_segments': False,
                                 'prefix': None,
+                                'header': [],
                             }
 
         :returns: A generator for returning the results of the delete
@@ -2250,6 +2252,8 @@ class SwiftService(object):
 
     def _delete_object(self, conn, container, obj, options,
                        results_queue=None):
+        _headers = {}
+        _headers = split_headers(options.get('header', []))
         res = {
             'action': 'delete_object',
             'container': container,
@@ -2261,7 +2265,8 @@ class SwiftService(object):
 
             if not options['leave_segments']:
                 try:
-                    headers = conn.head_object(container, obj)
+                    headers = conn.head_object(container, obj,
+                                               headers=_headers)
                     old_manifest = headers.get('x-object-manifest')
                     if config_true_value(headers.get('x-static-large-object')):
                         query_string = 'multipart-manifest=delete'
@@ -2270,7 +2275,9 @@ class SwiftService(object):
                         raise
 
             results_dict = {}
-            conn.delete_object(container, obj, query_string=query_string,
+            conn.delete_object(container, obj,
+                               headers=_headers,
+                               query_string=query_string,
                                response_dict=results_dict)
 
             if old_manifest:
@@ -2322,10 +2329,13 @@ class SwiftService(object):
         return res
 
     @staticmethod
-    def _delete_empty_container(conn, container):
+    def _delete_empty_container(conn, container, options):
         results_dict = {}
+        _headers = {}
+        _headers = split_headers(options.get('header', []))
         try:
-            conn.delete_container(container, response_dict=results_dict)
+            conn.delete_container(container, headers=_headers,
+                                  response_dict=results_dict)
             res = {'success': True}
         except Exception as err:
             traceback, err_time = report_traceback()
@@ -2363,7 +2373,7 @@ class SwiftService(object):
                 return
 
             con_del = self.thread_manager.container_pool.submit(
-                self._delete_empty_container, container
+                self._delete_empty_container, container, options
             )
             con_del_res = get_future_result(con_del)
 
diff --git a/swiftclient/shell.py b/swiftclient/shell.py
index 3f8aad6..9fd7aac 100755
--- a/swiftclient/shell.py
+++ b/swiftclient/shell.py
@@ -58,6 +58,7 @@ def immediate_exit(signum, frame):
 st_delete_options = '''[--all] [--leave-segments]
                     [--object-threads <threads>]
                     [--container-threads <threads>]
+                    [--header <header:value>]
                     [<container> [<object>] [...]]
 '''
 
@@ -72,6 +73,9 @@ Positional arguments:
 Optional arguments:
   -a, --all             Delete all containers and objects.
   --leave-segments      Do not delete segments of manifest objects.
+  -H, --header <header:value>
+                        Adds a custom request header to use for deleting
+                        objects or an entire container .
   --object-threads <threads>
                         Number of threads to use for deleting objects.
                         Default is 10.
@@ -89,6 +93,11 @@ def st_delete(parser, args, output_manager):
         '-p', '--prefix', dest='prefix',
         help='Only delete items beginning with the <prefix>.')
     parser.add_argument(
+        '-H', '--header', action='append', dest='header',
+        default=[],
+        help='Adds a custom request header to use for deleting objects '
+        'or an entire container.')
+    parser.add_argument(
         '--leave-segments', action='store_true',
         dest='leave_segments', default=False,
         help='Do not delete segments of manifest objects.')
@@ -445,7 +454,8 @@ def st_download(parser, args, output_manager):
 
 
 st_list_options = '''[--long] [--lh] [--totals] [--prefix <prefix>]
-                  [--delimiter <delimiter>] [<container>]
+                  [--delimiter <delimiter>] [--header <header:value>]
+                  [<container>]
 '''
 
 st_list_help = '''
@@ -465,6 +475,8 @@ Optional arguments:
                         Roll up items with the given delimiter. For containers
                         only. See OpenStack Swift API documentation for what
                         this means.
+  -H, --header <header:value>
+                        Adds a custom request header to use for listing.
 '''.strip('\n')
 
 
@@ -541,6 +553,10 @@ def st_list(parser, args, output_manager):
         help='Roll up items with the given delimiter. For containers '
              'only. See OpenStack Swift API documentation for '
              'what this means.')
+    parser.add_argument(
+        '-H', '--header', action='append', dest='header',
+        default=[],
+        help='Adds a custom request header to use for listing.')
     options, args = parse_args(parser, args)
     args = args[1:]
     if options['delimiter'] and not args:
@@ -580,7 +596,7 @@ def st_list(parser, args, output_manager):
             output_manager.error(e.value)
 
 
-st_stat_options = '''[--lh]
+st_stat_options = '''[--lh] [--header <header:value>]
                   [<container> [<object>]]
 '''
 
@@ -594,6 +610,8 @@ Positional arguments:
 Optional arguments:
   --lh                  Report sizes in human readable format similar to
                         ls -lh.
+  -H, --header <header:value>
+                        Adds a custom request header to use for stat.
 '''.strip('\n')
 
 
@@ -601,6 +619,11 @@ def st_stat(parser, args, output_manager):
     parser.add_argument(
         '--lh', dest='human', action='store_true', default=False,
         help='Report sizes in human readable format similar to ls -lh.')
+    parser.add_argument(
+        '-H', '--header', action='append', dest='header',
+        default=[],
+        help='Adds a custom request header to use for stat.')
+
     options, args = parse_args(parser, args)
     args = args[1:]
 
diff --git a/swiftclient/utils.py b/swiftclient/utils.py
index 0d1104e..e14602d 100644
--- a/swiftclient/utils.py
+++ b/swiftclient/utils.py
@@ -146,6 +146,20 @@ def parse_api_response(headers, body):
     return json.loads(body.decode(charset))
 
 
+def split_request_headers(options, prefix=''):
+    headers = {}
+    for item in options:
+        split_item = item.split(':', 1)
+        if len(split_item) == 2:
+            headers[(prefix + split_item[0]).title()] = split_item[1].strip()
+        else:
+            raise ValueError(
+                "Metadata parameter %s must contain a ':'.\n%s"
+                % (item, "Example: 'Color:Blue' or 'Size:Large'")
+            )
+    return headers
+
+
 def report_traceback():
     """
     Reports a timestamp and full traceback for a given exception.