Skip to content

Commit 826ea69

Browse files
dhermesdhermes
committed
Merge pull request #332 from dhermes/add-params-to-key-iterator
Adding extra parameters for querying Cloud Storage keys.
2 parents f7af09f + 2f28e9f commit 826ea69

File tree

5 files changed

+67
-25
lines changed

5 files changed

+67
-25
lines changed

‎gcloud/storage/bucket.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,7 @@ def upload_file(self, filename, key=None):
301301
if key is None:
302302
key = os.path.basename(filename)
303303
key = self.new_key(key)
304-
return key.set_contents_from_filename(filename)
304+
return key.upload_from_filename(filename)
305305

306306
def upload_file_object(self, file_obj, key=None):
307307
"""Shortcut method to upload a file object into this bucket.
@@ -341,7 +341,7 @@ def upload_file_object(self, file_obj, key=None):
341341
key = self.new_key(key)
342342
else:
343343
key = self.new_key(os.path.basename(file_obj.name))
344-
return key.set_contents_from_file(file_obj)
344+
return key.upload_from_file(file_obj)
345345

346346
def configure_website(self, main_page_suffix=None, not_found_page=None):
347347
"""Configure website-related metadata.

‎gcloud/storage/iterator.py‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,15 +38,24 @@ class Iterator(object):
3838
:type path: string
3939
:param path: The path to query for the list of items.
4040
"""
41-
def __init__(self, connection, path):
41+
42+
PAGE_TOKEN = 'pageToken'
43+
RESERVED_PARAMS = frozenset([PAGE_TOKEN])
44+
45+
def __init__(self, connection, path, extra_params=None):
4246
self.connection = connection
4347
self.path = path
4448
self.page_number = 0
4549
self.next_page_token = None
50+
self.extra_params = extra_params or {}
51+
reserved_in_use = self.RESERVED_PARAMS.intersection(
52+
self.extra_params.keys())
53+
if reserved_in_use:
54+
raise ValueError(('Using a reserved parameter',
55+
reserved_in_use))
4656

4757
def __iter__(self):
4858
"""Iterate through the list of items."""
49-
5059
while self.has_next_page():
5160
response = self.get_next_page_response()
5261
for item in self.get_items_from_response(response):
@@ -66,11 +75,13 @@ def has_next_page(self):
6675
def get_query_params(self):
6776
"""Getter for query parameters for the next request.
6877
69-
:rtype: dict or None
70-
:returns: A dictionary of query parameters or None if there are none.
78+
:rtype: dict
79+
:returns: A dictionary of query parameters.
7180
"""
72-
if self.next_page_token:
73-
return {'pageToken': self.next_page_token}
81+
result = ({self.PAGE_TOKEN: self.next_page_token}
82+
if self.next_page_token else {})
83+
result.update(self.extra_params)
84+
return result
7485

7586
def get_next_page_response(self):
7687
"""Requests the next page from the path provided.

‎gcloud/storage/key.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -342,9 +342,9 @@ def upload_from_string(self, data, content_type='text/plain'):
342342
"""
343343
string_buffer = StringIO()
344344
string_buffer.write(data)
345-
self.set_contents_from_file(file_obj=string_buffer, rewind=True,
346-
size=string_buffer.len,
347-
content_type=content_type)
345+
self.upload_from_file(file_obj=string_buffer, rewind=True,
346+
size=string_buffer.len,
347+
content_type=content_type)
348348
return self
349349

350350
# NOTE: Alias for boto-like API.
@@ -369,10 +369,11 @@ class _KeyIterator(Iterator):
369369
:type bucket: :class:`gcloud.storage.bucket.Bucket`
370370
:param bucket: The bucket from which to list keys.
371371
"""
372-
def __init__(self, bucket):
372+
def __init__(self, bucket, extra_params=None):
373373
self.bucket = bucket
374374
super(_KeyIterator, self).__init__(
375-
connection=bucket.connection, path=bucket.path + '/o')
375+
connection=bucket.connection, path=bucket.path + '/o',
376+
extra_params=extra_params)
376377

377378
def get_items_from_response(self, response):
378379
"""Factory method, yields :class:`.storage.key.Key` items from response.

‎gcloud/storage/test_bucket.py‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ def test___iter___empty(self):
7777
kw, = connection._requested
7878
self.assertEqual(kw['method'], 'GET')
7979
self.assertEqual(kw['path'], '/b/%s/o' % NAME)
80-
self.assertEqual(kw['query_params'], None)
80+
self.assertEqual(kw['query_params'], {})
8181

8282
def test___iter___non_empty(self):
8383
NAME = 'name'
@@ -91,7 +91,7 @@ def test___iter___non_empty(self):
9191
kw, = connection._requested
9292
self.assertEqual(kw['method'], 'GET')
9393
self.assertEqual(kw['path'], '/b/%s/o' % NAME)
94-
self.assertEqual(kw['query_params'], None)
94+
self.assertEqual(kw['query_params'], {})
9595

9696
def test___contains___miss(self):
9797
NAME = 'name'
@@ -154,7 +154,7 @@ def test_get_all_keys_empty(self):
154154
kw, = connection._requested
155155
self.assertEqual(kw['method'], 'GET')
156156
self.assertEqual(kw['path'], '/b/%s/o' % NAME)
157-
self.assertEqual(kw['query_params'], None)
157+
self.assertEqual(kw['query_params'], {})
158158

159159
def test_get_all_keys_non_empty(self):
160160
NAME = 'name'
@@ -168,7 +168,7 @@ def test_get_all_keys_non_empty(self):
168168
kw, = connection._requested
169169
self.assertEqual(kw['method'], 'GET')
170170
self.assertEqual(kw['path'], '/b/%s/o' % NAME)
171-
self.assertEqual(kw['query_params'], None)
171+
self.assertEqual(kw['query_params'], {})
172172

173173
def test_new_key_existing(self):
174174
from gcloud.storage.key import Key
@@ -334,8 +334,9 @@ def __init__(self, bucket, name):
334334
self._bucket = bucket
335335
self._name = name
336336

337-
def set_contents_from_filename(self, filename):
337+
def upload_from_filename(self, filename):
338338
_uploaded.append((self._bucket, self._name, filename))
339+
339340
bucket = self._makeOne()
340341
with _Monkey(MUT, Key=_Key):
341342
bucket.upload_file(FILENAME)
@@ -354,8 +355,9 @@ def __init__(self, bucket, name):
354355
self._bucket = bucket
355356
self._name = name
356357

357-
def set_contents_from_filename(self, filename):
358+
def upload_from_filename(self, filename):
358359
_uploaded.append((self._bucket, self._name, filename))
360+
359361
bucket = self._makeOne()
360362
with _Monkey(MUT, Key=_Key):
361363
bucket.upload_file(FILENAME, KEY)
@@ -374,8 +376,9 @@ def __init__(self, bucket, name):
374376
self._bucket = bucket
375377
self._name = name
376378

377-
def set_contents_from_file(self, fh):
379+
def upload_from_file(self, fh):
378380
_uploaded.append((self._bucket, self._name, fh))
381+
379382
bucket = self._makeOne()
380383
with _Monkey(MUT, Key=_Key):
381384
bucket.upload_file_object(FILEOBJECT)
@@ -395,8 +398,9 @@ def __init__(self, bucket, name):
395398
self._bucket = bucket
396399
self._name = name
397400

398-
def set_contents_from_file(self, fh):
401+
def upload_from_file(self, fh):
399402
_uploaded.append((self._bucket, self._name, fh))
403+
400404
bucket = self._makeOne()
401405
with _Monkey(MUT, Key=_Key):
402406
bucket.upload_file_object(FILEOBJECT, KEY)
@@ -728,7 +732,7 @@ def get_items_from_response(self, response):
728732
self.assertEqual(kw[0]['query_params'], {'projection': 'full'})
729733
self.assertEqual(kw[1]['method'], 'GET')
730734
self.assertEqual(kw[1]['path'], '/b/%s/o' % NAME)
731-
self.assertEqual(kw[1]['query_params'], None)
735+
self.assertEqual(kw[1]['query_params'], {})
732736

733737
def test_get_lifecycle_eager(self):
734738
NAME = 'name'

‎gcloud/storage/test_iterator.py‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ def _get_items(response):
3636
kw, = connection._requested
3737
self.assertEqual(kw['method'], 'GET')
3838
self.assertEqual(kw['path'], PATH)
39-
self.assertEqual(kw['query_params'], None)
39+
self.assertEqual(kw['query_params'], {})
4040

4141
def test_has_next_page_new(self):
4242
connection = _Connection()
@@ -64,7 +64,7 @@ def test_get_query_params_no_token(self):
6464
connection = _Connection()
6565
PATH = '/foo'
6666
iterator = self._makeOne(connection, PATH)
67-
self.assertEqual(iterator.get_query_params(), None)
67+
self.assertEqual(iterator.get_query_params(), {})
6868

6969
def test_get_query_params_w_token(self):
7070
connection = _Connection()
@@ -75,6 +75,32 @@ def test_get_query_params_w_token(self):
7575
self.assertEqual(iterator.get_query_params(),
7676
{'pageToken': TOKEN})
7777

78+
def test_get_query_params_extra_params(self):
79+
connection = _Connection()
80+
PATH = '/foo'
81+
extra_params = {'key': 'val'}
82+
iterator = self._makeOne(connection, PATH, extra_params=extra_params)
83+
self.assertEqual(iterator.get_query_params(), extra_params)
84+
85+
def test_get_query_params_w_token_and_extra_params(self):
86+
connection = _Connection()
87+
PATH = '/foo'
88+
TOKEN = 'token'
89+
extra_params = {'key': 'val'}
90+
iterator = self._makeOne(connection, PATH, extra_params=extra_params)
91+
iterator.next_page_token = TOKEN
92+
93+
expected_query = extra_params.copy()
94+
expected_query.update({'pageToken': TOKEN})
95+
self.assertEqual(iterator.get_query_params(), expected_query)
96+
97+
def test_get_query_params_w_token_collision(self):
98+
connection = _Connection()
99+
PATH = '/foo'
100+
extra_params = {'pageToken': 'val'}
101+
self.assertRaises(ValueError, self._makeOne, connection, PATH,
102+
extra_params=extra_params)
103+
78104
def test_get_next_page_response_new_no_token_in_response(self):
79105
PATH = '/foo'
80106
TOKEN = 'token'
@@ -90,7 +116,7 @@ def test_get_next_page_response_new_no_token_in_response(self):
90116
kw, = connection._requested
91117
self.assertEqual(kw['method'], 'GET')
92118
self.assertEqual(kw['path'], PATH)
93-
self.assertEqual(kw['query_params'], None)
119+
self.assertEqual(kw['query_params'], {})
94120

95121
def test_get_next_page_response_no_token(self):
96122
connection = _Connection()

0 commit comments

Comments
 (0)