paginate.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677
  1. # Copyright 2012-2014 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License"). You
  4. # may not use this file except in compliance with the License. A copy of
  5. # the License is located at
  6. #
  7. # http://aws.amazon.com/apache2.0/
  8. #
  9. # or in the "license" file accompanying this file. This file is
  10. # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
  11. # ANY KIND, either express or implied. See the License for the specific
  12. # language governing permissions and limitations under the License.
  13. from itertools import tee
  14. from botocore.compat import six
  15. import jmespath
  16. import json
  17. import base64
  18. import logging
  19. from botocore.exceptions import PaginationError
  20. from botocore.compat import zip
  21. from botocore.utils import set_value_from_jmespath, merge_dicts
  22. log = logging.getLogger(__name__)
  23. class TokenEncoder(object):
  24. """Encodes dictionaries into opaque strings.
  25. This for the most part json dumps + base64 encoding, but also supports
  26. having bytes in the dictionary in addition to the types that json can
  27. handle by default.
  28. This is intended for use in encoding pagination tokens, which in some
  29. cases can be complex structures and / or contain bytes.
  30. """
  31. def encode(self, token):
  32. """Encodes a dictionary to an opaque string.
  33. :type token: dict
  34. :param token: A dictionary containing pagination information,
  35. particularly the service pagination token(s) but also other boto
  36. metadata.
  37. :rtype: str
  38. :returns: An opaque string
  39. """
  40. try:
  41. # Try just using json dumps first to avoid having to traverse
  42. # and encode the dict. In 99.9999% of cases this will work.
  43. json_string = json.dumps(token)
  44. except (TypeError, UnicodeDecodeError):
  45. # If normal dumping failed, go through and base64 encode all bytes.
  46. encoded_token, encoded_keys = self._encode(token, [])
  47. # Save the list of all the encoded key paths. We can safely
  48. # assume that no service will ever use this key.
  49. encoded_token['boto_encoded_keys'] = encoded_keys
  50. # Now that the bytes are all encoded, dump the json.
  51. json_string = json.dumps(encoded_token)
  52. # base64 encode the json string to produce an opaque token string.
  53. return base64.b64encode(json_string.encode('utf-8')).decode('utf-8')
  54. def _encode(self, data, path):
  55. """Encode bytes in given data, keeping track of the path traversed."""
  56. if isinstance(data, dict):
  57. return self._encode_dict(data, path)
  58. elif isinstance(data, list):
  59. return self._encode_list(data, path)
  60. elif isinstance(data, six.binary_type):
  61. return self._encode_bytes(data, path)
  62. else:
  63. return data, []
  64. def _encode_list(self, data, path):
  65. """Encode any bytes in a list, noting the index of what is encoded."""
  66. new_data = []
  67. encoded = []
  68. for i, value in enumerate(data):
  69. new_path = path + [i]
  70. new_value, new_encoded = self._encode(value, new_path)
  71. new_data.append(new_value)
  72. encoded.extend(new_encoded)
  73. return new_data, encoded
  74. def _encode_dict(self, data, path):
  75. """Encode any bytes in a dict, noting the index of what is encoded."""
  76. new_data = {}
  77. encoded = []
  78. for key, value in data.items():
  79. new_path = path + [key]
  80. new_value, new_encoded = self._encode(value, new_path)
  81. new_data[key] = new_value
  82. encoded.extend(new_encoded)
  83. return new_data, encoded
  84. def _encode_bytes(self, data, path):
  85. """Base64 encode a byte string."""
  86. return base64.b64encode(data).decode('utf-8'), [path]
  87. class TokenDecoder(object):
  88. """Decodes token strings back into dictionaries.
  89. This performs the inverse operation to the TokenEncoder, accepting
  90. opaque strings and decoding them into a useable form.
  91. """
  92. def decode(self, token):
  93. """Decodes an opaque string to a dictionary.
  94. :type token: str
  95. :param token: A token string given by the botocore pagination
  96. interface.
  97. :rtype: dict
  98. :returns: A dictionary containing pagination information,
  99. particularly the service pagination token(s) but also other boto
  100. metadata.
  101. """
  102. json_string = base64.b64decode(token.encode('utf-8')).decode('utf-8')
  103. decoded_token = json.loads(json_string)
  104. # Remove the encoding metadata as it is read since it will no longer
  105. # be needed.
  106. encoded_keys = decoded_token.pop('boto_encoded_keys', None)
  107. if encoded_keys is None:
  108. return decoded_token
  109. else:
  110. return self._decode(decoded_token, encoded_keys)
  111. def _decode(self, token, encoded_keys):
  112. """Find each encoded value and decode it."""
  113. for key in encoded_keys:
  114. encoded = self._path_get(token, key)
  115. decoded = base64.b64decode(encoded.encode('utf-8'))
  116. self._path_set(token, key, decoded)
  117. return token
  118. def _path_get(self, data, path):
  119. """Return the nested data at the given path.
  120. For instance:
  121. data = {'foo': ['bar', 'baz']}
  122. path = ['foo', 0]
  123. ==> 'bar'
  124. """
  125. # jmespath isn't used here because it would be difficult to actually
  126. # create the jmespath query when taking all of the unknowns of key
  127. # structure into account. Gross though this is, it is simple and not
  128. # very error prone.
  129. d = data
  130. for step in path:
  131. d = d[step]
  132. return d
  133. def _path_set(self, data, path, value):
  134. """Set the value of a key in the given data.
  135. Example:
  136. data = {'foo': ['bar', 'baz']}
  137. path = ['foo', 1]
  138. value = 'bin'
  139. ==> data = {'foo': ['bar', 'bin']}
  140. """
  141. container = self._path_get(data, path[:-1])
  142. container[path[-1]] = value
  143. class PaginatorModel(object):
  144. def __init__(self, paginator_config):
  145. self._paginator_config = paginator_config['pagination']
  146. def get_paginator(self, operation_name):
  147. try:
  148. single_paginator_config = self._paginator_config[operation_name]
  149. except KeyError:
  150. raise ValueError("Paginator for operation does not exist: %s"
  151. % operation_name)
  152. return single_paginator_config
  153. class PageIterator(object):
  154. def __init__(self, method, input_token, output_token, more_results,
  155. result_keys, non_aggregate_keys, limit_key, max_items,
  156. starting_token, page_size, op_kwargs):
  157. self._method = method
  158. self._input_token = input_token
  159. self._output_token = output_token
  160. self._more_results = more_results
  161. self._result_keys = result_keys
  162. self._max_items = max_items
  163. self._limit_key = limit_key
  164. self._starting_token = starting_token
  165. self._page_size = page_size
  166. self._op_kwargs = op_kwargs
  167. self._resume_token = None
  168. self._non_aggregate_key_exprs = non_aggregate_keys
  169. self._non_aggregate_part = {}
  170. self._token_encoder = TokenEncoder()
  171. self._token_decoder = TokenDecoder()
  172. @property
  173. def result_keys(self):
  174. return self._result_keys
  175. @property
  176. def resume_token(self):
  177. """Token to specify to resume pagination."""
  178. return self._resume_token
  179. @resume_token.setter
  180. def resume_token(self, value):
  181. if not isinstance(value, dict):
  182. raise ValueError("Bad starting token: %s" % value)
  183. if 'boto_truncate_amount' in value:
  184. token_keys = sorted(self._input_token + ['boto_truncate_amount'])
  185. else:
  186. token_keys = sorted(self._input_token)
  187. dict_keys = sorted(value.keys())
  188. if token_keys == dict_keys:
  189. self._resume_token = self._token_encoder.encode(value)
  190. else:
  191. raise ValueError("Bad starting token: %s" % value)
  192. @property
  193. def non_aggregate_part(self):
  194. return self._non_aggregate_part
  195. def __iter__(self):
  196. current_kwargs = self._op_kwargs
  197. previous_next_token = None
  198. next_token = dict((key, None) for key in self._input_token)
  199. if self._starting_token is not None:
  200. # If the starting token exists, populate the next_token with the
  201. # values inside it. This ensures that we have the service's
  202. # pagination token on hand if we need to truncate after the
  203. # first response.
  204. next_token = self._parse_starting_token()[0]
  205. # The number of items from result_key we've seen so far.
  206. total_items = 0
  207. first_request = True
  208. primary_result_key = self.result_keys[0]
  209. starting_truncation = 0
  210. self._inject_starting_params(current_kwargs)
  211. while True:
  212. response = self._make_request(current_kwargs)
  213. parsed = self._extract_parsed_response(response)
  214. if first_request:
  215. # The first request is handled differently. We could
  216. # possibly have a resume/starting token that tells us where
  217. # to index into the retrieved page.
  218. if self._starting_token is not None:
  219. starting_truncation = self._handle_first_request(
  220. parsed, primary_result_key, starting_truncation)
  221. first_request = False
  222. self._record_non_aggregate_key_values(parsed)
  223. else:
  224. # If this isn't the first request, we have already sliced into
  225. # the first request and had to make additional requests after.
  226. # We no longer need to add this to truncation.
  227. starting_truncation = 0
  228. current_response = primary_result_key.search(parsed)
  229. if current_response is None:
  230. current_response = []
  231. num_current_response = len(current_response)
  232. truncate_amount = 0
  233. if self._max_items is not None:
  234. truncate_amount = (total_items + num_current_response) \
  235. - self._max_items
  236. if truncate_amount > 0:
  237. self._truncate_response(parsed, primary_result_key,
  238. truncate_amount, starting_truncation,
  239. next_token)
  240. yield response
  241. break
  242. else:
  243. yield response
  244. total_items += num_current_response
  245. next_token = self._get_next_token(parsed)
  246. if all(t is None for t in next_token.values()):
  247. break
  248. if self._max_items is not None and \
  249. total_items == self._max_items:
  250. # We're on a page boundary so we can set the current
  251. # next token to be the resume token.
  252. self.resume_token = next_token
  253. break
  254. if previous_next_token is not None and \
  255. previous_next_token == next_token:
  256. message = ("The same next token was received "
  257. "twice: %s" % next_token)
  258. raise PaginationError(message=message)
  259. self._inject_token_into_kwargs(current_kwargs, next_token)
  260. previous_next_token = next_token
  261. def search(self, expression):
  262. """Applies a JMESPath expression to a paginator
  263. Each page of results is searched using the provided JMESPath
  264. expression. If the result is not a list, it is yielded
  265. directly. If the result is a list, each element in the result
  266. is yielded individually (essentially implementing a flatmap in
  267. which the JMESPath search is the mapping function).
  268. :type expression: str
  269. :param expression: JMESPath expression to apply to each page.
  270. :return: Returns an iterator that yields the individual
  271. elements of applying a JMESPath expression to each page of
  272. results.
  273. """
  274. compiled = jmespath.compile(expression)
  275. for page in self:
  276. results = compiled.search(page)
  277. if isinstance(results, list):
  278. for element in results:
  279. yield element
  280. else:
  281. # Yield result directly if it is not a list.
  282. yield results
  283. def _make_request(self, current_kwargs):
  284. return self._method(**current_kwargs)
  285. def _extract_parsed_response(self, response):
  286. return response
  287. def _record_non_aggregate_key_values(self, response):
  288. non_aggregate_keys = {}
  289. for expression in self._non_aggregate_key_exprs:
  290. result = expression.search(response)
  291. set_value_from_jmespath(non_aggregate_keys,
  292. expression.expression,
  293. result)
  294. self._non_aggregate_part = non_aggregate_keys
  295. def _inject_starting_params(self, op_kwargs):
  296. # If the user has specified a starting token we need to
  297. # inject that into the operation's kwargs.
  298. if self._starting_token is not None:
  299. # Don't need to do anything special if there is no starting
  300. # token specified.
  301. next_token = self._parse_starting_token()[0]
  302. self._inject_token_into_kwargs(op_kwargs, next_token)
  303. if self._page_size is not None:
  304. # Pass the page size as the parameter name for limiting
  305. # page size, also known as the limit_key.
  306. op_kwargs[self._limit_key] = self._page_size
  307. def _inject_token_into_kwargs(self, op_kwargs, next_token):
  308. for name, token in next_token.items():
  309. if (token is not None) and (token != 'None'):
  310. op_kwargs[name] = token
  311. elif name in op_kwargs:
  312. del op_kwargs[name]
  313. def _handle_first_request(self, parsed, primary_result_key,
  314. starting_truncation):
  315. # If the payload is an array or string, we need to slice into it
  316. # and only return the truncated amount.
  317. starting_truncation = self._parse_starting_token()[1]
  318. all_data = primary_result_key.search(parsed)
  319. if isinstance(all_data, (list, six.string_types)):
  320. data = all_data[starting_truncation:]
  321. else:
  322. data = None
  323. set_value_from_jmespath(
  324. parsed,
  325. primary_result_key.expression,
  326. data
  327. )
  328. # We also need to truncate any secondary result keys
  329. # because they were not truncated in the previous last
  330. # response.
  331. for token in self.result_keys:
  332. if token == primary_result_key:
  333. continue
  334. sample = token.search(parsed)
  335. if isinstance(sample, list):
  336. empty_value = []
  337. elif isinstance(sample, six.string_types):
  338. empty_value = ''
  339. elif isinstance(sample, (int, float)):
  340. empty_value = 0
  341. else:
  342. empty_value = None
  343. set_value_from_jmespath(parsed, token.expression, empty_value)
  344. return starting_truncation
  345. def _truncate_response(self, parsed, primary_result_key, truncate_amount,
  346. starting_truncation, next_token):
  347. original = primary_result_key.search(parsed)
  348. if original is None:
  349. original = []
  350. amount_to_keep = len(original) - truncate_amount
  351. truncated = original[:amount_to_keep]
  352. set_value_from_jmespath(
  353. parsed,
  354. primary_result_key.expression,
  355. truncated
  356. )
  357. # The issue here is that even though we know how much we've truncated
  358. # we need to account for this globally including any starting
  359. # left truncation. For example:
  360. # Raw response: [0,1,2,3]
  361. # Starting index: 1
  362. # Max items: 1
  363. # Starting left truncation: [1, 2, 3]
  364. # End right truncation for max items: [1]
  365. # However, even though we only kept 1, this is post
  366. # left truncation so the next starting index should be 2, not 1
  367. # (left_truncation + amount_to_keep).
  368. next_token['boto_truncate_amount'] = \
  369. amount_to_keep + starting_truncation
  370. self.resume_token = next_token
  371. def _get_next_token(self, parsed):
  372. if self._more_results is not None:
  373. if not self._more_results.search(parsed):
  374. return {}
  375. next_tokens = {}
  376. for output_token, input_key in \
  377. zip(self._output_token, self._input_token):
  378. next_token = output_token.search(parsed)
  379. # We do not want to include any empty strings as actual tokens.
  380. # Treat them as None.
  381. if next_token:
  382. next_tokens[input_key] = next_token
  383. else:
  384. next_tokens[input_key] = None
  385. return next_tokens
  386. def result_key_iters(self):
  387. teed_results = tee(self, len(self.result_keys))
  388. return [ResultKeyIterator(i, result_key) for i, result_key
  389. in zip(teed_results, self.result_keys)]
  390. def build_full_result(self):
  391. complete_result = {}
  392. for response in self:
  393. page = response
  394. # We want to try to catch operation object pagination
  395. # and format correctly for those. They come in the form
  396. # of a tuple of two elements: (http_response, parsed_responsed).
  397. # We want the parsed_response as that is what the page iterator
  398. # uses. We can remove it though once operation objects are removed.
  399. if isinstance(response, tuple) and len(response) == 2:
  400. page = response[1]
  401. # We're incrementally building the full response page
  402. # by page. For each page in the response we need to
  403. # inject the necessary components from the page
  404. # into the complete_result.
  405. for result_expression in self.result_keys:
  406. # In order to incrementally update a result key
  407. # we need to search the existing value from complete_result,
  408. # then we need to search the _current_ page for the
  409. # current result key value. Then we append the current
  410. # value onto the existing value, and re-set that value
  411. # as the new value.
  412. result_value = result_expression.search(page)
  413. if result_value is None:
  414. continue
  415. existing_value = result_expression.search(complete_result)
  416. if existing_value is None:
  417. # Set the initial result
  418. set_value_from_jmespath(
  419. complete_result, result_expression.expression,
  420. result_value)
  421. continue
  422. # Now both result_value and existing_value contain something
  423. if isinstance(result_value, list):
  424. existing_value.extend(result_value)
  425. elif isinstance(result_value, (int, float, six.string_types)):
  426. # Modify the existing result with the sum or concatenation
  427. set_value_from_jmespath(
  428. complete_result, result_expression.expression,
  429. existing_value + result_value)
  430. merge_dicts(complete_result, self.non_aggregate_part)
  431. if self.resume_token is not None:
  432. complete_result['NextToken'] = self.resume_token
  433. return complete_result
  434. def _parse_starting_token(self):
  435. if self._starting_token is None:
  436. return None
  437. # The starting token is a dict passed as a base64 encoded string.
  438. next_token = self._starting_token
  439. try:
  440. next_token = self._token_decoder.decode(next_token)
  441. index = 0
  442. if 'boto_truncate_amount' in next_token:
  443. index = next_token.get('boto_truncate_amount')
  444. del next_token['boto_truncate_amount']
  445. except (ValueError, TypeError):
  446. next_token, index = self._parse_starting_token_deprecated()
  447. return next_token, index
  448. def _parse_starting_token_deprecated(self):
  449. """
  450. This handles parsing of old style starting tokens, and attempts to
  451. coerce them into the new style.
  452. """
  453. log.debug("Attempting to fall back to old starting token parser. For "
  454. "token: %s" % self._starting_token)
  455. if self._starting_token is None:
  456. return None
  457. parts = self._starting_token.split('___')
  458. next_token = []
  459. index = 0
  460. if len(parts) == len(self._input_token) + 1:
  461. try:
  462. index = int(parts.pop())
  463. except ValueError:
  464. # This doesn't look like a valid old-style token, so we're
  465. # passing it along as an opaque service token.
  466. parts = [self._starting_token]
  467. for part in parts:
  468. if part == 'None':
  469. next_token.append(None)
  470. else:
  471. next_token.append(part)
  472. return self._convert_deprecated_starting_token(next_token), index
  473. def _convert_deprecated_starting_token(self, deprecated_token):
  474. """
  475. This attempts to convert a deprecated starting token into the new
  476. style.
  477. """
  478. len_deprecated_token = len(deprecated_token)
  479. len_input_token = len(self._input_token)
  480. if len_deprecated_token > len_input_token:
  481. raise ValueError("Bad starting token: %s" % self._starting_token)
  482. elif len_deprecated_token < len_input_token:
  483. log.debug("Old format starting token does not contain all input "
  484. "tokens. Setting the rest, in order, as None.")
  485. for i in range(len_input_token - len_deprecated_token):
  486. deprecated_token.append(None)
  487. return dict(zip(self._input_token, deprecated_token))
  488. class Paginator(object):
  489. PAGE_ITERATOR_CLS = PageIterator
  490. def __init__(self, method, pagination_config, model):
  491. self._model = model
  492. self._method = method
  493. self._pagination_cfg = pagination_config
  494. self._output_token = self._get_output_tokens(self._pagination_cfg)
  495. self._input_token = self._get_input_tokens(self._pagination_cfg)
  496. self._more_results = self._get_more_results_token(self._pagination_cfg)
  497. self._non_aggregate_keys = self._get_non_aggregate_keys(
  498. self._pagination_cfg)
  499. self._result_keys = self._get_result_keys(self._pagination_cfg)
  500. self._limit_key = self._get_limit_key(self._pagination_cfg)
  501. @property
  502. def result_keys(self):
  503. return self._result_keys
  504. def _get_non_aggregate_keys(self, config):
  505. keys = []
  506. for key in config.get('non_aggregate_keys', []):
  507. keys.append(jmespath.compile(key))
  508. return keys
  509. def _get_output_tokens(self, config):
  510. output = []
  511. output_token = config['output_token']
  512. if not isinstance(output_token, list):
  513. output_token = [output_token]
  514. for config in output_token:
  515. output.append(jmespath.compile(config))
  516. return output
  517. def _get_input_tokens(self, config):
  518. input_token = self._pagination_cfg['input_token']
  519. if not isinstance(input_token, list):
  520. input_token = [input_token]
  521. return input_token
  522. def _get_more_results_token(self, config):
  523. more_results = config.get('more_results')
  524. if more_results is not None:
  525. return jmespath.compile(more_results)
  526. def _get_result_keys(self, config):
  527. result_key = config.get('result_key')
  528. if result_key is not None:
  529. if not isinstance(result_key, list):
  530. result_key = [result_key]
  531. result_key = [jmespath.compile(rk) for rk in result_key]
  532. return result_key
  533. def _get_limit_key(self, config):
  534. return config.get('limit_key')
  535. def paginate(self, **kwargs):
  536. """Create paginator object for an operation.
  537. This returns an iterable object. Iterating over
  538. this object will yield a single page of a response
  539. at a time.
  540. """
  541. page_params = self._extract_paging_params(kwargs)
  542. return self.PAGE_ITERATOR_CLS(
  543. self._method, self._input_token,
  544. self._output_token, self._more_results,
  545. self._result_keys, self._non_aggregate_keys,
  546. self._limit_key,
  547. page_params['MaxItems'],
  548. page_params['StartingToken'],
  549. page_params['PageSize'],
  550. kwargs)
  551. def _extract_paging_params(self, kwargs):
  552. pagination_config = kwargs.pop('PaginationConfig', {})
  553. max_items = pagination_config.get('MaxItems', None)
  554. if max_items is not None:
  555. max_items = int(max_items)
  556. page_size = pagination_config.get('PageSize', None)
  557. if page_size is not None:
  558. if self._limit_key is None:
  559. raise PaginationError(
  560. message="PageSize parameter is not supported for the "
  561. "pagination interface for this operation.")
  562. input_members = self._model.input_shape.members
  563. limit_key_shape = input_members.get(self._limit_key)
  564. if limit_key_shape.type_name == 'string':
  565. if not isinstance(page_size, six.string_types):
  566. page_size = str(page_size)
  567. else:
  568. page_size = int(page_size)
  569. return {
  570. 'MaxItems': max_items,
  571. 'StartingToken': pagination_config.get('StartingToken', None),
  572. 'PageSize': page_size,
  573. }
  574. class ResultKeyIterator(object):
  575. """Iterates over the results of paginated responses.
  576. Each iterator is associated with a single result key.
  577. Iterating over this object will give you each element in
  578. the result key list.
  579. :param pages_iterator: An iterator that will give you
  580. pages of results (a ``PageIterator`` class).
  581. :param result_key: The JMESPath expression representing
  582. the result key.
  583. """
  584. def __init__(self, pages_iterator, result_key):
  585. self._pages_iterator = pages_iterator
  586. self.result_key = result_key
  587. def __iter__(self):
  588. for page in self._pages_iterator:
  589. results = self.result_key.search(page)
  590. if results is None:
  591. results = []
  592. for result in results:
  593. yield result