| OLD | NEW |
| 1 # Copyright 2016 The Chromium Authors. All rights reserved. | 1 # Copyright 2016 The Chromium Authors. All rights reserved. |
| 2 # Use of this source code is governed by a BSD-style license that can be | 2 # Use of this source code is governed by a BSD-style license that can be |
| 3 # found in the LICENSE file. | 3 # found in the LICENSE file. |
| 4 | 4 |
| 5 """Fetches entities and iterate over and process them.""" | 5 """Fetches entities and iterate over and process them.""" |
| 6 | 6 |
| 7 import os | 7 import os |
| 8 | 8 |
| 9 import remote_api # pylint: disable=W | 9 import remote_api # pylint: disable=W |
| 10 | 10 |
| 11 _DEFAULT_BATCH_SIZE = 1000 | 11 _DEFAULT_BATCH_SIZE = 1000 |
| 12 | 12 |
| 13 | 13 |
| 14 # TODO(crbug.com/662540): Add unittests. | 14 # TODO(crbug.com/662540): Add unittests. |
| 15 def ProjectEntity(entity, fields): # pragma: no cover. | 15 def ProjectEntity(entity, fields): # pragma: no cover. |
| 16 """Projects fields from entity. Returns dict.""" | 16 """Projects fields from entity. Returns dict.""" |
| 17 entity_info = {} | 17 entity_info = {} |
| 18 for field in fields: | 18 for field in fields: |
| 19 if hasattr(entity, field): | 19 if hasattr(entity, field): |
| 20 entity_info[field] = getattr(entity, field) | 20 entity_info[field] = getattr(entity, field) |
| 21 else: | 21 else: |
| 22 entity_info[field] = None | 22 entity_info[field] = None |
| 23 entity_info['id'] = entity.key.id() | 23 entity_info['id'] = entity.key.urlsafe() |
| 24 return entity_info | 24 return entity_info |
| 25 | 25 |
| 26 | 26 |
| 27 # TODO(crbug.com/662540): Add unittests. | 27 # TODO(crbug.com/662540): Add unittests. |
| 28 def Iterate(query, | 28 def Iterate(query, |
| 29 fields, | |
| 30 app_id, | 29 app_id, |
| 30 fields=None, |
| 31 filter_func=None, | 31 filter_func=None, |
| 32 batch_size=_DEFAULT_BATCH_SIZE, | 32 batch_size=_DEFAULT_BATCH_SIZE, |
| 33 batch_run=False): # pragma: no cover. | 33 batch_run=False): # pragma: no cover. |
| 34 """Iterates entities queried by query. | 34 """Iterates entities queried by query. |
| 35 | 35 |
| 36 Args: | 36 Args: |
| 37 query (ndb.Query): The query to fetch entities. | 37 query (ndb.Query): The query to fetch entities. |
| 38 fields (list): Field names of an entity to be projected to a dict. | |
| 39 If a given field name is not available, it is set to None. | |
| 40 'id' is always added by default as an integer. | |
| 41 app_id (str): App engine app id. | 38 app_id (str): App engine app id. |
| 39 fields (list): Field names of an datastore model entity to be projected |
| 40 to a dict. If None provided, return the entities instead of projected |
| 41 dict. |
| 42 filter_func (function): A function that does in memory filtering. | 42 filter_func (function): A function that does in memory filtering. |
| 43 batch_size (int): The number of entities to query at one time. | 43 batch_size (int): The number of entities to query at one time. |
| 44 batch_run (bool): If True, iterate batches of entities, if | 44 batch_run (bool): If True, iterate batches of entities, if |
| 45 False, iterate each entity. | 45 False, iterate each entity. |
| 46 | 46 |
| 47 An exmaple is available in crash_printer/print_crash.py. | 47 An exmaple is available in crash_printer/print_crash.py. |
| 48 """ | 48 """ |
| 49 remote_api.EnableRemoteApi(app_id) | 49 remote_api.EnableRemoteApi(app_id) |
| 50 | 50 |
| 51 cursor = None | 51 cursor = None |
| 52 while True: | 52 while True: |
| 53 entities, next_cursor, more = query.fetch_page(batch_size, | 53 entities, next_cursor, more = query.fetch_page(batch_size, |
| 54 start_cursor=cursor) | 54 start_cursor=cursor) |
| 55 if not more and not entities: | 55 if not more and not entities: |
| 56 break | 56 break |
| 57 | 57 |
| 58 if filter_func: | 58 if filter_func: |
| 59 entities = filter_func(entities) | 59 entities = filter_func(entities) |
| 60 | 60 |
| 61 entities = [ProjectEntity(entity, fields) for entity in entities] | 61 if fields: |
| 62 entities = [ProjectEntity(entity, fields) for entity in entities] |
| 63 |
| 62 if batch_run: | 64 if batch_run: |
| 63 yield entities | 65 yield entities |
| 64 else: | 66 else: |
| 65 for entity in entities: | 67 for entity in entities: |
| 66 yield entity | 68 yield entity |
| 67 | 69 |
| 68 if not more: | 70 if not more: |
| 69 break | 71 break |
| 70 | 72 |
| 71 cursor = next_cursor | 73 cursor = next_cursor |
| OLD | NEW |