Merge pull request #332 from dhermes/add-params-to-key-iterator

Adding extra parameters for querying Cloud Storage keys.

Author: dhermes

gcloud/storage/bucket.py

@@ -301,7 +301,7 @@ def upload_file(self, filename, key=None):
         if key is None:
             key = os.path.basename(filename)
         key = self.new_key(key)
-        return key.set_contents_from_filename(filename)
+        return key.upload_from_filename(filename)
 
     def upload_file_object(self, file_obj, key=None):
         """Shortcut method to upload a file object into this bucket.
@@ -341,7 +341,7 @@ def upload_file_object(self, file_obj, key=None):
             key = self.new_key(key)
         else:
             key = self.new_key(os.path.basename(file_obj.name))
-        return key.set_contents_from_file(file_obj)
+        return key.upload_from_file(file_obj)
 
     def configure_website(self, main_page_suffix=None, not_found_page=None):
         """Configure website-related metadata.

gcloud/storage/iterator.py

@@ -38,15 +38,24 @@ class Iterator(object):
     :type path: string
     :param path: The path to query for the list of items.
     """
-    def __init__(self, connection, path):
+
+    PAGE_TOKEN = 'pageToken'
+    RESERVED_PARAMS = frozenset([PAGE_TOKEN])
+
+    def __init__(self, connection, path, extra_params=None):
         self.connection = connection
         self.path = path
         self.page_number = 0
         self.next_page_token = None
+        self.extra_params = extra_params or {}
+        reserved_in_use = self.RESERVED_PARAMS.intersection(
+            self.extra_params.keys())
+        if reserved_in_use:
+            raise ValueError(('Using a reserved parameter',
+                              reserved_in_use))
 
     def __iter__(self):
         """Iterate through the list of items."""
-
         while self.has_next_page():
             response = self.get_next_page_response()
             for item in self.get_items_from_response(response):
@@ -66,11 +75,13 @@ def has_next_page(self):
     def get_query_params(self):
         """Getter for query parameters for the next request.
 
-        :rtype: dict or None
-        :returns: A dictionary of query parameters or None if there are none.
+        :rtype: dict
+        :returns: A dictionary of query parameters.
         """
-        if self.next_page_token:
-            return {'pageToken': self.next_page_token}
+        result = ({self.PAGE_TOKEN: self.next_page_token}
+                  if self.next_page_token else {})
+        result.update(self.extra_params)
+        return result
 
     def get_next_page_response(self):
         """Requests the next page from the path provided.

gcloud/storage/key.py

@@ -342,9 +342,9 @@ def upload_from_string(self, data, content_type='text/plain'):
         """
         string_buffer = StringIO()
         string_buffer.write(data)
-        self.set_contents_from_file(file_obj=string_buffer, rewind=True,
-                                    size=string_buffer.len,
-                                    content_type=content_type)
+        self.upload_from_file(file_obj=string_buffer, rewind=True,
+                              size=string_buffer.len,
+                              content_type=content_type)
         return self
 
     # NOTE: Alias for boto-like API.
@@ -369,10 +369,11 @@ class _KeyIterator(Iterator):
     :type bucket: :class:`gcloud.storage.bucket.Bucket`
     :param bucket: The bucket from which to list keys.
     """
-    def __init__(self, bucket):
+    def __init__(self, bucket, extra_params=None):
         self.bucket = bucket
         super(_KeyIterator, self).__init__(
-            connection=bucket.connection, path=bucket.path + '/o')
+            connection=bucket.connection, path=bucket.path + '/o',
+            extra_params=extra_params)
 
     def get_items_from_response(self, response):
         """Factory method, yields :class:`.storage.key.Key` items from response.

gcloud/storage/test_bucket.py

@@ -77,7 +77,7 @@ def test___iter___empty(self):
         kw, = connection._requested
         self.assertEqual(kw['method'], 'GET')
         self.assertEqual(kw['path'], '/b/%s/o' % NAME)
-        self.assertEqual(kw['query_params'], None)
+        self.assertEqual(kw['query_params'], {})
 
     def test___iter___non_empty(self):
         NAME = 'name'
@@ -91,7 +91,7 @@ def test___iter___non_empty(self):
         kw, = connection._requested
         self.assertEqual(kw['method'], 'GET')
         self.assertEqual(kw['path'], '/b/%s/o' % NAME)
-        self.assertEqual(kw['query_params'], None)
+        self.assertEqual(kw['query_params'], {})
 
     def test___contains___miss(self):
         NAME = 'name'
@@ -154,7 +154,7 @@ def test_get_all_keys_empty(self):
         kw, = connection._requested
         self.assertEqual(kw['method'], 'GET')
         self.assertEqual(kw['path'], '/b/%s/o' % NAME)
-        self.assertEqual(kw['query_params'], None)
+        self.assertEqual(kw['query_params'], {})
 
     def test_get_all_keys_non_empty(self):
         NAME = 'name'
@@ -168,7 +168,7 @@ def test_get_all_keys_non_empty(self):
         kw, = connection._requested
         self.assertEqual(kw['method'], 'GET')
         self.assertEqual(kw['path'], '/b/%s/o' % NAME)
-        self.assertEqual(kw['query_params'], None)
+        self.assertEqual(kw['query_params'], {})
 
     def test_new_key_existing(self):
         from gcloud.storage.key import Key
@@ -334,8 +334,9 @@ def __init__(self, bucket, name):
                 self._bucket = bucket
                 self._name = name
 
-            def set_contents_from_filename(self, filename):
+            def upload_from_filename(self, filename):
                 _uploaded.append((self._bucket, self._name, filename))
+
         bucket = self._makeOne()
         with _Monkey(MUT, Key=_Key):
             bucket.upload_file(FILENAME)
@@ -354,8 +355,9 @@ def __init__(self, bucket, name):
                 self._bucket = bucket
                 self._name = name
 
-            def set_contents_from_filename(self, filename):
+            def upload_from_filename(self, filename):
                 _uploaded.append((self._bucket, self._name, filename))
+
         bucket = self._makeOne()
         with _Monkey(MUT, Key=_Key):
             bucket.upload_file(FILENAME, KEY)
@@ -374,8 +376,9 @@ def __init__(self, bucket, name):
                 self._bucket = bucket
                 self._name = name
 
-            def set_contents_from_file(self, fh):
+            def upload_from_file(self, fh):
                 _uploaded.append((self._bucket, self._name, fh))
+
         bucket = self._makeOne()
         with _Monkey(MUT, Key=_Key):
             bucket.upload_file_object(FILEOBJECT)
@@ -395,8 +398,9 @@ def __init__(self, bucket, name):
                 self._bucket = bucket
                 self._name = name
 
-            def set_contents_from_file(self, fh):
+            def upload_from_file(self, fh):
                 _uploaded.append((self._bucket, self._name, fh))
+
         bucket = self._makeOne()
         with _Monkey(MUT, Key=_Key):
             bucket.upload_file_object(FILEOBJECT, KEY)
@@ -728,7 +732,7 @@ def get_items_from_response(self, response):
         self.assertEqual(kw[0]['query_params'], {'projection': 'full'})
         self.assertEqual(kw[1]['method'], 'GET')
         self.assertEqual(kw[1]['path'], '/b/%s/o' % NAME)
-        self.assertEqual(kw[1]['query_params'], None)
+        self.assertEqual(kw[1]['query_params'], {})
 
     def test_get_lifecycle_eager(self):
         NAME = 'name'

gcloud/storage/test_iterator.py

@@ -36,7 +36,7 @@ def _get_items(response):
         kw, = connection._requested
         self.assertEqual(kw['method'], 'GET')
         self.assertEqual(kw['path'], PATH)
-        self.assertEqual(kw['query_params'], None)
+        self.assertEqual(kw['query_params'], {})
 
     def test_has_next_page_new(self):
         connection = _Connection()
@@ -64,7 +64,7 @@ def test_get_query_params_no_token(self):
         connection = _Connection()
         PATH = '/foo'
         iterator = self._makeOne(connection, PATH)
-        self.assertEqual(iterator.get_query_params(), None)
+        self.assertEqual(iterator.get_query_params(), {})
 
     def test_get_query_params_w_token(self):
         connection = _Connection()
@@ -75,6 +75,32 @@ def test_get_query_params_w_token(self):
         self.assertEqual(iterator.get_query_params(),
                          {'pageToken': TOKEN})
 
+    def test_get_query_params_extra_params(self):
+        connection = _Connection()
+        PATH = '/foo'
+        extra_params = {'key': 'val'}
+        iterator = self._makeOne(connection, PATH, extra_params=extra_params)
+        self.assertEqual(iterator.get_query_params(), extra_params)
+
+    def test_get_query_params_w_token_and_extra_params(self):
+        connection = _Connection()
+        PATH = '/foo'
+        TOKEN = 'token'
+        extra_params = {'key': 'val'}
+        iterator = self._makeOne(connection, PATH, extra_params=extra_params)
+        iterator.next_page_token = TOKEN
+
+        expected_query = extra_params.copy()
+        expected_query.update({'pageToken': TOKEN})
+        self.assertEqual(iterator.get_query_params(), expected_query)
+
+    def test_get_query_params_w_token_collision(self):
+        connection = _Connection()
+        PATH = '/foo'
+        extra_params = {'pageToken': 'val'}
+        self.assertRaises(ValueError, self._makeOne, connection, PATH,
+                          extra_params=extra_params)
+
     def test_get_next_page_response_new_no_token_in_response(self):
         PATH = '/foo'
         TOKEN = 'token'
@@ -90,7 +116,7 @@ def test_get_next_page_response_new_no_token_in_response(self):
         kw, = connection._requested
         self.assertEqual(kw['method'], 'GET')
         self.assertEqual(kw['path'], PATH)
-        self.assertEqual(kw['query_params'], None)
+        self.assertEqual(kw['query_params'], {})
 
     def test_get_next_page_response_no_token(self):
         connection = _Connection()