| #!/usr/bin/env python |
| # |
| # Copyright 2007 Google Inc. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| |
| |
| |
| """Used render templates for datastore admin.""" |
| |
| |
| import base64 |
| import datetime |
| import logging |
| import os |
| import random |
| |
| from google.appengine.datastore import entity_pb |
| from google.appengine.api import datastore |
| from google.appengine.api import memcache |
| from google.appengine.api import users |
| from google.appengine.datastore import datastore_rpc |
| from google.appengine.ext import db |
| from google.appengine.ext import webapp |
| from google.appengine.ext.datastore_admin import config |
| from google.appengine.ext.db import stats |
| from google.appengine.ext.webapp import _template |
| |
| |
| try: |
| |
| from google.appengine.ext.mapreduce import context |
| from google.appengine.ext.mapreduce import control |
| from google.appengine.ext.mapreduce import model |
| from google.appengine.ext.mapreduce import operation as mr_operation |
| from google.appengine.ext.mapreduce import util |
| except ImportError: |
| |
| from google.appengine._internal.mapreduce import context |
| from google.appengine._internal.mapreduce import control |
| from google.appengine._internal.mapreduce import model |
| from google.appengine._internal.mapreduce import operation as mr_operation |
| from google.appengine._internal.mapreduce import util |
| |
| |
| MEMCACHE_NAMESPACE = '_ah-datastore_admin' |
| XSRF_VALIDITY_TIME = 600 |
| KINDS_AND_SIZES_VAR = 'kinds_and_sizes' |
| MAPREDUCE_MIN_SHARDS = 8 |
| MAPREDUCE_DEFAULT_SHARDS = 32 |
| MAPREDUCE_MAX_SHARDS = 256 |
| RESERVE_KEY_POOL_MAX_SIZE = 1000 |
| |
| DATASTORE_ADMIN_OPERATION_KIND = '_AE_DatastoreAdmin_Operation' |
| BACKUP_INFORMATION_KIND = '_AE_Backup_Information' |
| BACKUP_INFORMATION_FILES_KIND = '_AE_Backup_Information_Kind_Files' |
| BACKUP_INFORMATION_KIND_TYPE_INFO = '_AE_Backup_Information_Kind_Type_Info' |
| DATASTORE_ADMIN_KINDS = (DATASTORE_ADMIN_OPERATION_KIND, |
| BACKUP_INFORMATION_KIND, |
| BACKUP_INFORMATION_FILES_KIND, |
| BACKUP_INFORMATION_KIND_TYPE_INFO) |
| |
| |
| def IsKindNameVisible(kind_name): |
| return not (kind_name.startswith('__') or |
| kind_name in DATASTORE_ADMIN_KINDS or |
| kind_name in model._MAP_REDUCE_KINDS) |
| |
| |
| def RenderToResponse(handler, template_file, template_params): |
| """Render the given template_file using template_vals and write to response. |
| |
| Args: |
| handler: the handler whose response we should render to |
| template_file: the file name only of the template file we are using |
| template_params: the parameters used to render the given template |
| """ |
| template_params = _GetDefaultParams(template_params) |
| |
| |
| |
| |
| |
| |
| |
| |
| handler.response.headers['X-FRAME-OPTIONS'] = ('ALLOW-FROM %s' % |
| config.ADMIN_CONSOLE_URL) |
| template_params['admin_console_url'] = config.ADMIN_CONSOLE_URL |
| |
| rendered = _template.render(_GetTemplatePath(template_file), template_params) |
| handler.response.out.write(rendered) |
| |
| |
| def _GetTemplatePath(template_file): |
| """Return the expected path for the template to render. |
| |
| Args: |
| template_file: simple file name of template to render. |
| |
| Returns: |
| path of template to render. |
| """ |
| return os.path.join( |
| os.path.dirname(__file__), 'templates', template_file) |
| |
| |
| def _GetDefaultParams(template_params): |
| """Update template_params to always contain necessary paths and never None.""" |
| if not template_params: |
| template_params = {} |
| template_params.update({ |
| 'base_path': config.BASE_PATH, |
| 'mapreduce_path': config.MAPREDUCE_PATH, |
| }) |
| return template_params |
| |
| |
| def CreateXsrfToken(action): |
| """Generate a token to be passed with a form for XSRF protection. |
| |
| Args: |
| action: action to restrict token to |
| |
| Returns: |
| suitably random token which is only valid for ten minutes and, if the user |
| is authenticated, is only valid for the user that generated it. |
| """ |
| user_str = _MakeUserStr() |
| |
| token = base64.b64encode( |
| ''.join(chr(int(random.random()*255)) for _ in range(0, 64))) |
| |
| memcache.set(token, |
| (user_str, action), |
| time=XSRF_VALIDITY_TIME, |
| namespace=MEMCACHE_NAMESPACE) |
| |
| return token |
| |
| |
| def ValidateXsrfToken(token, action): |
| """Validate a given XSRF token by retrieving it from memcache. |
| |
| If the token has not been evicted from memcache (past ten minutes) and the |
| user strings are equal, then this is a valid token. |
| |
| Args: |
| token: token to validate from memcache. |
| action: action that token should correspond to |
| |
| Returns: |
| True if the token exists in memcache and the user strings are equal, |
| False otherwise. |
| """ |
| user_str = _MakeUserStr() |
| token_obj = memcache.get(token, namespace=MEMCACHE_NAMESPACE) |
| |
| if not token_obj: |
| return False |
| |
| token_str, token_action = token_obj |
| if user_str != token_str or action != token_action: |
| return False |
| |
| return True |
| |
| |
| def CacheStats(formatted_results): |
| """Cache last retrieved kind size values in memcache. |
| |
| Args: |
| formatted_results: list of dictionaries of the form returnned by |
| main._PresentableKindStats. |
| """ |
| kinds_and_sizes = dict((kind['kind_name'], kind['total_bytes']) |
| for kind in formatted_results) |
| |
| memcache.set(KINDS_AND_SIZES_VAR, |
| kinds_and_sizes, |
| namespace=MEMCACHE_NAMESPACE) |
| |
| |
| def RetrieveCachedStats(): |
| """Retrieve cached kind sizes from last datastore stats call. |
| |
| Returns: |
| Dictionary mapping kind names to total bytes. |
| """ |
| return memcache.get(KINDS_AND_SIZES_VAR, namespace=MEMCACHE_NAMESPACE) |
| |
| |
| def _MakeUserStr(): |
| """Make a user string to use to represent the user. 'noauth' by default.""" |
| user = users.get_current_user() |
| return user.nickname() if user else 'noauth' |
| |
| |
| def GetPrettyBytes(bytes_num, significant_digits=0): |
| """Get a pretty print view of the given number of bytes. |
| |
| This will give a string like 'X MBytes'. |
| |
| Args: |
| bytes_num: the original number of bytes to pretty print. |
| significant_digits: number of digits to display after the decimal point. |
| |
| Returns: |
| A string that has the pretty print version of the given bytes. |
| If bytes_num is to big the string 'Alot' will be returned. |
| """ |
| byte_prefixes = ['', 'K', 'M', 'G', 'T', 'P', 'E'] |
| for i in range(0, 7): |
| exp = i * 10 |
| if bytes_num < 1<<(exp + 10): |
| if i == 0: |
| formatted_bytes = str(bytes_num) |
| else: |
| formatted_bytes = '%.*f' % (significant_digits, |
| (bytes_num * 1.0 / (1<<exp))) |
| if formatted_bytes != '1': |
| plural = 's' |
| else: |
| plural = '' |
| return '%s %sByte%s' % (formatted_bytes, byte_prefixes[i], plural) |
| |
| logging.error('Number too high to convert: %d', bytes_num) |
| return 'Alot' |
| |
| |
| def FormatThousands(value): |
| """Format a numerical value, inserting commas as thousands separators. |
| |
| Args: |
| value: An integer, float, or string representation thereof. |
| If the argument is a float, it is converted to a string using '%.2f'. |
| |
| Returns: |
| A string with groups of 3 digits before the decimal point (if any) |
| separated by commas. |
| |
| NOTE: We don't deal with whitespace, and we don't insert |
| commas into long strings of digits after the decimal point. |
| """ |
| if isinstance(value, float): |
| value = '%.2f' % value |
| else: |
| value = str(value) |
| if '.' in value: |
| head, tail = value.split('.', 1) |
| tail = '.' + tail |
| elif 'e' in value: |
| head, tail = value.split('e', 1) |
| tail = 'e' + tail |
| else: |
| head = value |
| tail = '' |
| sign = '' |
| if head.startswith('-'): |
| sign = '-' |
| head = head[1:] |
| while len(head) > 3: |
| tail = ',' + head[-3:] + tail |
| head = head[:-3] |
| return sign + head + tail |
| |
| |
| def TruncDelta(delta): |
| """Strips microseconds from a timedelta.""" |
| return datetime.timedelta(days=delta.days, seconds=delta.seconds) |
| |
| |
| def GetPrintableStrs(namespace, kinds): |
| """Returns tuples describing affected kinds and namespace. |
| |
| Args: |
| namespace: namespace being targeted. |
| kinds: list of kinds being targeted. |
| |
| Returns: |
| (namespace_str, kind_str) tuple used for display to user. |
| """ |
| namespace_str = namespace or '' |
| if kinds: |
| kind_str = 'all %s entities' % ', '.join(kinds) |
| else: |
| kind_str = '' |
| return (namespace_str, kind_str) |
| |
| |
| def ParseKindsAndSizes(kinds): |
| """Parses kind|size list and returns template parameters. |
| |
| Args: |
| kinds: list of kinds to process. |
| |
| Returns: |
| sizes_known: whether or not all kind objects have known sizes. |
| size_total: total size of objects with known sizes. |
| len(kinds) - 2: for template rendering of greater than 3 kinds. |
| """ |
| sizes_known = True |
| size_total = 0 |
| kinds_and_sizes = RetrieveCachedStats() |
| |
| if kinds_and_sizes: |
| for kind in kinds: |
| if kind in kinds_and_sizes: |
| size_total += kinds_and_sizes[kind] |
| else: |
| sizes_known = False |
| else: |
| sizes_known = False |
| |
| if size_total: |
| size_total = GetPrettyBytes(size_total) |
| |
| return sizes_known, size_total, len(kinds) - 2 |
| |
| |
| def _CreateDatastoreConfig(): |
| """Create datastore config for use during datastore admin operations.""" |
| return datastore_rpc.Configuration(force_writes=True) |
| |
| |
| def GenerateHomeUrl(request): |
| """Generates a link to the Datastore Admin main page. |
| |
| Primarily intended to be used for cancel buttons or links on error pages. To |
| avoid any XSS security vulnerabilities the URL should not use any |
| user-defined strings (unless proper precautions are taken). |
| |
| Args: |
| request: the webapp.Request object (to determine if certain query |
| parameters need to be used). |
| |
| Returns: |
| domain-relative URL for the main Datastore Admin page. |
| """ |
| datastore_admin_home = config.BASE_PATH |
| if request and request.get('run_as_a_service'): |
| datastore_admin_home += '?run_as_a_service=True' |
| return datastore_admin_home |
| |
| |
| class MapreduceDoneHandler(webapp.RequestHandler): |
| """Handler to delete data associated with successful MapReduce jobs.""" |
| |
| SUFFIX = 'mapreduce_done' |
| |
| def post(self): |
| """Mapreduce done callback to delete job data if it was successful.""" |
| if 'Mapreduce-Id' in self.request.headers: |
| mapreduce_id = self.request.headers['Mapreduce-Id'] |
| mapreduce_state = model.MapreduceState.get_by_job_id(mapreduce_id) |
| mapreduce_params = mapreduce_state.mapreduce_spec.params |
| |
| db_config = _CreateDatastoreConfig() |
| if mapreduce_state.result_status == model.MapreduceState.RESULT_SUCCESS: |
| operation_key = mapreduce_params.get( |
| DatastoreAdminOperation.PARAM_DATASTORE_ADMIN_OPERATION) |
| if operation_key is None: |
| logging.error('Done callback for job %s without operation key.', |
| mapreduce_id) |
| else: |
| |
| def tx(): |
| operation = DatastoreAdminOperation.get(operation_key) |
| if mapreduce_id in operation.active_job_ids: |
| operation.active_jobs -= 1 |
| operation.completed_jobs += 1 |
| operation.active_job_ids.remove(mapreduce_id) |
| if not operation.active_jobs: |
| if operation.status == DatastoreAdminOperation.STATUS_ACTIVE: |
| operation.status = DatastoreAdminOperation.STATUS_COMPLETED |
| db.delete(DatastoreAdminOperationJob.all().ancestor(operation), |
| config=db_config) |
| operation.put(config=db_config) |
| if 'done_callback_handler' in mapreduce_params: |
| done_callback_handler = util.for_name( |
| mapreduce_params['done_callback_handler']) |
| if done_callback_handler: |
| done_callback_handler(operation, mapreduce_id, mapreduce_state) |
| else: |
| logging.error('done_callbackup_handler %s was not found', |
| mapreduce_params['done_callback_handler']) |
| db.run_in_transaction(tx) |
| if config.CLEANUP_MAPREDUCE_STATE: |
| keys = [] |
| keys = model.ShardState.calculate_keys_by_mapreduce_state( |
| mapreduce_state) |
| |
| keys.append(model.MapreduceControl.get_key_by_job_id(mapreduce_id)) |
| db.delete(keys, config=db_config) |
| db.delete(mapreduce_state, config=db_config) |
| logging.info('State for successful job %s was deleted.', mapreduce_id) |
| else: |
| logging.info('Job %s was not successful so no state was deleted.', |
| mapreduce_id) |
| else: |
| logging.error('Done callback called without Mapreduce Id.') |
| |
| |
| class Error(Exception): |
| """Base DatastoreAdmin error type.""" |
| |
| |
| |
| class DatastoreAdminOperation(db.Model): |
| """An entity to keep progress and status of datastore admin operation.""" |
| STATUS_CREATED = 'Created' |
| STATUS_ACTIVE = 'Active' |
| STATUS_COMPLETED = 'Completed' |
| STATUS_FAILED = 'Failed' |
| STATUS_ABORTED = 'Aborted' |
| |
| |
| PARAM_DATASTORE_ADMIN_OPERATION = 'datastore_admin_operation' |
| DEFAULT_LAST_UPDATED_VALUE = datetime.datetime(1970, 1, 1) |
| |
| description = db.TextProperty() |
| status = db.StringProperty(default=STATUS_CREATED) |
| active_jobs = db.IntegerProperty(default=0) |
| active_job_ids = db.StringListProperty() |
| completed_jobs = db.IntegerProperty(default=0) |
| last_updated = db.DateTimeProperty(default=DEFAULT_LAST_UPDATED_VALUE, |
| auto_now=True) |
| status_info = db.StringProperty(default='', indexed=False) |
| service_job_id = db.StringProperty() |
| |
| @classmethod |
| def kind(cls): |
| return DATASTORE_ADMIN_OPERATION_KIND |
| |
| |
| class DatastoreAdminOperationJob(db.Model): |
| """An entity to keep track of started jobs to ensure idempotency. |
| |
| This entity can be used during spawning additional jobs. It is |
| always stored as a child entity of DatastoreAdminOperation. |
| Entity key name is job unique id. |
| """ |
| pass |
| |
| |
| def StartOperation(description): |
| """Start datastore admin operation. |
| |
| Args: |
| description: operation description to be displayed to user. |
| |
| Returns: |
| an instance of DatastoreAdminOperation. |
| """ |
| |
| |
| |
| operation = DatastoreAdminOperation( |
| description=description, |
| id=db.allocate_ids( |
| db.Key.from_path(DatastoreAdminOperation.kind(), 1), 1)[0]) |
| operation.put(config=_CreateDatastoreConfig()) |
| return operation |
| |
| |
| @db.non_transactional(allow_existing=False) |
| def StartMap(operation_key, |
| job_name, |
| handler_spec, |
| reader_spec, |
| writer_spec, |
| mapper_params, |
| mapreduce_params=None, |
| queue_name=None, |
| shard_count=MAPREDUCE_DEFAULT_SHARDS): |
| """Start map as part of datastore admin operation. |
| |
| Will increase number of active jobs inside the operation and start new map. |
| |
| Args: |
| operation_key: Key of the DatastoreAdminOperation for current operation. |
| job_name: Map job name. |
| handler_spec: Map handler specification. |
| reader_spec: Input reader specification. |
| writer_spec: Output writer specification. |
| mapper_params: Custom mapper parameters. |
| mapreduce_params: Custom mapreduce parameters. |
| queue_name: the name of the queue that will be used by the M/R. |
| shard_count: the number of shards the M/R will try to use. |
| |
| Returns: |
| resulting map job id as string. |
| """ |
| |
| if not mapreduce_params: |
| mapreduce_params = {} |
| mapreduce_params[DatastoreAdminOperation.PARAM_DATASTORE_ADMIN_OPERATION] = ( |
| str(operation_key)) |
| mapreduce_params['done_callback'] = '%s/%s' % (config.BASE_PATH, |
| MapreduceDoneHandler.SUFFIX) |
| if queue_name is not None: |
| mapreduce_params['done_callback_queue'] = queue_name |
| mapreduce_params['force_writes'] = 'True' |
| |
| def tx(is_xg_transaction): |
| """Start MapReduce job and update datastore admin state. |
| |
| Args: |
| is_xg_transaction: True if we are running inside a xg-enabled |
| transaction, else False if we are running inside a non-xg-enabled |
| transaction (which means the datastore admin state is updated in one |
| transaction and the MapReduce job in an indepedent transaction). |
| Returns: |
| result MapReduce job id as a string. |
| """ |
| job_id = control.start_map( |
| job_name, handler_spec, reader_spec, |
| mapper_params, |
| output_writer_spec=writer_spec, |
| mapreduce_parameters=mapreduce_params, |
| base_path=config.MAPREDUCE_PATH, |
| shard_count=shard_count, |
| in_xg_transaction=is_xg_transaction, |
| queue_name=queue_name) |
| operation = DatastoreAdminOperation.get(operation_key) |
| operation.status = DatastoreAdminOperation.STATUS_ACTIVE |
| operation.active_jobs += 1 |
| operation.active_job_ids = list(set(operation.active_job_ids + [job_id])) |
| operation.put(config=_CreateDatastoreConfig()) |
| return job_id |
| |
| |
| |
| |
| |
| |
| datastore_type = datastore_rpc._GetDatastoreType() |
| |
| if datastore_type != datastore_rpc.BaseConnection.MASTER_SLAVE_DATASTORE: |
| return db.run_in_transaction_options( |
| db.create_transaction_options(xg=True), tx, True) |
| else: |
| return db.run_in_transaction(tx, False) |
| |
| |
| def RunMapForKinds(operation_key, |
| kinds, |
| job_name_template, |
| handler_spec, |
| reader_spec, |
| writer_spec, |
| mapper_params, |
| mapreduce_params=None, |
| queue_name=None, |
| max_shard_count=None): |
| """Run mapper job for all entities in specified kinds. |
| |
| Args: |
| operation_key: The key of the DatastoreAdminOperation to record all jobs. |
| kinds: list of entity kinds as strings. |
| job_name_template: template for naming individual mapper jobs. Can |
| reference %(kind)s and %(namespace)s formatting variables. |
| handler_spec: mapper handler specification. |
| reader_spec: reader specification. |
| writer_spec: writer specification. |
| mapper_params: custom parameters to pass to mapper. |
| mapreduce_params: dictionary parameters relevant to the whole job. |
| queue_name: the name of the queue that will be used by the M/R. |
| max_shard_count: maximum value for shards count. |
| |
| Returns: |
| Ids of all started mapper jobs as list of strings. |
| """ |
| jobs = [] |
| try: |
| for kind in kinds: |
| mapper_params['entity_kind'] = kind |
| job_name = job_name_template % {'kind': kind, 'namespace': |
| mapper_params.get('namespace', '')} |
| shard_count = GetShardCount(kind, max_shard_count) |
| jobs.append(StartMap(operation_key, job_name, handler_spec, reader_spec, |
| writer_spec, mapper_params, mapreduce_params, |
| queue_name=queue_name, shard_count=shard_count)) |
| return jobs |
| |
| except BaseException, ex: |
| AbortAdminOperation(operation_key, |
| _status=DatastoreAdminOperation.STATUS_FAILED, |
| _status_info='%s: %s' % (ex.__class__.__name__, ex)) |
| raise |
| |
| |
| def GetShardCount(kind, max_shard_count=None): |
| stat = stats.KindStat.all().filter('kind_name =', kind).get() |
| if stat: |
| |
| shard_count = min(max(MAPREDUCE_MIN_SHARDS, |
| stat.bytes // (32 * 1024 * 1024)), |
| MAPREDUCE_MAX_SHARDS) |
| if max_shard_count and max_shard_count < shard_count: |
| shard_count = max_shard_count |
| return shard_count |
| |
| return MAPREDUCE_DEFAULT_SHARDS |
| |
| |
| def AbortAdminOperation(operation_key, |
| _status=DatastoreAdminOperation.STATUS_ABORTED, |
| _status_info=''): |
| """Aborts active jobs.""" |
| operation = DatastoreAdminOperation.get(operation_key) |
| operation.status = _status |
| operation.status_info = _status_info |
| operation.put(config=_CreateDatastoreConfig()) |
| for job in operation.active_job_ids: |
| logging.info('Aborting Job %s', job) |
| model.MapreduceControl.abort(job, config=_CreateDatastoreConfig()) |
| |
| |
| def get_kind_from_entity_pb(entity): |
| element_list = entity.key().path().element_list() |
| return element_list[-1].type() if element_list else None |
| |
| |
| def FixKeys(entity_proto, app_id): |
| """Go over keys in the given entity and update the application id. |
| |
| Args: |
| entity_proto: An EntityProto to be fixed up. All identifiable keys in the |
| proto will have the 'app' field reset to match app_id. |
| app_id: The desired application id, typically os.getenv('APPLICATION_ID'). |
| """ |
| |
| def FixKey(mutable_key): |
| mutable_key.set_app(app_id) |
| |
| def FixPropertyList(property_list): |
| for prop in property_list: |
| prop_value = prop.mutable_value() |
| if prop_value.has_referencevalue(): |
| FixKey(prop_value.mutable_referencevalue()) |
| elif prop.meaning() == entity_pb.Property.ENTITY_PROTO: |
| embedded_entity_proto = entity_pb.EntityProto() |
| try: |
| embedded_entity_proto.ParsePartialFromString(prop_value.stringvalue()) |
| except Exception: |
| logging.exception('Failed to fix-keys for property %s of %s', |
| prop.name(), |
| entity_proto.key()) |
| else: |
| FixKeys(embedded_entity_proto, app_id) |
| prop_value.set_stringvalue( |
| embedded_entity_proto.SerializePartialToString()) |
| |
| |
| if entity_proto.has_key() and entity_proto.key().path().element_size(): |
| FixKey(entity_proto.mutable_key()) |
| |
| FixPropertyList(entity_proto.property_list()) |
| FixPropertyList(entity_proto.raw_property_list()) |
| |
| |
| class ReserveKeyPool(object): |
| """Mapper pool which buffers keys with ids to reserve. |
| |
| Runs v4 AllocateIds rpc(s) when flushed. |
| """ |
| |
| def __init__(self): |
| self.keys = [] |
| |
| def reserve_key(self, key): |
| for id_or_name in key.to_path()[1::2]: |
| if isinstance(id_or_name, (int, long)): |
| self.keys.append(key) |
| if len(self.keys) >= RESERVE_KEY_POOL_MAX_SIZE: |
| self.flush() |
| return |
| |
| def flush(self): |
| |
| datastore._GetConnection()._reserve_keys(self.keys) |
| self.keys = [] |
| |
| |
| class ReserveKey(mr_operation.Operation): |
| """Mapper operation to reserve key ids.""" |
| |
| def __init__(self, key): |
| self.key = key |
| self.app_id = key.app() |
| self.pool_id = 'reserve_key_%s_pool' % self.app_id |
| |
| def __call__(self, ctx): |
| pool = ctx.get_pool(self.pool_id) |
| if not pool: |
| pool = ReserveKeyPool() |
| ctx.register_pool(self.pool_id, pool) |
| pool.reserve_key(self.key) |
| |
| |
| class PutPool(context.Pool): |
| """A trimmed copy of the MutationPool class. |
| |
| Properties: |
| puts: a list of entities to put to datastore. |
| max_entity_count: maximum number of entities before flushing it to db. |
| """ |
| POOL_NAME = 'put_pool' |
| |
| def __init__(self, max_entity_count=context.MAX_ENTITY_COUNT): |
| """Constructor. |
| |
| Args: |
| max_entity_count: maximum number of entities before flushing it to db. |
| """ |
| self.max_entity_count = max_entity_count |
| self.puts = [] |
| |
| def Put(self, entity): |
| """Registers entity to put to datastore. |
| |
| Args: |
| entity: The EntityProto for the entity to be put. |
| """ |
| if len(self.puts) >= self.max_entity_count: |
| self.flush() |
| self.puts.append(entity) |
| |
| def flush(self): |
| """Flush all puts to datastore.""" |
| if self.puts: |
| datastore_rpc.Connection( |
| config=datastore_rpc.Configuration(deadline=60)).put(self.puts) |
| self.puts = [] |
| |
| |
| class Put(mr_operation.Operation): |
| """Mapper operation to batch puts.""" |
| |
| def __init__(self, entity): |
| """Constructor. |
| |
| Args: |
| entity: The EntityProto of the entity to put. |
| """ |
| self.entity = entity |
| |
| def __call__(self, ctx): |
| pool = ctx.get_pool(PutPool.POOL_NAME) |
| if not pool: |
| pool = PutPool( |
| max_entity_count=(context.MAX_ENTITY_COUNT/(2**ctx.task_retry_count))) |
| ctx.register_pool(PutPool.POOL_NAME, pool) |
| pool.Put(self.entity) |
