Merge pull request #3758 from markhjelm/develop

GCS Requester Pays changes

5 files changed, 143 insertions, 13 deletions

diff --git a/boto/gs/bucket.py b/boto/gs/bucket.py
index 67c601e6..dcce1d08 100644
--- a/boto/gs/bucket.py
+++ b/boto/gs/bucket.py
@@ -50,6 +50,9 @@ ERROR_DETAILS_REGEX = re.compile(r'<Details>(?P<details>.*)</Details>')
 class Bucket(S3Bucket):
     """Represents a Google Cloud Storage bucket."""
 
+    BillingBody = ('<?xml version="1.0" encoding="UTF-8"?>\n'
+                   '<BillingConfiguration><RequesterPays>%s</RequesterPays>'
+                   '</BillingConfiguration>')
     StorageClassBody = ('<?xml version="1.0" encoding="UTF-8"?>\n'
                         '<StorageClass>%s</StorageClass>')
     VersioningBody = ('<?xml version="1.0" encoding="UTF-8"?>\n'
@@ -594,7 +597,7 @@ class Bucket(S3Bucket):
             raise self.connection.provider.storage_response_error(
                 response.status, response.reason, body)
 
-    def get_storage_class(self):
+    def get_storage_class(self, headers=None):
         """
         Returns the StorageClass for the bucket.
 
@@ -602,7 +605,8 @@ class Bucket(S3Bucket):
         :return: The StorageClass for the bucket.
         """
         response = self.connection.make_request('GET', self.name,
-                                                query_args=STORAGE_CLASS_ARG)
+                                                query_args=STORAGE_CLASS_ARG,
+                                                headers=headers)
         body = response.read()
         if response.status == 200:
             rs = ResultSet(self)
@@ -999,3 +1003,65 @@ class Bucket(S3Bucket):
         else:
             raise self.connection.provider.storage_response_error(
                 response.status, response.reason, body)
+
+    def get_billing_config(self, headers=None):
+        """Returns the current status of billing configuration on the bucket.
+
+        :param dict headers: Additional headers to send with the request.
+
+        :rtype: dict
+        :returns: A dictionary containing the parsed XML response from GCS. The
+            overall structure is:
+
+            * BillingConfiguration
+
+              * RequesterPays: Enabled/Disabled.
+        """
+        return self.get_billing_configuration_with_xml(headers)[0]
+
+    def get_billing_configuration_with_xml(self, headers=None):
+        """Returns the current status of billing configuration on the bucket as
+        unparsed XML.
+
+        :param dict headers: Additional headers to send with the request.
+
+        :rtype: 2-Tuple
+        :returns: 2-tuple containing:
+
+            1) A dictionary containing the parsed XML response from GCS. The
+              overall structure is:
+
+              * BillingConfiguration
+
+                * RequesterPays: Enabled/Disabled.
+
+            2) Unparsed XML describing the bucket's website configuration.
+        """
+        response = self.connection.make_request('GET', self.name,
+                                                query_args='billing',
+                                                headers=headers)
+        body = response.read()
+        boto.log.debug(body)
+
+        if response.status != 200:
+            raise self.connection.provider.storage_response_error(
+                response.status, response.reason, body)
+
+        e = boto.jsonresponse.Element()
+        h = boto.jsonresponse.XmlHandler(e, None);
+        h.parse(body)
+        return e, body
+
+    def configure_billing(self, requester_pays=False, headers=None):
+        """Configure billing for this bucket.
+
+        :param bool requester_pays: If set to True, enables requester pays on
+            this bucket. If set to False, disables requester pays.
+
+        :param dict headers: Additional headers to send with the request.
+        """
+        if requester_pays == True:
+            req_body = self.BillingBody % ('Enabled')
+        else:
+            req_body = self.BillingBody % ('Disabled')
+        self.set_subresource('billing', req_body, headers=headers)
diff --git a/boto/s3/bucket.py b/boto/s3/bucket.py
index 88fb29ff..1adae4b1 100644
--- a/boto/s3/bucket.py
+++ b/boto/s3/bucket.py
@@ -852,8 +852,8 @@ class Bucket(object):
                 src_bucket = self
             else:
                 src_bucket = self.connection.get_bucket(
-                    src_bucket_name, validate=False)
-            acl = src_bucket.get_xml_acl(src_key_name)
+                    src_bucket_name, validate=False, headers=headers)
+            acl = src_bucket.get_xml_acl(src_key_name, headers=headers)
         if encrypt_key:
             headers[provider.server_side_encryption_header] = 'AES256'
         src = '%s/%s' % (src_bucket_name, urllib.parse.quote(src_key_name))
@@ -1123,7 +1123,7 @@ class Bucket(object):
         policy = self.get_acl(headers=headers)
         return policy.acl.grants
 
-    def get_location(self):
+    def get_location(self, headers=None):
         """
         Returns the LocationConstraint for the bucket.
 
@@ -1132,6 +1132,7 @@ class Bucket(object):
             string if no constraint was specified when bucket was created.
         """
         response = self.connection.make_request('GET', self.name,
+                                                headers=headers,
                                                 query_args='location')
         body = response.read()
         if response.status == 200:
@@ -1824,8 +1825,8 @@ class Bucket(object):
     def delete(self, headers=None):
         return self.connection.delete_bucket(self.name, headers=headers)
 
-    def get_tags(self):
-        response = self.get_xml_tags()
+    def get_tags(self, headers=None):
+        response = self.get_xml_tags(headers)
         tags = Tags()
         h = handler.XmlHandler(tags, self)
         if not isinstance(response, bytes):
@@ -1833,10 +1834,10 @@ class Bucket(object):
         xml.sax.parseString(response, h)
         return tags
 
-    def get_xml_tags(self):
+    def get_xml_tags(self, headers=None):
         response = self.connection.make_request('GET', self.name,
                                                 query_args='tagging',
-                                                headers=None)
+                                                headers=headers)
         body = response.read()
         if response.status == 200:
             return body
diff --git a/boto/storage_uri.py b/boto/storage_uri.py
index 128b2ca4..feecc0f0 100755
--- a/boto/storage_uri.py
+++ b/boto/storage_uri.py
@@ -440,7 +440,7 @@ class BucketStorageUri(StorageUri):
     def get_location(self, validate=False, headers=None):
         self._check_bucket_uri('get_location')
         bucket = self.get_bucket(validate, headers)
-        return bucket.get_location()
+        return bucket.get_location(headers)
 
     def get_storage_class(self, validate=False, headers=None):
         self._check_bucket_uri('get_storage_class')
@@ -450,7 +450,7 @@ class BucketStorageUri(StorageUri):
             raise ValueError('get_storage_class() not supported for %s '
                              'URIs.' % self.scheme)
         bucket = self.get_bucket(validate, headers)
-        return bucket.get_storage_class()
+        return bucket.get_storage_class(headers)
 
     def set_storage_class(self, storage_class, validate=False, headers=None):
         """Updates a bucket's storage class."""
@@ -801,11 +801,31 @@ class BucketStorageUri(StorageUri):
         bucket = self.get_bucket(validate, headers)
         bucket.configure_lifecycle(lifecycle_config, headers)
 
+    def get_billing_config(self, headers=None):
+        self._check_bucket_uri('get_billing_config')
+        # billing is defined as a bucket param for GCS, but not for S3.
+        if self.scheme != 'gs':
+            raise ValueError('get_billing_config() not supported for %s '
+                             'URIs.' % self.scheme)
+        bucket = self.get_bucket(False, headers)
+        return bucket.get_billing_config(headers)
+
+    def configure_billing(self, requester_pays=False, validate=False,
+                          headers=None):
+        """Sets or updates a bucket's billing configuration."""
+        self._check_bucket_uri('configure_billing')
+        # billing is defined as a bucket param for GCS, but not for S3.
+        if self.scheme != 'gs':
+            raise ValueError('configure_billing() not supported for %s '
+                             'URIs.' % self.scheme)
+        bucket = self.get_bucket(validate, headers)
+        bucket.configure_billing(requester_pays=requester_pays, headers=headers)
+
     def exists(self, headers=None):
         """Returns True if the object exists or False if it doesn't"""
         if not self.object_name:
             raise InvalidUriError('exists on object-less URI (%s)' % self.uri)
-        bucket = self.get_bucket()
+        bucket = self.get_bucket(headers)
         key = bucket.get_key(self.object_name, headers=headers)
         return bool(key)
 
diff --git a/boto/utils.py b/boto/utils.py
index 39a8cf77..f8801817 100644
--- a/boto/utils.py
+++ b/boto/utils.py
@@ -89,7 +89,11 @@ qsa_of_interest = ['acl', 'cors', 'defaultObjectAcl', 'location', 'logging',
                    # Storage.
                    'websiteConfig',
                    # compose is a QSA for objects in Google Cloud Storage.
-                   'compose']
+                   'compose',
+                   # billing is a QSA for buckets in Google Cloud Storage.
+                   'billing',
+                   # userProject is a QSA for requests in Google Cloud Storage.
+                   'userProject']
 
 
 _first_cap_regex = re.compile('(.)([A-Z][a-z]+)')
diff --git a/tests/integration/gs/test_basic.py b/tests/integration/gs/test_basic.py
index e12519de..f26f87a2 100644
--- a/tests/integration/gs/test_basic.py
+++ b/tests/integration/gs/test_basic.py
@@ -73,6 +73,10 @@ LIFECYCLE_CONDITIONS_FOR_DELETE_RULE = {
     'MatchesStorageClass': ['STANDARD']}
 LIFECYCLE_CONDITIONS_FOR_SET_STORAGE_CLASS_RULE = {'Age': '366'}
 
+BILLING_EMPTY = {'BillingConfiguration': {}}
+BILLING_ENABLED = {'BillingConfiguration': {'RequesterPays': 'Enabled'}}
+BILLING_DISABLED = {'BillingConfiguration': {'RequesterPays': 'Disabled'}}
+
 # Regexp for matching project-private default object ACL.
 PROJECT_PRIVATE_RE = ('\s*<AccessControlList>\s*<Entries>\s*<Entry>'
   '\s*<Scope type="GroupById">\s*<ID>[-a-zA-Z0-9]+</ID>'
@@ -452,3 +456,38 @@ class GSBasicTest(GSTestCase):
         uri.configure_lifecycle(lifecycle_config)
         xml = uri.get_lifecycle_config().to_xml()
         self.assertEqual(xml, LIFECYCLE_DOC)
+
+    def test_billing_config_bucket(self):
+        """Test setting and getting of billing config on Bucket."""
+        # create a new bucket
+        bucket = self._MakeBucket()
+        bucket_name = bucket.name
+        # get billing config and make sure it's empty
+        billing = bucket.get_billing_config()
+        self.assertEqual(billing, BILLING_EMPTY)
+        # set requester pays to enabled
+        bucket.configure_billing(requester_pays=True)
+        billing = bucket.get_billing_config()
+        self.assertEqual(billing, BILLING_ENABLED)
+        # set requester pays to disabled
+        bucket.configure_billing(requester_pays=False)
+        billing = bucket.get_billing_config()
+        self.assertEqual(billing, BILLING_DISABLED)
+
+    def test_billing_config_storage_uri(self):
+        """Test setting and getting of billing config with storage_uri."""
+        # create a new bucket
+        bucket = self._MakeBucket()
+        bucket_name = bucket.name
+        uri = storage_uri('gs://' + bucket_name)
+        # get billing config and make sure it's empty
+        billing = uri.get_billing_config()
+        self.assertEqual(billing, BILLING_EMPTY)
+        # set requester pays to enabled
+        uri.configure_billing(requester_pays=True)
+        billing = uri.get_billing_config()
+        self.assertEqual(billing, BILLING_ENABLED)
+        # set requester pays to disabled
+        uri.configure_billing(requester_pays=False)
+        billing = uri.get_billing_config()
+        self.assertEqual(billing, BILLING_DISABLED)