files.new_contents: import array
  import datetime
  import itertools
  import logging
  import md5
  import sys
  import uuid
  
  from tornado import gen
  from tornado.ioloop import IOLoop
  
  from appscale.datastore import dbconstants, helper_functions
  
  from appscale.common.datastore_index import DatastoreIndex, merge_indexes
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from kazoo.client import KazooState
  from appscale.datastore.dbconstants import (
    APP_ENTITY_SCHEMA, BadRequest, ID_KEY_LENGTH, InternalError, MAX_TX_DURATION,
    Timeout
  )
  from appscale.datastore.cassandra_env.cassandra_interface import (
    batch_size, LARGE_BATCH_THRESHOLD)
  from appscale.datastore.cassandra_env.entity_id_allocator import EntityIDAllocator
  from appscale.datastore.cassandra_env.entity_id_allocator import ScatteredAllocator
  from appscale.datastore.cassandra_env.large_batch import BatchNotApplied
  from appscale.datastore.cassandra_env.utils import deletions_for_entity
  from appscale.datastore.cassandra_env.utils import mutations_for_entity
  from appscale.datastore.index_manager import IndexInaccessible
  from appscale.datastore.taskqueue_client import EnqueueError, TaskQueueClient
  from appscale.datastore.utils import _FindIndexToUse
  from appscale.datastore.utils import clean_app_id
  from appscale.datastore.utils import decode_path
  from appscale.datastore.utils import encode_entity_table_key
  from appscale.datastore.utils import encode_index_pb
  from appscale.datastore.utils import encode_path_from_filter
  from appscale.datastore.utils import get_composite_index_keys
  from appscale.datastore.utils import get_entity_key
  from appscale.datastore.utils import get_entity_kind
  from appscale.datastore.utils import get_index_key_from_params
  from appscale.datastore.utils import get_kind_key
  from appscale.datastore.utils import group_for_key
  from appscale.datastore.utils import kind_from_encoded_key
  from appscale.datastore.utils import reference_property_to_reference
  from appscale.datastore.utils import UnprocessedQueryCursor
  from appscale.datastore.range_iterator import RangeExhausted, RangeIterator
  from appscale.datastore.zkappscale import entity_lock
  from appscale.datastore.zkappscale import zktransaction
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.api import datastore_errors
  from google.appengine.api.datastore_distributed import _MAX_ACTIONS_PER_TXN
  from google.appengine.datastore import appscale_stub_util
  from google.appengine.datastore import datastore_pb
  from google.appengine.datastore import datastore_index
  from google.appengine.datastore import entity_pb
  from google.appengine.datastore import sortable_pb_encoder
  from google.appengine.datastore.datastore_stub_util import IdToCounter
  from google.appengine.datastore.datastore_stub_util import SEQUENTIAL
  from google.appengine.runtime import apiproxy_errors
  from google.appengine.ext import db
  from google.appengine.ext.db.metadata import Namespace
  from google.net.proto.ProtocolBuffer import ProtocolBufferDecodeError
  
  logger = logging.getLogger(__name__)
  
  
  class DatastoreDistributed():
    """ AppScale persistent layer for the datastore API. It is the
        replacement for the AppServers to persist their data into
        a distributed datastore instead of a flat file.
    """
    # Max number of results for a query
    _MAXIMUM_RESULTS = 10000
  
    # Maximum amount of filter and orderings allowed within a query
    _MAX_QUERY_COMPONENTS = 63
  
    # For enabling and disabling range inclusivity
    _ENABLE_INCLUSIVITY = True
    _DISABLE_INCLUSIVITY = False
  
    # Delimiter between app names and namespace and the rest of an entity key
    _NAMESPACE_SEPARATOR = dbconstants.KEY_DELIMITER
  
    # Delimiter between parameters in index keys.
    _SEPARATOR = dbconstants.KEY_DELIMITER
  
    # This is the terminating string for range queries
    _TERM_STRING = dbconstants.TERMINATING_STRING
  
    # Smallest possible value that is considered non-null and indexable.
    MIN_INDEX_VALUE = '\x01'
  
    # When assigning the first allocated ID, give this value
    _FIRST_VALID_ALLOCATED_ID = 1
  
    # The key we use to lock for allocating new IDs
    _ALLOCATE_ROOT_KEY = "__allocate__"
  
    # Number of times to retry acquiring a lock for non transactions.
    NON_TRANS_LOCK_RETRY_COUNT = 5
  
    # How long to wait before retrying to grab a lock
    LOCK_RETRY_TIME = .5
  
    # Maximum number of allowed composite indexes any one application can
    # register.
    _MAX_NUM_INDEXES = dbconstants.MAX_NUMBER_OF_COMPOSITE_INDEXES
  
    # The position of the prop name when splitting an index entry by the
    # delimiter.
    PROP_NAME_IN_SINGLE_PROP_INDEX = 3
  
    # The cassandra index column that stores the reference to the entity.
    INDEX_REFERENCE_COLUMN = 'reference'
  
    # The number of entities to fetch at a time when updating indices.
    BATCH_SIZE = 100
  
    def __init__(self, datastore_batch, transaction_manager, zookeeper=None,
                 log_level=logging.INFO, taskqueue_locations=()):
      """
         Constructor.
  
       Args:
         datastore_batch: A reference to the batch datastore interface.
         zookeeper: A reference to the zookeeper interface.
      """
      class_name = self.__class__.__name__
      self.logger = logging.getLogger(class_name)
      self.logger.setLevel(log_level)
  
      assert datastore_batch.valid_data_version_sync()
  
      self.logger.info('Starting {}'.format(class_name))
  
      # datastore accessor used by this class to do datastore operations.
      self.datastore_batch = datastore_batch
  
      # zookeeper instance for accesing ZK functionality.
      self.zookeeper = zookeeper
  
      # Maintain a scattered allocator for each project.
      self.scattered_allocators = {}
  
      # Maintain a sequential allocator for each project.
      self.sequential_allocators = {}
  
      self.taskqueue_client = TaskQueueClient(taskqueue_locations)
      self.transaction_manager = transaction_manager
      self.index_manager = None
      self.zookeeper.handle.add_listener(self._zk_state_listener)
  
    def get_limit(self, query):
      """ Returns the limit that should be used for the given query.
  
      Args:
        query: A datastore_pb.Query.
      Returns:
        An int, the limit to be used when accessing the datastore.
      """
      limit = self._MAXIMUM_RESULTS
      if query.has_count():
        limit = min(query.count(), self._MAXIMUM_RESULTS)
  
      if query.has_limit():
        limit = min(query.limit(), limit)
  
      if query.has_offset():
        limit = limit + min(query.offset(), self._MAXIMUM_RESULTS)
      # We can not scan with 0 or less, hence we set it to one.
      if limit <= 0:
        limit = 1
      return limit
  
    @staticmethod
    def __decode_index_str(value, prop_value):
      """ Takes an encoded string and converts it to a PropertyValue.
  
      Args:
        value: An encoded str.
        prop_value: PropertyValue to fill in.
      """
      value = str(value).replace('\x01\x01', '\x00').replace('\x01\x02', '\x01')
      decoded_value = sortable_pb_encoder.Decoder(
        array.array('B', str(value)))
      prop_value.Merge(decoded_value)
  
    @staticmethod
    def validate_app_id(app_id):
      """ Verify that this is the stub for app_id.
  
      Args:
        app_id: An application ID.
      Raises:
        AppScaleBadArg: If the application id is not set.
      """
      if not app_id:
        raise dbconstants.AppScaleBadArg("Application name must be set")
  
    @staticmethod
    def validate_key(key):
      """ Validate this key by checking to see if it has a name or id.
  
      Args:
        key: entity_pb.Reference
      Raises:
        datastore_errors.BadRequestError: if the key is invalid
        TypeError: if key is not of entity_pb.Reference
      """
  
      if not isinstance(key, entity_pb.Reference):
        raise TypeError("Expected type Reference")
  
      DatastoreDistributed.validate_app_id(key.app())
  
      for elem in key.path().element_list():
        if elem.has_id() and elem.has_name():
          raise datastore_errors.BadRequestError(
              'Each key path element should have id or name but not both: {0}' \
              .format(key))
  
    def get_table_prefix(self, data):
      """ Returns the namespace prefix for a query.
  
      Args:
        data: An Entity, Key or Query PB, or an (app_id, ns) tuple.
      Returns:
        A valid table prefix.
      """
      if isinstance(data, entity_pb.EntityProto):
        app_id = clean_app_id(data.key().app())
        namespace = data.key().name_space()
      elif isinstance(data, tuple):
        app_id = data[0]
        namespace = data[1]
      else:
        app_id = clean_app_id(data.app())
        namespace = data.name_space()
  
      return self._SEPARATOR.join([app_id, namespace])
  
    @staticmethod
    def get_ancestor_key_from_ent_key(ent_key):
      """ Get the key string for the ancestor portion of a composite key.
  
      Args:
        ent_key: A string of the entire path of an entity.
      Returns:
        A str of the path of the ancestor.
      """
      ancestor = ""
      tokens = str(ent_key).split(dbconstants.KIND_SEPARATOR)
      # Strip off the empty placeholder and also do not include the last kind.
      for token in tokens[:-2]:
        ancestor += token + dbconstants.KIND_SEPARATOR
      return ancestor
  
    @staticmethod
    def get_composite_index_key(index, entity, position_list=None,
      filters=None):
      """ Creates a key to the composite index table for a given entity
      for a composite cursor.
  
      Keys are built as such:
        app_id/ns/composite_id/ancestor/valuevaluevalue..../entity_key
      Components explained:
      ns: The namespace of the entity.
      composite_id: The composite ID assigned to this index upon creation.
      ancestor: The root ancestor path (only if the query this index is for
        has an ancestor)
      value(s): The string representation of mulitiple properties.
      entity_key: The entity key (full path) used as a means of having a unique
        identifier. This prevents two entities with the same values from
        colliding.
  
      Args:
        index: A datstore_pb.CompositeIndex.
        entity: A entity_pb.EntityProto.
        position_list: A list of datastore_pb.CompiledCursor_Position items.
          Contains values for property items from a cursor.
        filters: A list of datastore_pb.Query_Filters, used to attain equality
          values not present in position_list.
      Returns:
        A string representing a key to the composite table.
      """
      composite_id = index.id()
      definition = index.definition()
      app_id = clean_app_id(entity.key().app())
      name_space = entity.key().name_space()
      ent_key = encode_index_pb(entity.key().path())
      pre_comp_index_key = "{0}{1}{2}{4}{3}{4}".format(app_id,
        DatastoreDistributed._NAMESPACE_SEPARATOR, name_space, composite_id,
        DatastoreDistributed._SEPARATOR)
      if definition.ancestor() == 1:
        ancestor = DatastoreDistributed.get_ancestor_key_from_ent_key(ent_key)
        pre_comp_index_key += "{0}{1}".format(ancestor,
          DatastoreDistributed._SEPARATOR)
  
      value_dict = {}
      for prop in entity.property_list():
        value_dict[prop.name()] = \
          str(encode_index_pb(prop.value()))
  
      # Position list and filters are used if we're creating a composite
      # key for a cursor.
      if position_list:
        for indexvalue in position_list[0].indexvalue_list():
          value_dict[indexvalue.property()] = \
            str(encode_index_pb(indexvalue.value()))
      if filters:
        for filt in filters:
          if filt.op() == datastore_pb.Query_Filter.EQUAL:
            value_dict[filt.property(0).name()] = \
              str(encode_index_pb(filt.property(0).value()))
  
      index_value = ""
      for prop in definition.property_list():
        name = prop.name()
        value = ''
        if name in value_dict:
          value = value_dict[name]
        elif name == "__key__":
          value = ent_key
        else:
          logger.warning("Given entity {0} is missing a property value {1}.".\
            format(entity, prop.name()))
        if prop.direction() == entity_pb.Index_Property.DESCENDING:
          value = helper_functions.reverse_lex(value)
  
        index_value += str(value) + DatastoreDistributed._SEPARATOR
  
      # We append the ent key to have unique keys if entities happen
      # to share the same index values (and ancestor).
      composite_key = "{0}{1}{2}".format(pre_comp_index_key, index_value,
        ent_key)
      return composite_key
  
    @gen.coroutine
    def insert_composite_indexes(self, entities, composite_indexes):
      """ Creates composite indexes for a set of entities.
  
      Args:
        entities: A list entities.
        composite_indexes: A list of datastore_pb.CompositeIndex.
      """
      if not composite_indexes:
        return
      row_keys = []
      row_values = {}
      # Create default composite index for all entities. Here we take each
      # of the properties in one
      for ent in entities:
        for index_def in composite_indexes:
          # Skip any indexes if the kind does not match.
          kind = get_entity_kind(ent.key())
          if index_def.definition().entity_type() != kind:
            continue
  
          # Make sure the entity contains the required entities for the composite
          # definition.
          prop_name_def_list = [index_prop.name() for index_prop in \
            index_def.definition().property_list()]
          all_prop_names_in_ent = [prop.name() for prop in \
            ent.property_list()]
          has_values = True
          for index_prop_name in prop_name_def_list:
            if index_prop_name not in all_prop_names_in_ent:
              has_values = False
            # Special property name which does not show up in the list but
            # is apart of the key of the entity.
            if index_prop_name == "__key__":
              has_values = True
          if not has_values:
            continue
  
          # Get the composite index key.
          composite_index_keys = get_composite_index_keys(index_def, ent)
          row_keys.extend(composite_index_keys)
  
          # Get the reference value for the composite table.
          entity_key = str(encode_index_pb(ent.key().path()))
          prefix = self.get_table_prefix(ent.key())
          reference = "{0}{1}{2}".format(prefix, self._SEPARATOR,  entity_key)
          for composite_key in composite_index_keys:
            row_values[composite_key] = {'reference': reference}
  
      yield self.datastore_batch.batch_put_entity(
        dbconstants.COMPOSITE_TABLE, row_keys,
        dbconstants.COMPOSITE_SCHEMA, row_values)
  
    @gen.coroutine
    def delete_composite_index_metadata(self, app_id, index):
      """ Deletes a index for the given application identifier.
  
      Args:
        app_id: A string representing the application identifier.
        index: A entity_pb.CompositeIndex object.
      """
      self.logger.info('Deleting composite index:\n{}'.format(index))
      try:
        project_index_manager = self.index_manager.projects[app_id]
      except KeyError:
        raise BadRequest('project_id: {} not found'.format(app_id))
  
      # TODO: Remove actual index entries.
      project_index_manager.delete_index_definition(index.id())
  
    @gen.coroutine
    def create_composite_index(self, app_id, index):
      """ Stores a new index for the given application identifier.
  
      Args:
        app_id: A string representing the application identifier.
        index: A entity_pb.CompositeIndex object.
      Returns:
        A unique number representing the composite index ID.
      """
      new_index = DatastoreIndex.from_pb(index)
  
      # The ID must be a positive number that fits in a signed 64-bit int.
      new_index.id = uuid.uuid1().int >> 65
      merge_indexes(self.zookeeper.handle, app_id, [new_index])
      raise gen.Return(new_index.id)
  
    @gen.coroutine
    def update_composite_index(self, app_id, index):
      """ Updates an index for a given app ID.
  
      Args:
        app_id: A string containing the app ID.
        index: An entity_pb.CompositeIndex object.
      """
      self.logger.info('Updating index: {}'.format(index))
      entries_updated = 0
      entity_type = index.definition().entity_type()
  
      # TODO: Adjust prefix based on ancestor.
      prefix = '{app}{delimiter}{entity_type}{kind_separator}'.format(
        app=app_id,
        delimiter=self._SEPARATOR * 2,
        entity_type=entity_type,
        kind_separator=dbconstants.KIND_SEPARATOR,
      )
      start_row = prefix
      end_row = prefix + self._TERM_STRING
      start_inclusive = True
  
      while True:
        # Fetch references from the kind table since entity keys can have a
        # parent prefix.
        references = yield self.datastore_batch.range_query(
          table_name=dbconstants.APP_KIND_TABLE,
          column_names=dbconstants.APP_KIND_SCHEMA,
          start_key=start_row,
          end_key=end_row,
          limit=self.BATCH_SIZE,
          offset=0,
          start_inclusive=start_inclusive,
        )
  
        pb_entities = yield self.__fetch_entities(references)
        entities = [entity_pb.EntityProto(entity) for entity in pb_entities]
  
        yield self.insert_composite_indexes(entities, [index])
        entries_updated += len(entities)
  
        # If we fetched fewer references than we asked for, we're done.
        if len(references) < self.BATCH_SIZE:
          break
  
        start_row = references[-1].keys()[0]
        start_inclusive = self._DISABLE_INCLUSIVITY
  
      self.logger.info('Updated {} index entries.'.format(entries_updated))
  
    @gen.coroutine
    def allocate_size(self, project, size):
      """ Allocates a block of IDs for a project.
  
      Args:
        project: A string specifying the project ID.
        size: An integer specifying the number of IDs to reserve.
      Returns:
        A tuple of integers specifying the start and end ID.
      """
      if project not in self.sequential_allocators:
        self.sequential_allocators[project] = EntityIDAllocator(
          self.datastore_batch.session, project)
  
      allocator = self.sequential_allocators[project]
      start_id, end_id = yield allocator.allocate_size(size)
      raise gen.Return((start_id, end_id))
  
    @gen.coroutine
    def allocate_max(self, project, max_id):
      """ Reserves all IDs up to the one given.
  
      Args:
        project: A string specifying the project ID.
        max_id: An integer specifying the maximum ID to allocated.
      Returns:
        A tuple of integers specifying the start and end ID.
      """
      if project not in self.sequential_allocators:
        self.sequential_allocators[project] = EntityIDAllocator(
          self.datastore_batch.session, project)
  
      allocator = self.sequential_allocators[project]
      start_id, end_id = yield allocator.allocate_max(max_id)
      raise gen.Return((start_id, end_id))
  
    @gen.coroutine
    def reserve_ids(self, project_id, ids):
      """ Ensures the given IDs are not re-allocated.
  
      Args:
        project_id: A string specifying the project ID.
        ids: An iterable of integers specifying entity IDs.
      """
      if project_id not in self.sequential_allocators:
        self.sequential_allocators[project_id] = EntityIDAllocator(
          self.datastore_batch.session, project_id)
  
      if project_id not in self.scattered_allocators:
        self.scattered_allocators[project_id] = ScatteredAllocator(
          self.datastore_batch.session, project_id)
  
      for id_ in ids:
        counter, space = IdToCounter(id_)
        if space == SEQUENTIAL:
          allocator = self.sequential_allocators[project_id]
        else:
          allocator = self.scattered_allocators[project_id]
  
        yield allocator.set_min_counter(counter)
  
    @gen.coroutine
    def put_entities(self, app, entities):
      """ Updates indexes of existing entities, inserts new entities and
          indexes for them.
  
      Args:
        app: A string containing the application ID.
        entities: List of entities.
      """
      self.logger.debug('Inserting {} entities'.format(len(entities)))
  
      composite_indexes = self.get_indexes(app)
  
      by_group = {}
      for entity in entities:
        group_key = group_for_key(entity.key()).Encode()
        if group_key not in by_group:
          by_group[group_key] = []
        by_group[group_key].append(entity)
  
      for encoded_group_key, entity_list in by_group.iteritems():
        group_key = entity_pb.Reference(encoded_group_key)
  
        txid = self.transaction_manager.create_transaction_id(app, xg=False)
        self.transaction_manager.set_groups(app, txid, [group_key])
  
        # Allow the lock to stick around if there is an issue applying the batch.
        lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
        try:
          yield lock.acquire()
        except entity_lock.LockTimeout:
          raise Timeout('Unable to acquire entity group lock')
  
        try:
          entity_keys = [
            get_entity_key(self.get_table_prefix(entity), entity.key().path())
            for entity in entity_list]
          try:
            current_values = yield self.datastore_batch.batch_get_entity(
              dbconstants.APP_ENTITY_TABLE, entity_keys, APP_ENTITY_SCHEMA)
          except dbconstants.AppScaleDBConnectionError:
            lock.release()
            self.transaction_manager.delete_transaction_id(app, txid)
            raise
  
          batch = []
          entity_changes = []
          for entity in entity_list:
            prefix = self.get_table_prefix(entity)
            entity_key = get_entity_key(prefix, entity.key().path())
  
            current_value = None
            if current_values[entity_key]:
              current_value = entity_pb.EntityProto(
                current_values[entity_key][APP_ENTITY_SCHEMA[0]])
  
            batch.extend(mutations_for_entity(entity, txid, current_value,
                                              composite_indexes))
  
            batch.append({'table': 'group_updates',
                          'key': bytearray(encoded_group_key),
                          'last_update': txid})
  
            entity_changes.append(
              {'key': entity.key(), 'old': current_value, 'new': entity})
  
          if batch_size(batch) > LARGE_BATCH_THRESHOLD:
            try:
              yield self.datastore_batch.large_batch(app, batch, entity_changes,
                                                     txid)
            except BatchNotApplied as error:
              # If the "applied" switch has not been flipped, the lock can be
              # released. The transaction ID is kept so that the groomer can
              # clean up the batch tables.
              lock.release()
              raise dbconstants.AppScaleDBConnectionError(str(error))
          else:
            try:
              yield self.datastore_batch.normal_batch(batch, txid)
            except dbconstants.AppScaleDBConnectionError:
              # Since normal batches are guaranteed to be atomic, the lock can
              # be released.
              lock.release()
              self.transaction_manager.delete_transaction_id(app, txid)
              raise
  
          lock.release()
  
        finally:
          # In case of failure entity group lock should stay acquired
          # as transaction groomer will handle it later.
          # But tornado lock must be released.
          lock.ensure_release_tornado_lock()
  
        self.transaction_manager.delete_transaction_id(app, txid)
  
    @gen.coroutine
    def delete_entities(self, group, txid, keys, composite_indexes=()):
      """ Deletes the entities and the indexes associated with them.
  
      Args:
        group: An entity group Reference object.
        txid: An integer specifying a transaction ID.
        keys: An interable containing entity Reference objects.
        composite_indexes: A list or tuple of CompositeIndex objects.
      """
      entity_keys = []
      for key in keys:
        prefix = self.get_table_prefix(key)
        entity_keys.append(get_entity_key(prefix, key.path()))
  
      # Must fetch the entities to get the keys of indexes before deleting.
      current_values = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, entity_keys, APP_ENTITY_SCHEMA)
  
      for key in entity_keys:
        if not current_values[key]:
          continue
  
        current_value = entity_pb.EntityProto(
          current_values[key][APP_ENTITY_SCHEMA[0]])
        batch = deletions_for_entity(current_value, composite_indexes)
  
        batch.append({'table': 'group_updates',
                      'key': bytearray(group.Encode()),
                      'last_update': txid})
  
        yield self.datastore_batch.normal_batch(batch, txid)
  
    @gen.coroutine
    def dynamic_put(self, app_id, put_request, put_response):
      """ Stores and entity and its indexes in the datastore.
  
      Args:
        app_id: Application ID.
        put_request: Request with entities to store.
        put_response: The response sent back to the app server.
      Raises:
        ZKTransactionException: If we are unable to acquire/release ZooKeeper locks.
      """
      if app_id not in self.scattered_allocators:
        self.scattered_allocators[app_id] = ScatteredAllocator(
          self.datastore_batch.session, app_id)
      allocator = self.scattered_allocators[app_id]
  
      entities = put_request.entity_list()
  
      for entity in entities:
        self.validate_key(entity.key())
  
        for prop in itertools.chain(entity.property_list(),
                                    entity.raw_property_list()):
          if prop.value().has_uservalue():
            uid = md5.new(prop.value().uservalue().email().lower()).digest()
            uid = '1' + ''.join(['%02d' % ord(x) for x in uid])[:20]
            prop.mutable_value().mutable_uservalue().set_obfuscated_gaiaid(uid)
  
        last_path = entity.key().path().element_list()[-1]
        if last_path.id() == 0 and not last_path.has_name():
          allocated_id = yield allocator.next()
          last_path.set_id(allocated_id)
          group = entity.mutable_entity_group()
          root = entity.key().path().element(0)
          group.add_element().CopyFrom(root)
  
      if put_request.has_transaction():
        yield self.datastore_batch.put_entities_tx(
          app_id, put_request.transaction().handle(), entities)
      else:
        yield self.put_entities(app_id, entities)
        self.logger.debug('Updated {} entities'.format(len(entities)))
  
      put_response.key_list().extend([e.key() for e in entities])
  
    def get_root_key_from_entity_key(self, entity_key):
      """ Extract the root key from an entity key. We
          remove any excess children from a string to get to
          the root key.
  
      Args:
        entity_key: A string or Key object representing a row key.
      Returns:
        The root key extracted from the row key.
      Raises:
        TypeError: If the type is not supported.
      """
      if isinstance(entity_key, str):
        tokens = entity_key.split(dbconstants.KIND_SEPARATOR)
        return tokens[0] + dbconstants.KIND_SEPARATOR
      elif isinstance(entity_key, entity_pb.Reference):
        app_id = clean_app_id(entity_key.app())
        path = entity_key.path()
        element_list = path.element_list()
        return self.get_root_key(app_id, entity_key.name_space(), element_list)
      else:
        raise TypeError("Unable to get root key from given type of %s" % \
                        entity_key.__class__)
  
    def get_root_key(self, app_id, ns, ancestor_list):
      """ Gets the root key string from an ancestor listing.
  
      Args:
        app_id: The app ID of the listing.
        ns: The namespace of the entity.
        ancestor_list: The ancestry of a given entity.
      Returns:
        A string representing the root key of an entity.
      """
      prefix = self.get_table_prefix((app_id, ns))
      first_ent = ancestor_list[0]
      if first_ent.has_name():
        key_id = first_ent.name()
      elif first_ent.has_id():
        # Make sure ids are ordered lexigraphically by making sure they
        # are of set size i.e. 2 > 0003 but 0002 < 0003.
        key_id = str(first_ent.id()).zfill(ID_KEY_LENGTH)
      return "{0}{1}{2}:{3}{4}".format(prefix, self._NAMESPACE_SEPARATOR,
        first_ent.type(), key_id, dbconstants.KIND_SEPARATOR)
  
    def is_instance_wrapper(self, obj, expected_type):
      """ A wrapper for isinstance for mocking purposes.
  
      Return whether an object is an instance of a class or of a subclass thereof.
      With a type as second argument, return whether that is the object's type.
  
      Args:
        obj: The object to check.
        expected_type: A instance type we are comparing obj's type to.
      Returns:
        True if obj is of type expected_type, False otherwise.
      """
      return isinstance(obj, expected_type)
  
    def acquire_locks_for_trans(self, entities, txnid):
      """ Acquires locks for entities for one particular entity group.
  
      Args:
        entities: A list of entities (entity_pb.EntityProto or entity_pb.Reference)
                  for which are are getting a lock for.
        txnid: The transaction ID handler.
      Returns:
        A hash mapping root keys to transaction IDs.
      Raises:
        ZKTransactionException: If lock is not obtainable.
        TypeError: If args are of incorrect types.
      """
      # Key tuples are the prefix and the root key for which we're getting locks.
      root_keys = []
      txn_hash = {}
      if not self.is_instance_wrapper(entities, list):
        raise TypeError("Expected a list and got {0}".format(entities.__class__))
      for ent in entities:
        if self.is_instance_wrapper(ent, entity_pb.Reference):
          root_keys.append(self.get_root_key_from_entity_key(ent))
        elif self.is_instance_wrapper(ent, entity_pb.EntityProto):
          root_keys.append(self.get_root_key_from_entity_key(ent.key()))
        else:
          raise TypeError("Excepted either a reference or an EntityProto"
             "got {0}".format(ent.__class__))
  
      if entities == []:
        return {}
  
      if self.is_instance_wrapper(entities[0], entity_pb.Reference):
        app_id = entities[0].app()
      else:
        app_id = entities[0].key().app()
      app_id = clean_app_id(app_id)
      # Remove all duplicate root keys.
      root_keys = list(set(root_keys))
      try:
        for root_key in root_keys:
          txn_hash[root_key] = txnid
          self.zookeeper.acquire_lock(app_id, txnid, root_key)
      except zktransaction.ZKTransactionException as zkte:
        self.logger.warning('Concurrent transaction: {}'.format(txnid))
        for root_key in txn_hash:
          self.zookeeper.notify_failed_transaction(app_id, txn_hash[root_key])
        raise zkte
  
      return txn_hash
  
    def release_locks_for_nontrans(self, app_id, entities, txn_hash):
      """  Releases locks for non-transactional puts.
  
      Args:
        entities: List of entities for which we are releasing locks. Can
                  be either entity_pb.EntityProto or entity_pb.Reference.
        txn_hash: A hash mapping root keys to transaction IDs.
      Raises:
        ZKTransactionException: If we are unable to release locks.
      """
      root_keys = []
      for ent in entities:
        if isinstance(ent, entity_pb.EntityProto):
          ent = ent.key()
        root_keys.append(self.get_root_key_from_entity_key(ent))
  
      # Remove all duplicate root keys
      root_keys = list(set(root_keys))
      for root_key in root_keys:
        txnid = txn_hash[root_key]
        self.zookeeper.release_lock(app_id, txnid)
  
    @gen.coroutine
    def fetch_keys(self, key_list):
      """ Given a list of keys fetch the entities.
  
      Args:
        key_list: A list of keys to fetch.
      Returns:
        A tuple of entities from the datastore and key list.
      """
      row_keys = []
      for key in key_list:
        self.validate_app_id(key.app())
        index_key = str(encode_index_pb(key.path()))
        prefix = self.get_table_prefix(key)
        row_keys.append(self._SEPARATOR.join([prefix, index_key]))
      result = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, row_keys, APP_ENTITY_SCHEMA)
      raise gen.Return((result, row_keys))
  
    @gen.coroutine
    def dynamic_get(self, app_id, get_request, get_response):
      """ Fetch keys from the datastore.
  
      Args:
         app_id: The application ID.
         get_request: Request with list of keys.
         get_response: Response to application server.
      Raises:
        ZKTransactionException: If a lock was unable to get acquired.
      """
      keys = get_request.key_list()
      if len(keys) < 5:
        self.logger.debug('Get:\n{}'.format(get_request))
      else:
        self.logger.debug('Get: {} keys'.format(len(keys)))
  
      if get_request.has_transaction():
        results, row_keys = yield self.fetch_keys(keys)
        fetched_groups = {group_for_key(key).Encode() for key in keys}
        yield self.datastore_batch.record_reads(
          app_id, get_request.transaction().handle(), fetched_groups)
      else:
        results, row_keys = yield self.fetch_keys(keys)
  
      result_count = 0
      for r in row_keys:
        group = get_response.add_entity()
        if r in results and APP_ENTITY_SCHEMA[0] in results[r]:
          result_count += 1
          group.mutable_entity().CopyFrom(
            entity_pb.EntityProto(results[r][APP_ENTITY_SCHEMA[0]]))
      self.logger.debug('Returning {} results'.format(result_count))
  
    @gen.coroutine
    def dynamic_delete(self, app_id, delete_request):
      """ Deletes a set of rows.
  
      Args:
        app_id: The application ID.
        delete_request: Request with a list of keys.
      """
      keys = delete_request.key_list()
      if not keys:
        return
  
      ent_kinds = []
      for key in delete_request.key_list():
        last_path = key.path().element_list()[-1]
        if last_path.type() not in ent_kinds:
          ent_kinds.append(last_path.type())
  
      filtered_indexes = [index for index in self.get_indexes(app_id)
                          if index.definition().entity_type() in ent_kinds]
  
      if delete_request.has_transaction():
        txid = delete_request.transaction().handle()
        yield self.datastore_batch.delete_entities_tx(app_id, txid, keys)
      else:
        by_group = {}
        for key in keys:
          group_key = group_for_key(key).Encode()
          if group_key not in by_group:
            by_group[group_key] = []
          by_group[group_key].append(key)
  
        for encoded_group_key, key_list in by_group.iteritems():
          group_key = entity_pb.Reference(encoded_group_key)
  
          txid = self.transaction_manager.create_transaction_id(app_id, xg=False)
          self.transaction_manager.set_groups(app_id, txid, [group_key])
  
          # Allow the lock to stick around if there is an issue applying the batch.
          lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
          try:
            yield lock.acquire()
          except entity_lock.LockTimeout:
            raise Timeout('Unable to acquire entity group lock')
  
          try:
            yield self.delete_entities(
              group_key,
              txid,
              key_list,
              composite_indexes=filtered_indexes
            )
            lock.release()
          finally:
            # In case of failure entity group lock should stay acquired
            # as transaction groomer will handle it later.
            # But tornado lock must be released.
            lock.ensure_release_tornado_lock()
  
          self.logger.debug('Removed {} entities'.format(len(key_list)))
          self.transaction_manager.delete_transaction_id(app_id, txid)
  
    def generate_filter_info(self, filters):
      """Transform a list of filters into a more usable form.
  
      Args:
        filters: A list of filter PBs.
      Returns:
        A dict mapping property names to lists of (op, value) tuples.
      """
      filter_info = {}
      for filt in filters:
        prop = filt.property(0)
        value = prop.value()
        if prop.name() == '__key__':
          value = reference_property_to_reference(value.referencevalue())
          value = value.path()
        filter_info.setdefault(prop.name(), []).\
          append((filt.op(), encode_index_pb(value)))
      return filter_info
  
    def generate_order_info(self, orders):
      """Transform a list of orders into a more usable form which
         is a tuple of properties and ordering directions.
  
      Args:
        orders: A list of order PBs.
      Returns:
        A list of (property, direction) tuples.
      """
      orders = [(order.property(), order.direction()) for order in orders]
      if orders and orders[-1] == ('__key__', datastore_pb.Query_Order.ASCENDING):
        orders.pop()
      return orders
  
    @gen.coroutine
    def __get_start_key(self, prefix, prop_name, order, last_result, query=None):
      """ Builds the start key for cursor query.
  
      Args:
        prefix: The start key prefix (app id and namespace).
        prop_name: Property name of the filter.
        order: Sort order the query requires.
        last_result: Last result encoded in cursor.
        query: A datastore_pb.Query object.
      Raises:
        AppScaleDBError if unable to retrieve original entity or original entity
          no longer has the requested property.
      """
      e = last_result
      path = str(encode_index_pb(e.key().path()))
      last_result_key = self._SEPARATOR.join([prefix, path])
      if not prop_name and not order:
        raise gen.Return(last_result_key)
      if e.property_list():
        plist = e.property_list()
      else:
        # Fetch the entity from the datastore in order to get the property
        # values.
        ret = yield self.datastore_batch.batch_get_entity(
          dbconstants.APP_ENTITY_TABLE, [last_result_key], APP_ENTITY_SCHEMA)
  
        if APP_ENTITY_SCHEMA[0] not in ret[last_result_key]:
          message = '{} not found in {}'.format(
            last_result_key, dbconstants.APP_ENTITY_TABLE)
          raise dbconstants.AppScaleDBError(message)
  
        ent = entity_pb.EntityProto(ret[last_result_key][APP_ENTITY_SCHEMA[0]])
        plist = ent.property_list()
  
      val = None
  
      # Use the value from the query if possible. This reduces ambiguity when
      # the entity provided by the cursor has multiple values for the property.
      if (query is not None and query.filter_size() == 1 and
          query.filter(0).op() == datastore_pb.Query_Filter.EQUAL):
        query_prop = query.filter(0).property(0)
        val = str(encode_index_pb(query_prop.value()))
  
      if val is None:
        for p in plist:
          if p.name() == prop_name:
            val = str(encode_index_pb(p.value()))
            break
  
      if val is None:
        raise dbconstants.AppScaleDBError('{} not in entity'.format(prop_name))
  
      if order == datastore_pb.Query_Order.DESCENDING:
        val = helper_functions.reverse_lex(val)
      params = [prefix, get_entity_kind(e), prop_name, val, path]
      raise gen.Return(get_index_key_from_params(params))
  
    def is_zigzag_merge_join(self, query, filter_info, order_info):
      """ Checks to see if the current query can be executed as a zigzag
      merge join.
  
      Args:
        query: A datastore_pb.Query.
        filter_info: dict of property names mapping to tuples of filter
          operators and values.
        order_info: tuple with property name and the sort order.
      Returns:
        True if it qualifies as a zigzag merge join, and false otherwise.
      """
      filter_info = self.remove_exists_filters(filter_info)
  
      order_properties = []
      for order in order_info:
        order_properties.append(order[0])
  
      property_names = []
      for property_name in filter_info:
        filt = filter_info[property_name]
        property_names.append(property_name)
        # We only handle equality filters for zigzag merge join queries.
        if filt[0][0] != datastore_pb.Query_Filter.EQUAL:
          return False
  
      if len(filter_info) < 2:
        return False
  
      for order_property_name in order_properties:
        if order_property_name not in property_names:
          return False
  
      return True
  
    @gen.coroutine
    def __fetch_entities_from_row_list(self, rowkeys):
      """ Given a list of keys fetch the entities from the entity table.
  
      Args:
        rowkeys: A list of strings which are keys to the entitiy table.
      Returns:
        A list of entities.
      """
      result = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, rowkeys, APP_ENTITY_SCHEMA)
      entities = []
      for key in rowkeys:
        if key in result and APP_ENTITY_SCHEMA[0] in result[key]:
          entities.append(result[key][APP_ENTITY_SCHEMA[0]])
      raise gen.Return(entities)
  
    def __extract_rowkeys_from_refs(self, refs):
      """ Extract the rowkeys to fetch from a list of references.
  
      Args:
        refs: key/value pairs where the values contain a reference to the
              entitiy table.
      Returns:
        A list of rowkeys.
      """
      if len(refs) == 0:
        return []
      keys = [item.keys()[0] for item in refs]
      rowkeys = []
      for index, ent in enumerate(refs):
        key = keys[index]
        ent = ent[key]['reference']
        # Make sure not to fetch the same entity more than once.
        if ent not in rowkeys:
          rowkeys.append(ent)
      return rowkeys
  
    @gen.coroutine
    def __fetch_entities(self, refs):
      """ Given a list of references, get the entities.
  
      Args:
        refs: key/value pairs where the values contain a reference to
              the entitiy table.
      Returns:
        A list of validated entities.
      """
      rowkeys = self.__extract_rowkeys_from_refs(refs)
      result = yield self.__fetch_entities_from_row_list(rowkeys)
      raise gen.Return(result)
  
    @gen.coroutine
    def __fetch_entities_dict(self, refs):
      """ Given a list of references, return the entities as a dictionary.
  
      Args:
        refs: key/value pairs where the values contain a reference to
              the entitiy table.
      Returns:
        A dictionary of validated entities.
      """
      rowkeys = self.__extract_rowkeys_from_refs(refs)
      result = yield self.__fetch_entities_dict_from_row_list(rowkeys)
      raise gen.Return(result)
  
    @gen.coroutine
    def __fetch_entities_dict_from_row_list(self, rowkeys):
      """ Given a list of rowkeys, return the entities as a dictionary.
  
      Args:
        rowkeys: A list of strings which are keys to the entitiy table.
      Returns:
        A dictionary of validated entities.
      """
      results = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, rowkeys, APP_ENTITY_SCHEMA)
  
      clean_results = {}
      for key in rowkeys:
        if key in results and APP_ENTITY_SCHEMA[0] in results[key]:
          clean_results[key] = results[key][APP_ENTITY_SCHEMA[0]]
  
      raise gen.Return(clean_results)
  
    @gen.coroutine
    def __fetch_and_validate_entity_set(self, index_dict, limit, app_id,
      direction):
      """ Fetch all the valid entities as needed from references.
  
      Args:
        index_dict: A dictionary containing a list of index entries for each
          reference.
        limit: An integer specifying the max number of entities needed.
        app_id: A string, the application identifier.
        direction: The direction of the index.
      Returns:
        A list of valid entities.
      """
      references = index_dict.keys()
      # Prevent duplicate entities across queries with a cursor.
      references.sort()
      offset = 0
      results = []
      to_fetch = limit
      added_padding = False
      while True:
        refs_to_fetch = references[offset:offset + to_fetch]
  
        # If we've exhausted the list of references, we can return.
        if len(refs_to_fetch) == 0:
          raise gen.Return(results[:limit])
  
        entities = yield self.__fetch_entities_dict_from_row_list(refs_to_fetch)
  
        # Prevent duplicate entities across queries with a cursor.
        entity_keys = entities.keys()
        entity_keys.sort()
  
        for reference in entity_keys:
          use_result = False
          indexes_to_check = index_dict[reference]
          for index_info in indexes_to_check:
            index = index_info['index']
            prop_name = index_info['prop_name']
            entry = {index: {'reference': reference}}
            if self.__valid_index_entry(entry, entities, direction, prop_name):
              use_result = True
            else:
              use_result = False
              break
  
          if use_result:
            results.append(entities[reference])
            if len(results) >= limit:
              raise gen.Return(results[:limit])
  
        offset = offset + to_fetch
  
        to_fetch -= len(results)
  
        # Pad the number of references to fetch to increase the likelihood of
        # getting all the valid references that we need.
        if not added_padding:
          to_fetch += dbconstants.MAX_GROUPS_FOR_XG
  
  
  [... the rest is omitted ...]
files.new_path: AppDB/appscale/datastore/datastore_distributed.py
files.old_contents: import array
  import datetime
  import itertools
  import logging
  import md5
  import sys
  import uuid
  
  from tornado import gen
  from tornado.ioloop import IOLoop
  
  from appscale.datastore import dbconstants, helper_functions
  
  from appscale.common.datastore_index import DatastoreIndex, merge_indexes
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from kazoo.client import KazooState
  from appscale.datastore.dbconstants import (
    APP_ENTITY_SCHEMA, BadRequest, ID_KEY_LENGTH, InternalError, MAX_TX_DURATION,
    Timeout
  )
  from appscale.datastore.cassandra_env.cassandra_interface import (
    batch_size, LARGE_BATCH_THRESHOLD)
  from appscale.datastore.cassandra_env.entity_id_allocator import EntityIDAllocator
  from appscale.datastore.cassandra_env.entity_id_allocator import ScatteredAllocator
  from appscale.datastore.cassandra_env.large_batch import BatchNotApplied
  from appscale.datastore.cassandra_env.utils import deletions_for_entity
  from appscale.datastore.cassandra_env.utils import mutations_for_entity
  from appscale.datastore.index_manager import IndexInaccessible
  from appscale.datastore.taskqueue_client import EnqueueError, TaskQueueClient
  from appscale.datastore.utils import _FindIndexToUse
  from appscale.datastore.utils import clean_app_id
  from appscale.datastore.utils import decode_path
  from appscale.datastore.utils import encode_entity_table_key
  from appscale.datastore.utils import encode_index_pb
  from appscale.datastore.utils import encode_path_from_filter
  from appscale.datastore.utils import get_composite_index_keys
  from appscale.datastore.utils import get_entity_key
  from appscale.datastore.utils import get_entity_kind
  from appscale.datastore.utils import get_index_key_from_params
  from appscale.datastore.utils import get_kind_key
  from appscale.datastore.utils import group_for_key
  from appscale.datastore.utils import kind_from_encoded_key
  from appscale.datastore.utils import reference_property_to_reference
  from appscale.datastore.utils import UnprocessedQueryCursor
  from appscale.datastore.range_iterator import RangeExhausted, RangeIterator
  from appscale.datastore.zkappscale import entity_lock
  from appscale.datastore.zkappscale import zktransaction
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.api import datastore_errors
  from google.appengine.api.datastore_distributed import _MAX_ACTIONS_PER_TXN
  from google.appengine.datastore import appscale_stub_util
  from google.appengine.datastore import datastore_pb
  from google.appengine.datastore import datastore_index
  from google.appengine.datastore import entity_pb
  from google.appengine.datastore import sortable_pb_encoder
  from google.appengine.datastore.datastore_stub_util import IdToCounter
  from google.appengine.datastore.datastore_stub_util import SEQUENTIAL
  from google.appengine.runtime import apiproxy_errors
  from google.appengine.ext import db
  from google.appengine.ext.db.metadata import Namespace
  from google.net.proto.ProtocolBuffer import ProtocolBufferDecodeError
  
  logger = logging.getLogger(__name__)
  
  
  class DatastoreDistributed():
    """ AppScale persistent layer for the datastore API. It is the
        replacement for the AppServers to persist their data into
        a distributed datastore instead of a flat file.
    """
    # Max number of results for a query
    _MAXIMUM_RESULTS = 10000
  
    # Maximum amount of filter and orderings allowed within a query
    _MAX_QUERY_COMPONENTS = 63
  
    # For enabling and disabling range inclusivity
    _ENABLE_INCLUSIVITY = True
    _DISABLE_INCLUSIVITY = False
  
    # Delimiter between app names and namespace and the rest of an entity key
    _NAMESPACE_SEPARATOR = dbconstants.KEY_DELIMITER
  
    # Delimiter between parameters in index keys.
    _SEPARATOR = dbconstants.KEY_DELIMITER
  
    # This is the terminating string for range queries
    _TERM_STRING = dbconstants.TERMINATING_STRING
  
    # Smallest possible value that is considered non-null and indexable.
    MIN_INDEX_VALUE = '\x01'
  
    # When assigning the first allocated ID, give this value
    _FIRST_VALID_ALLOCATED_ID = 1
  
    # The key we use to lock for allocating new IDs
    _ALLOCATE_ROOT_KEY = "__allocate__"
  
    # Number of times to retry acquiring a lock for non transactions.
    NON_TRANS_LOCK_RETRY_COUNT = 5
  
    # How long to wait before retrying to grab a lock
    LOCK_RETRY_TIME = .5
  
    # Maximum number of allowed composite indexes any one application can
    # register.
    _MAX_NUM_INDEXES = dbconstants.MAX_NUMBER_OF_COMPOSITE_INDEXES
  
    # The position of the prop name when splitting an index entry by the
    # delimiter.
    PROP_NAME_IN_SINGLE_PROP_INDEX = 3
  
    # The cassandra index column that stores the reference to the entity.
    INDEX_REFERENCE_COLUMN = 'reference'
  
    # The number of entities to fetch at a time when updating indices.
    BATCH_SIZE = 100
  
    def __init__(self, datastore_batch, transaction_manager, zookeeper=None,
                 log_level=logging.INFO, taskqueue_locations=()):
      """
         Constructor.
  
       Args:
         datastore_batch: A reference to the batch datastore interface.
         zookeeper: A reference to the zookeeper interface.
      """
      class_name = self.__class__.__name__
      self.logger = logging.getLogger(class_name)
      self.logger.setLevel(log_level)
  
      assert datastore_batch.valid_data_version_sync()
  
      self.logger.info('Starting {}'.format(class_name))
  
      # datastore accessor used by this class to do datastore operations.
      self.datastore_batch = datastore_batch
  
      # zookeeper instance for accesing ZK functionality.
      self.zookeeper = zookeeper
  
      # Maintain a scattered allocator for each project.
      self.scattered_allocators = {}
  
      # Maintain a sequential allocator for each project.
      self.sequential_allocators = {}
  
      self.taskqueue_client = TaskQueueClient(taskqueue_locations)
      self.transaction_manager = transaction_manager
      self.index_manager = None
      self.zookeeper.handle.add_listener(self._zk_state_listener)
  
    def get_limit(self, query):
      """ Returns the limit that should be used for the given query.
  
      Args:
        query: A datastore_pb.Query.
      Returns:
        An int, the limit to be used when accessing the datastore.
      """
      limit = self._MAXIMUM_RESULTS
      if query.has_count():
        limit = min(query.count(), self._MAXIMUM_RESULTS)
  
      if query.has_limit():
        limit = min(query.limit(), limit)
  
      if query.has_offset():
        limit = limit + min(query.offset(), self._MAXIMUM_RESULTS)
      # We can not scan with 0 or less, hence we set it to one.
      if limit <= 0:
        limit = 1
      return limit
  
    @staticmethod
    def __decode_index_str(value, prop_value):
      """ Takes an encoded string and converts it to a PropertyValue.
  
      Args:
        value: An encoded str.
        prop_value: PropertyValue to fill in.
      """
      value = str(value).replace('\x01\x01', '\x00').replace('\x01\x02', '\x01')
      decoded_value = sortable_pb_encoder.Decoder(
        array.array('B', str(value)))
      prop_value.Merge(decoded_value)
  
    @staticmethod
    def validate_app_id(app_id):
      """ Verify that this is the stub for app_id.
  
      Args:
        app_id: An application ID.
      Raises:
        AppScaleBadArg: If the application id is not set.
      """
      if not app_id:
        raise dbconstants.AppScaleBadArg("Application name must be set")
  
    @staticmethod
    def validate_key(key):
      """ Validate this key by checking to see if it has a name or id.
  
      Args:
        key: entity_pb.Reference
      Raises:
        datastore_errors.BadRequestError: if the key is invalid
        TypeError: if key is not of entity_pb.Reference
      """
  
      if not isinstance(key, entity_pb.Reference):
        raise TypeError("Expected type Reference")
  
      DatastoreDistributed.validate_app_id(key.app())
  
      for elem in key.path().element_list():
        if elem.has_id() and elem.has_name():
          raise datastore_errors.BadRequestError(
              'Each key path element should have id or name but not both: {0}' \
              .format(key))
  
    def get_table_prefix(self, data):
      """ Returns the namespace prefix for a query.
  
      Args:
        data: An Entity, Key or Query PB, or an (app_id, ns) tuple.
      Returns:
        A valid table prefix.
      """
      if isinstance(data, entity_pb.EntityProto):
        app_id = clean_app_id(data.key().app())
        namespace = data.key().name_space()
      elif isinstance(data, tuple):
        app_id = data[0]
        namespace = data[1]
      else:
        app_id = clean_app_id(data.app())
        namespace = data.name_space()
  
      return self._SEPARATOR.join([app_id, namespace])
  
    @staticmethod
    def get_ancestor_key_from_ent_key(ent_key):
      """ Get the key string for the ancestor portion of a composite key.
  
      Args:
        ent_key: A string of the entire path of an entity.
      Returns:
        A str of the path of the ancestor.
      """
      ancestor = ""
      tokens = str(ent_key).split(dbconstants.KIND_SEPARATOR)
      # Strip off the empty placeholder and also do not include the last kind.
      for token in tokens[:-2]:
        ancestor += token + dbconstants.KIND_SEPARATOR
      return ancestor
  
    @staticmethod
    def get_composite_index_key(index, entity, position_list=None,
      filters=None):
      """ Creates a key to the composite index table for a given entity
      for a composite cursor.
  
      Keys are built as such:
        app_id/ns/composite_id/ancestor/valuevaluevalue..../entity_key
      Components explained:
      ns: The namespace of the entity.
      composite_id: The composite ID assigned to this index upon creation.
      ancestor: The root ancestor path (only if the query this index is for
        has an ancestor)
      value(s): The string representation of mulitiple properties.
      entity_key: The entity key (full path) used as a means of having a unique
        identifier. This prevents two entities with the same values from
        colliding.
  
      Args:
        index: A datstore_pb.CompositeIndex.
        entity: A entity_pb.EntityProto.
        position_list: A list of datastore_pb.CompiledCursor_Position items.
          Contains values for property items from a cursor.
        filters: A list of datastore_pb.Query_Filters, used to attain equality
          values not present in position_list.
      Returns:
        A string representing a key to the composite table.
      """
      composite_id = index.id()
      definition = index.definition()
      app_id = clean_app_id(entity.key().app())
      name_space = entity.key().name_space()
      ent_key = encode_index_pb(entity.key().path())
      pre_comp_index_key = "{0}{1}{2}{4}{3}{4}".format(app_id,
        DatastoreDistributed._NAMESPACE_SEPARATOR, name_space, composite_id,
        DatastoreDistributed._SEPARATOR)
      if definition.ancestor() == 1:
        ancestor = DatastoreDistributed.get_ancestor_key_from_ent_key(ent_key)
        pre_comp_index_key += "{0}{1}".format(ancestor,
          DatastoreDistributed._SEPARATOR)
  
      value_dict = {}
      for prop in entity.property_list():
        value_dict[prop.name()] = \
          str(encode_index_pb(prop.value()))
  
      # Position list and filters are used if we're creating a composite
      # key for a cursor.
      if position_list:
        for indexvalue in position_list[0].indexvalue_list():
          value_dict[indexvalue.property()] = \
            str(encode_index_pb(indexvalue.value()))
      if filters:
        for filt in filters:
          if filt.op() == datastore_pb.Query_Filter.EQUAL:
            value_dict[filt.property(0).name()] = \
              str(encode_index_pb(filt.property(0).value()))
  
      index_value = ""
      for prop in definition.property_list():
        name = prop.name()
        value = ''
        if name in value_dict:
          value = value_dict[name]
        elif name == "__key__":
          value = ent_key
        else:
          logger.warning("Given entity {0} is missing a property value {1}.".\
            format(entity, prop.name()))
        if prop.direction() == entity_pb.Index_Property.DESCENDING:
          value = helper_functions.reverse_lex(value)
  
        index_value += str(value) + DatastoreDistributed._SEPARATOR
  
      # We append the ent key to have unique keys if entities happen
      # to share the same index values (and ancestor).
      composite_key = "{0}{1}{2}".format(pre_comp_index_key, index_value,
        ent_key)
      return composite_key
  
    @gen.coroutine
    def insert_composite_indexes(self, entities, composite_indexes):
      """ Creates composite indexes for a set of entities.
  
      Args:
        entities: A list entities.
        composite_indexes: A list of datastore_pb.CompositeIndex.
      """
      if not composite_indexes:
        return
      row_keys = []
      row_values = {}
      # Create default composite index for all entities. Here we take each
      # of the properties in one
      for ent in entities:
        for index_def in composite_indexes:
          # Skip any indexes if the kind does not match.
          kind = get_entity_kind(ent.key())
          if index_def.definition().entity_type() != kind:
            continue
  
          # Make sure the entity contains the required entities for the composite
          # definition.
          prop_name_def_list = [index_prop.name() for index_prop in \
            index_def.definition().property_list()]
          all_prop_names_in_ent = [prop.name() for prop in \
            ent.property_list()]
          has_values = True
          for index_prop_name in prop_name_def_list:
            if index_prop_name not in all_prop_names_in_ent:
              has_values = False
            # Special property name which does not show up in the list but
            # is apart of the key of the entity.
            if index_prop_name == "__key__":
              has_values = True
          if not has_values:
            continue
  
          # Get the composite index key.
          composite_index_keys = get_composite_index_keys(index_def, ent)
          row_keys.extend(composite_index_keys)
  
          # Get the reference value for the composite table.
          entity_key = str(encode_index_pb(ent.key().path()))
          prefix = self.get_table_prefix(ent.key())
          reference = "{0}{1}{2}".format(prefix, self._SEPARATOR,  entity_key)
          for composite_key in composite_index_keys:
            row_values[composite_key] = {'reference': reference}
  
      yield self.datastore_batch.batch_put_entity(
        dbconstants.COMPOSITE_TABLE, row_keys,
        dbconstants.COMPOSITE_SCHEMA, row_values)
  
    @gen.coroutine
    def delete_composite_index_metadata(self, app_id, index):
      """ Deletes a index for the given application identifier.
  
      Args:
        app_id: A string representing the application identifier.
        index: A entity_pb.CompositeIndex object.
      """
      self.logger.info('Deleting composite index:\n{}'.format(index))
      try:
        project_index_manager = self.index_manager.projects[app_id]
      except KeyError:
        raise BadRequest('project_id: {} not found'.format(app_id))
  
      # TODO: Remove actual index entries.
      project_index_manager.delete_index_definition(index.id())
  
    @gen.coroutine
    def create_composite_index(self, app_id, index):
      """ Stores a new index for the given application identifier.
  
      Args:
        app_id: A string representing the application identifier.
        index: A entity_pb.CompositeIndex object.
      Returns:
        A unique number representing the composite index ID.
      """
      new_index = DatastoreIndex.from_pb(index)
  
      # The ID must be a positive number that fits in a signed 64-bit int.
      new_index.id = uuid.uuid1().int >> 65
      merge_indexes(self.zookeeper.handle, app_id, [new_index])
      raise gen.Return(new_index.id)
  
    @gen.coroutine
    def update_composite_index(self, app_id, index):
      """ Updates an index for a given app ID.
  
      Args:
        app_id: A string containing the app ID.
        index: An entity_pb.CompositeIndex object.
      """
      self.logger.info('Updating index: {}'.format(index))
      entries_updated = 0
      entity_type = index.definition().entity_type()
  
      # TODO: Adjust prefix based on ancestor.
      prefix = '{app}{delimiter}{entity_type}{kind_separator}'.format(
        app=app_id,
        delimiter=self._SEPARATOR * 2,
        entity_type=entity_type,
        kind_separator=dbconstants.KIND_SEPARATOR,
      )
      start_row = prefix
      end_row = prefix + self._TERM_STRING
      start_inclusive = True
  
      while True:
        # Fetch references from the kind table since entity keys can have a
        # parent prefix.
        references = yield self.datastore_batch.range_query(
          table_name=dbconstants.APP_KIND_TABLE,
          column_names=dbconstants.APP_KIND_SCHEMA,
          start_key=start_row,
          end_key=end_row,
          limit=self.BATCH_SIZE,
          offset=0,
          start_inclusive=start_inclusive,
        )
  
        pb_entities = yield self.__fetch_entities(references)
        entities = [entity_pb.EntityProto(entity) for entity in pb_entities]
  
        yield self.insert_composite_indexes(entities, [index])
        entries_updated += len(entities)
  
        # If we fetched fewer references than we asked for, we're done.
        if len(references) < self.BATCH_SIZE:
          break
  
        start_row = references[-1].keys()[0]
        start_inclusive = self._DISABLE_INCLUSIVITY
  
      self.logger.info('Updated {} index entries.'.format(entries_updated))
  
    @gen.coroutine
    def allocate_size(self, project, size):
      """ Allocates a block of IDs for a project.
  
      Args:
        project: A string specifying the project ID.
        size: An integer specifying the number of IDs to reserve.
      Returns:
        A tuple of integers specifying the start and end ID.
      """
      if project not in self.sequential_allocators:
        self.sequential_allocators[project] = EntityIDAllocator(
          self.datastore_batch.session, project)
  
      allocator = self.sequential_allocators[project]
      start_id, end_id = yield allocator.allocate_size(size)
      raise gen.Return((start_id, end_id))
  
    @gen.coroutine
    def allocate_max(self, project, max_id):
      """ Reserves all IDs up to the one given.
  
      Args:
        project: A string specifying the project ID.
        max_id: An integer specifying the maximum ID to allocated.
      Returns:
        A tuple of integers specifying the start and end ID.
      """
      if project not in self.sequential_allocators:
        self.sequential_allocators[project] = EntityIDAllocator(
          self.datastore_batch.session, project)
  
      allocator = self.sequential_allocators[project]
      start_id, end_id = yield allocator.allocate_max(max_id)
      raise gen.Return((start_id, end_id))
  
    @gen.coroutine
    def reserve_ids(self, project_id, ids):
      """ Ensures the given IDs are not re-allocated.
  
      Args:
        project_id: A string specifying the project ID.
        ids: An iterable of integers specifying entity IDs.
      """
      if project_id not in self.sequential_allocators:
        self.sequential_allocators[project_id] = EntityIDAllocator(
          self.datastore_batch.session, project_id)
  
      if project_id not in self.scattered_allocators:
        self.scattered_allocators[project_id] = ScatteredAllocator(
          self.datastore_batch.session, project_id)
  
      for id_ in ids:
        counter, space = IdToCounter(id_)
        if space == SEQUENTIAL:
          allocator = self.sequential_allocators[project_id]
        else:
          allocator = self.scattered_allocators[project_id]
  
        yield allocator.set_min_counter(counter)
  
    @gen.coroutine
    def put_entities(self, app, entities):
      """ Updates indexes of existing entities, inserts new entities and
          indexes for them.
  
      Args:
        app: A string containing the application ID.
        entities: List of entities.
      """
      self.logger.debug('Inserting {} entities'.format(len(entities)))
  
      composite_indexes = self.get_indexes(app)
  
      by_group = {}
      for entity in entities:
        group_key = group_for_key(entity.key()).Encode()
        if group_key not in by_group:
          by_group[group_key] = []
        by_group[group_key].append(entity)
  
      for encoded_group_key, entity_list in by_group.iteritems():
        group_key = entity_pb.Reference(encoded_group_key)
  
        txid = self.transaction_manager.create_transaction_id(app, xg=False)
        self.transaction_manager.set_groups(app, txid, [group_key])
  
        # Allow the lock to stick around if there is an issue applying the batch.
        lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
        try:
          yield lock.acquire()
        except entity_lock.LockTimeout:
          raise Timeout('Unable to acquire entity group lock')
  
        try:
          entity_keys = [
            get_entity_key(self.get_table_prefix(entity), entity.key().path())
            for entity in entity_list]
          try:
            current_values = yield self.datastore_batch.batch_get_entity(
              dbconstants.APP_ENTITY_TABLE, entity_keys, APP_ENTITY_SCHEMA)
          except dbconstants.AppScaleDBConnectionError:
            lock.release()
            self.transaction_manager.delete_transaction_id(app, txid)
            raise
  
          batch = []
          entity_changes = []
          for entity in entity_list:
            prefix = self.get_table_prefix(entity)
            entity_key = get_entity_key(prefix, entity.key().path())
  
            current_value = None
            if current_values[entity_key]:
              current_value = entity_pb.EntityProto(
                current_values[entity_key][APP_ENTITY_SCHEMA[0]])
  
            batch.extend(mutations_for_entity(entity, txid, current_value,
                                              composite_indexes))
  
            batch.append({'table': 'group_updates',
                          'key': bytearray(encoded_group_key),
                          'last_update': txid})
  
            entity_changes.append(
              {'key': entity.key(), 'old': current_value, 'new': entity})
  
          if batch_size(batch) > LARGE_BATCH_THRESHOLD:
            try:
              yield self.datastore_batch.large_batch(app, batch, entity_changes,
                                                     txid)
            except BatchNotApplied as error:
              # If the "applied" switch has not been flipped, the lock can be
              # released. The transaction ID is kept so that the groomer can
              # clean up the batch tables.
              lock.release()
              raise dbconstants.AppScaleDBConnectionError(str(error))
          else:
            try:
              yield self.datastore_batch.normal_batch(batch, txid)
            except dbconstants.AppScaleDBConnectionError:
              # Since normal batches are guaranteed to be atomic, the lock can
              # be released.
              lock.release()
              self.transaction_manager.delete_transaction_id(app, txid)
              raise
  
          lock.release()
  
        finally:
          # In case of failure entity group lock should stay acquired
          # as transaction groomer will handle it later.
          # But tornado lock must be released.
          lock.ensure_release_tornado_lock()
  
        self.transaction_manager.delete_transaction_id(app, txid)
  
    @gen.coroutine
    def delete_entities(self, group, txid, keys, composite_indexes=()):
      """ Deletes the entities and the indexes associated with them.
  
      Args:
        group: An entity group Reference object.
        txid: An integer specifying a transaction ID.
        keys: An interable containing entity Reference objects.
        composite_indexes: A list or tuple of CompositeIndex objects.
      """
      entity_keys = []
      for key in keys:
        prefix = self.get_table_prefix(key)
        entity_keys.append(get_entity_key(prefix, key.path()))
  
      # Must fetch the entities to get the keys of indexes before deleting.
      current_values = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, entity_keys, APP_ENTITY_SCHEMA)
  
      for key in entity_keys:
        if not current_values[key]:
          continue
  
        current_value = entity_pb.EntityProto(
          current_values[key][APP_ENTITY_SCHEMA[0]])
        batch = deletions_for_entity(current_value, composite_indexes)
  
        batch.append({'table': 'group_updates',
                      'key': bytearray(group.Encode()),
                      'last_update': txid})
  
        yield self.datastore_batch.normal_batch(batch, txid)
  
    @gen.coroutine
    def dynamic_put(self, app_id, put_request, put_response):
      """ Stores and entity and its indexes in the datastore.
  
      Args:
        app_id: Application ID.
        put_request: Request with entities to store.
        put_response: The response sent back to the app server.
      Raises:
        ZKTransactionException: If we are unable to acquire/release ZooKeeper locks.
      """
      if app_id not in self.scattered_allocators:
        self.scattered_allocators[app_id] = ScatteredAllocator(
          self.datastore_batch.session, app_id)
      allocator = self.scattered_allocators[app_id]
  
      entities = put_request.entity_list()
  
      for entity in entities:
        self.validate_key(entity.key())
  
        for prop in itertools.chain(entity.property_list(),
                                    entity.raw_property_list()):
          if prop.value().has_uservalue():
            uid = md5.new(prop.value().uservalue().email().lower()).digest()
            uid = '1' + ''.join(['%02d' % ord(x) for x in uid])[:20]
            prop.mutable_value().mutable_uservalue().set_obfuscated_gaiaid(uid)
  
        last_path = entity.key().path().element_list()[-1]
        if last_path.id() == 0 and not last_path.has_name():
          allocated_id = yield allocator.next()
          last_path.set_id(allocated_id)
          group = entity.mutable_entity_group()
          root = entity.key().path().element(0)
          group.add_element().CopyFrom(root)
  
      if put_request.has_transaction():
        yield self.datastore_batch.put_entities_tx(
          app_id, put_request.transaction().handle(), entities)
      else:
        yield self.put_entities(app_id, entities)
        self.logger.debug('Updated {} entities'.format(len(entities)))
  
      put_response.key_list().extend([e.key() for e in entities])
  
    def get_root_key_from_entity_key(self, entity_key):
      """ Extract the root key from an entity key. We
          remove any excess children from a string to get to
          the root key.
  
      Args:
        entity_key: A string or Key object representing a row key.
      Returns:
        The root key extracted from the row key.
      Raises:
        TypeError: If the type is not supported.
      """
      if isinstance(entity_key, str):
        tokens = entity_key.split(dbconstants.KIND_SEPARATOR)
        return tokens[0] + dbconstants.KIND_SEPARATOR
      elif isinstance(entity_key, entity_pb.Reference):
        app_id = clean_app_id(entity_key.app())
        path = entity_key.path()
        element_list = path.element_list()
        return self.get_root_key(app_id, entity_key.name_space(), element_list)
      else:
        raise TypeError("Unable to get root key from given type of %s" % \
                        entity_key.__class__)
  
    def get_root_key(self, app_id, ns, ancestor_list):
      """ Gets the root key string from an ancestor listing.
  
      Args:
        app_id: The app ID of the listing.
        ns: The namespace of the entity.
        ancestor_list: The ancestry of a given entity.
      Returns:
        A string representing the root key of an entity.
      """
      prefix = self.get_table_prefix((app_id, ns))
      first_ent = ancestor_list[0]
      if first_ent.has_name():
        key_id = first_ent.name()
      elif first_ent.has_id():
        # Make sure ids are ordered lexigraphically by making sure they
        # are of set size i.e. 2 > 0003 but 0002 < 0003.
        key_id = str(first_ent.id()).zfill(ID_KEY_LENGTH)
      return "{0}{1}{2}:{3}{4}".format(prefix, self._NAMESPACE_SEPARATOR,
        first_ent.type(), key_id, dbconstants.KIND_SEPARATOR)
  
    def is_instance_wrapper(self, obj, expected_type):
      """ A wrapper for isinstance for mocking purposes.
  
      Return whether an object is an instance of a class or of a subclass thereof.
      With a type as second argument, return whether that is the object's type.
  
      Args:
        obj: The object to check.
        expected_type: A instance type we are comparing obj's type to.
      Returns:
        True if obj is of type expected_type, False otherwise.
      """
      return isinstance(obj, expected_type)
  
    def acquire_locks_for_trans(self, entities, txnid):
      """ Acquires locks for entities for one particular entity group.
  
      Args:
        entities: A list of entities (entity_pb.EntityProto or entity_pb.Reference)
                  for which are are getting a lock for.
        txnid: The transaction ID handler.
      Returns:
        A hash mapping root keys to transaction IDs.
      Raises:
        ZKTransactionException: If lock is not obtainable.
        TypeError: If args are of incorrect types.
      """
      # Key tuples are the prefix and the root key for which we're getting locks.
      root_keys = []
      txn_hash = {}
      if not self.is_instance_wrapper(entities, list):
        raise TypeError("Expected a list and got {0}".format(entities.__class__))
      for ent in entities:
        if self.is_instance_wrapper(ent, entity_pb.Reference):
          root_keys.append(self.get_root_key_from_entity_key(ent))
        elif self.is_instance_wrapper(ent, entity_pb.EntityProto):
          root_keys.append(self.get_root_key_from_entity_key(ent.key()))
        else:
          raise TypeError("Excepted either a reference or an EntityProto"
             "got {0}".format(ent.__class__))
  
      if entities == []:
        return {}
  
      if self.is_instance_wrapper(entities[0], entity_pb.Reference):
        app_id = entities[0].app()
      else:
        app_id = entities[0].key().app()
      app_id = clean_app_id(app_id)
      # Remove all duplicate root keys.
      root_keys = list(set(root_keys))
      try:
        for root_key in root_keys:
          txn_hash[root_key] = txnid
          self.zookeeper.acquire_lock(app_id, txnid, root_key)
      except zktransaction.ZKTransactionException as zkte:
        self.logger.warning('Concurrent transaction: {}'.format(txnid))
        for root_key in txn_hash:
          self.zookeeper.notify_failed_transaction(app_id, txn_hash[root_key])
        raise zkte
  
      return txn_hash
  
    def release_locks_for_nontrans(self, app_id, entities, txn_hash):
      """  Releases locks for non-transactional puts.
  
      Args:
        entities: List of entities for which we are releasing locks. Can
                  be either entity_pb.EntityProto or entity_pb.Reference.
        txn_hash: A hash mapping root keys to transaction IDs.
      Raises:
        ZKTransactionException: If we are unable to release locks.
      """
      root_keys = []
      for ent in entities:
        if isinstance(ent, entity_pb.EntityProto):
          ent = ent.key()
        root_keys.append(self.get_root_key_from_entity_key(ent))
  
      # Remove all duplicate root keys
      root_keys = list(set(root_keys))
      for root_key in root_keys:
        txnid = txn_hash[root_key]
        self.zookeeper.release_lock(app_id, txnid)
  
    @gen.coroutine
    def fetch_keys(self, key_list):
      """ Given a list of keys fetch the entities.
  
      Args:
        key_list: A list of keys to fetch.
      Returns:
        A tuple of entities from the datastore and key list.
      """
      row_keys = []
      for key in key_list:
        self.validate_app_id(key.app())
        index_key = str(encode_index_pb(key.path()))
        prefix = self.get_table_prefix(key)
        row_keys.append(self._SEPARATOR.join([prefix, index_key]))
      result = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, row_keys, APP_ENTITY_SCHEMA)
      raise gen.Return((result, row_keys))
  
    @gen.coroutine
    def dynamic_get(self, app_id, get_request, get_response):
      """ Fetch keys from the datastore.
  
      Args:
         app_id: The application ID.
         get_request: Request with list of keys.
         get_response: Response to application server.
      Raises:
        ZKTransactionException: If a lock was unable to get acquired.
      """
      keys = get_request.key_list()
      if len(keys) < 5:
        self.logger.debug('Get:\n{}'.format(get_request))
      else:
        self.logger.debug('Get: {} keys'.format(len(keys)))
  
      if get_request.has_transaction():
        results, row_keys = yield self.fetch_keys(keys)
        fetched_groups = {group_for_key(key).Encode() for key in keys}
        yield self.datastore_batch.record_reads(
          app_id, get_request.transaction().handle(), fetched_groups)
      else:
        results, row_keys = yield self.fetch_keys(keys)
  
      result_count = 0
      for r in row_keys:
        group = get_response.add_entity()
        if r in results and APP_ENTITY_SCHEMA[0] in results[r]:
          result_count += 1
          group.mutable_entity().CopyFrom(
            entity_pb.EntityProto(results[r][APP_ENTITY_SCHEMA[0]]))
      self.logger.debug('Returning {} results'.format(result_count))
  
    @gen.coroutine
    def dynamic_delete(self, app_id, delete_request):
      """ Deletes a set of rows.
  
      Args:
        app_id: The application ID.
        delete_request: Request with a list of keys.
      """
      keys = delete_request.key_list()
      if not keys:
        return
  
      ent_kinds = []
      for key in delete_request.key_list():
        last_path = key.path().element_list()[-1]
        if last_path.type() not in ent_kinds:
          ent_kinds.append(last_path.type())
  
      filtered_indexes = [index for index in self.get_indexes(app_id)
                          if index.definition().entity_type() in ent_kinds]
  
      if delete_request.has_transaction():
        txid = delete_request.transaction().handle()
        yield self.datastore_batch.delete_entities_tx(app_id, txid, keys)
      else:
        by_group = {}
        for key in keys:
          group_key = group_for_key(key).Encode()
          if group_key not in by_group:
            by_group[group_key] = []
          by_group[group_key].append(key)
  
        for encoded_group_key, key_list in by_group.iteritems():
          group_key = entity_pb.Reference(encoded_group_key)
  
          txid = self.transaction_manager.create_transaction_id(app_id, xg=False)
          self.transaction_manager.set_groups(app_id, txid, [group_key])
  
          # Allow the lock to stick around if there is an issue applying the batch.
          lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
          try:
            yield lock.acquire()
          except entity_lock.LockTimeout:
            raise Timeout('Unable to acquire entity group lock')
  
          try:
            yield self.delete_entities(
              group_key,
              txid,
              key_list,
              composite_indexes=filtered_indexes
            )
            lock.release()
          finally:
            # In case of failure entity group lock should stay acquired
            # as transaction groomer will handle it later.
            # But tornado lock must be released.
            lock.ensure_release_tornado_lock()
  
          self.logger.debug('Removed {} entities'.format(len(key_list)))
          self.transaction_manager.delete_transaction_id(app_id, txid)
  
    def generate_filter_info(self, filters):
      """Transform a list of filters into a more usable form.
  
      Args:
        filters: A list of filter PBs.
      Returns:
        A dict mapping property names to lists of (op, value) tuples.
      """
      filter_info = {}
      for filt in filters:
        prop = filt.property(0)
        value = prop.value()
        if prop.name() == '__key__':
          value = reference_property_to_reference(value.referencevalue())
          value = value.path()
        filter_info.setdefault(prop.name(), []).\
          append((filt.op(), encode_index_pb(value)))
      return filter_info
  
    def generate_order_info(self, orders):
      """Transform a list of orders into a more usable form which
         is a tuple of properties and ordering directions.
  
      Args:
        orders: A list of order PBs.
      Returns:
        A list of (property, direction) tuples.
      """
      orders = [(order.property(), order.direction()) for order in orders]
      if orders and orders[-1] == ('__key__', datastore_pb.Query_Order.ASCENDING):
        orders.pop()
      return orders
  
    @gen.coroutine
    def __get_start_key(self, prefix, prop_name, order, last_result, query=None):
      """ Builds the start key for cursor query.
  
      Args:
        prefix: The start key prefix (app id and namespace).
        prop_name: Property name of the filter.
        order: Sort order the query requires.
        last_result: Last result encoded in cursor.
        query: A datastore_pb.Query object.
      Raises:
        AppScaleDBError if unable to retrieve original entity or original entity
          no longer has the requested property.
      """
      e = last_result
      path = str(encode_index_pb(e.key().path()))
      last_result_key = self._SEPARATOR.join([prefix, path])
      if not prop_name and not order:
        raise gen.Return(last_result_key)
      if e.property_list():
        plist = e.property_list()
      else:
        # Fetch the entity from the datastore in order to get the property
        # values.
        ret = yield self.datastore_batch.batch_get_entity(
          dbconstants.APP_ENTITY_TABLE, [last_result_key], APP_ENTITY_SCHEMA)
  
        if APP_ENTITY_SCHEMA[0] not in ret[last_result_key]:
          message = '{} not found in {}'.format(
            last_result_key, dbconstants.APP_ENTITY_TABLE)
          raise dbconstants.AppScaleDBError(message)
  
        ent = entity_pb.EntityProto(ret[last_result_key][APP_ENTITY_SCHEMA[0]])
        plist = ent.property_list()
  
      val = None
  
      # Use the value from the query if possible. This reduces ambiguity when
      # the entity provided by the cursor has multiple values for the property.
      if (query is not None and query.filter_size() == 1 and
          query.filter(0).op() == datastore_pb.Query_Filter.EQUAL):
        query_prop = query.filter(0).property(0)
        val = str(encode_index_pb(query_prop.value()))
  
      if val is None:
        for p in plist:
          if p.name() == prop_name:
            val = str(encode_index_pb(p.value()))
            break
  
      if val is None:
        raise dbconstants.AppScaleDBError('{} not in entity'.format(prop_name))
  
      if order == datastore_pb.Query_Order.DESCENDING:
        val = helper_functions.reverse_lex(val)
      params = [prefix, get_entity_kind(e), prop_name, val, path]
      raise gen.Return(get_index_key_from_params(params))
  
    def is_zigzag_merge_join(self, query, filter_info, order_info):
      """ Checks to see if the current query can be executed as a zigzag
      merge join.
  
      Args:
        query: A datastore_pb.Query.
        filter_info: dict of property names mapping to tuples of filter
          operators and values.
        order_info: tuple with property name and the sort order.
      Returns:
        True if it qualifies as a zigzag merge join, and false otherwise.
      """
      filter_info = self.remove_exists_filters(filter_info)
  
      order_properties = []
      for order in order_info:
        order_properties.append(order[0])
  
      property_names = []
      for property_name in filter_info:
        filt = filter_info[property_name]
        property_names.append(property_name)
        # We only handle equality filters for zigzag merge join queries.
        if filt[0][0] != datastore_pb.Query_Filter.EQUAL:
          return False
  
      if len(filter_info) < 2:
        return False
  
      for order_property_name in order_properties:
        if order_property_name not in property_names:
          return False
  
      return True
  
    @gen.coroutine
    def __fetch_entities_from_row_list(self, rowkeys):
      """ Given a list of keys fetch the entities from the entity table.
  
      Args:
        rowkeys: A list of strings which are keys to the entitiy table.
      Returns:
        A list of entities.
      """
      result = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, rowkeys, APP_ENTITY_SCHEMA)
      entities = []
      for key in rowkeys:
        if key in result and APP_ENTITY_SCHEMA[0] in result[key]:
          entities.append(result[key][APP_ENTITY_SCHEMA[0]])
      raise gen.Return(entities)
  
    def __extract_rowkeys_from_refs(self, refs):
      """ Extract the rowkeys to fetch from a list of references.
  
      Args:
        refs: key/value pairs where the values contain a reference to the
              entitiy table.
      Returns:
        A list of rowkeys.
      """
      if len(refs) == 0:
        return []
      keys = [item.keys()[0] for item in refs]
      rowkeys = []
      for index, ent in enumerate(refs):
        key = keys[index]
        ent = ent[key]['reference']
        # Make sure not to fetch the same entity more than once.
        if ent not in rowkeys:
          rowkeys.append(ent)
      return rowkeys
  
    @gen.coroutine
    def __fetch_entities(self, refs):
      """ Given a list of references, get the entities.
  
      Args:
        refs: key/value pairs where the values contain a reference to
              the entitiy table.
      Returns:
        A list of validated entities.
      """
      rowkeys = self.__extract_rowkeys_from_refs(refs)
      result = yield self.__fetch_entities_from_row_list(rowkeys)
      raise gen.Return(result)
  
    @gen.coroutine
    def __fetch_entities_dict(self, refs):
      """ Given a list of references, return the entities as a dictionary.
  
      Args:
        refs: key/value pairs where the values contain a reference to
              the entitiy table.
      Returns:
        A dictionary of validated entities.
      """
      rowkeys = self.__extract_rowkeys_from_refs(refs)
      result = yield self.__fetch_entities_dict_from_row_list(rowkeys)
      raise gen.Return(result)
  
    @gen.coroutine
    def __fetch_entities_dict_from_row_list(self, rowkeys):
      """ Given a list of rowkeys, return the entities as a dictionary.
  
      Args:
        rowkeys: A list of strings which are keys to the entitiy table.
      Returns:
        A dictionary of validated entities.
      """
      results = yield self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, rowkeys, APP_ENTITY_SCHEMA)
  
      clean_results = {}
      for key in rowkeys:
        if key in results and APP_ENTITY_SCHEMA[0] in results[key]:
          clean_results[key] = results[key][APP_ENTITY_SCHEMA[0]]
  
      raise gen.Return(clean_results)
  
    @gen.coroutine
    def __fetch_and_validate_entity_set(self, index_dict, limit, app_id,
      direction):
      """ Fetch all the valid entities as needed from references.
  
      Args:
        index_dict: A dictionary containing a list of index entries for each
          reference.
        limit: An integer specifying the max number of entities needed.
        app_id: A string, the application identifier.
        direction: The direction of the index.
      Returns:
        A list of valid entities.
      """
      references = index_dict.keys()
      # Prevent duplicate entities across queries with a cursor.
      references.sort()
      offset = 0
      results = []
      to_fetch = limit
      added_padding = False
      while True:
        refs_to_fetch = references[offset:offset + to_fetch]
  
        # If we've exhausted the list of references, we can return.
        if len(refs_to_fetch) == 0:
          raise gen.Return(results[:limit])
  
        entities = yield self.__fetch_entities_dict_from_row_list(refs_to_fetch)
  
        # Prevent duplicate entities across queries with a cursor.
        entity_keys = entities.keys()
        entity_keys.sort()
  
        for reference in entity_keys:
          use_result = False
          indexes_to_check = index_dict[reference]
          for index_info in indexes_to_check:
            index = index_info['index']
            prop_name = index_info['prop_name']
            entry = {index: {'reference': reference}}
            if self.__valid_index_entry(entry, entities, direction, prop_name):
              use_result = True
            else:
              use_result = False
              break
  
          if use_result:
            results.append(entities[reference])
            if len(results) >= limit:
              raise gen.Return(results[:limit])
  
        offset = offset + to_fetch
  
        to_fetch -= len(results)
  
        # Pad the number of references to fetch to increase the likelihood of
        # getting all the valid references that we need.
        if not added_padding:
          to_fetch += dbconstants.MAX_GROUPS_FOR_XG
  
  
  [... the rest is omitted ...]
files.old_path: AppDB/appscale/datastore/datastore_distributed.py
files.new_contents: # See LICENSE file
  #
  """
  This web service interfaces with the datastore. It takes protocol buffer
  requests from AppServers and responds according to the type of request its
  given (Put, Get, Delete, Query, etc).
  """
  import argparse
  import json
  import logging
  import os
  import sys
  import time
  import tornado.httpserver
  import tornado.web
  
  from appscale.common import appscale_info
  from appscale.common.appscale_info import get_load_balancer_ips
  from appscale.common.async_retrying import retry_data_watch_coroutine
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from kazoo.client import KazooState
  from kazoo.exceptions import NodeExistsError
  from tornado import gen
  from tornado.httpclient import AsyncHTTPClient
  from tornado.ioloop import IOLoop
  from tornado.options import options
  from .. import dbconstants
  from ..appscale_datastore_batch import DatastoreFactory
  from ..datastore_distributed import DatastoreDistributed
  from ..index_manager import IndexManager
  from ..utils import (clean_app_id,
                       logger,
                       UnprocessedQueryResult)
  from ..zkappscale import zktransaction
  from ..zkappscale.transaction_manager import TransactionManager
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.api import api_base_pb
  from google.appengine.api.taskqueue import taskqueue_service_pb
  from google.appengine.datastore import datastore_pb
  from google.appengine.datastore import datastore_v4_pb
  from google.appengine.datastore import entity_pb
  from google.appengine.ext.remote_api import remote_api_pb
  from google.net.proto.ProtocolBuffer import ProtocolBufferDecodeError
  
  # Global for accessing the datastore. An instance of DatastoreDistributed.
  datastore_access = None
  
  # A record of active datastore servers.
  datastore_servers = set()
  
  # The ZooKeeper path where this server registers its availability.
  server_node = None
  
  # ZooKeeper global variable for locking
  zookeeper = None
  
  # Determines whether or not to allow datastore writes. Note: After enabling,
  # datastore processes must be restarted and the groomer must be stopped.
  READ_ONLY = False
  
  # Global stats.
  STATS = {}
  
  # The ZooKeeper path where a list of active datastore servers is stored.
  DATASTORE_SERVERS_NODE = '/appscale/datastore/servers'
  
  
  class ClearHandler(tornado.web.RequestHandler):
    """ Defines what to do when the webserver receives a /clear HTTP request. """
    def set_default_headers(self):
      """ Instructs clients to close the connection after each response. """
      self.set_header('Connection', 'close')
  
    @tornado.web.asynchronous
    def post(self):
      """ Handles POST requests for clearing datastore server stats. """
      global STATS
      STATS = {}
      self.write({"message": "Statistics for this server cleared."})
      self.finish()
  
  
  class ReadOnlyHandler(tornado.web.RequestHandler):
    """ Handles requests to check or set read-only mode. """
    def set_default_headers(self):
      """ Instructs clients to close the connection after each response. """
      self.set_header('Connection', 'close')
  
    @tornado.web.asynchronous
    def post(self):
      """ Handle requests to turn read-only mode on or off. """
      global READ_ONLY
  
      payload = self.request.body
      data = json.loads(payload)
      if 'readOnly' not in data:
        self.set_status(dbconstants.HTTP_BAD_REQUEST)
  
      if data['readOnly']:
        READ_ONLY = True
        message = 'Write operations now disabled.'
      else:
        READ_ONLY = False
        message = 'Write operations now enabled.'
  
      logger.info(message)
      self.write({'message': message})
      self.finish()
  
  
  class ReserveKeysHandler(tornado.web.RequestHandler):
    """ Handles v4 AllocateIds requests from other servers. """
    @gen.coroutine
    def post(self):
      """ Prevents the provided IDs from being re-allocated. """
      project_id = self.request.headers['appdata']
      request = datastore_v4_pb.AllocateIdsRequest(self.request.body)
      ids = [key.path_element_list()[-1].id() for key in request.reserve_list()]
      yield datastore_access.reserve_ids(project_id, ids)
  
  
  class MainHandler(tornado.web.RequestHandler):
    """
    Defines what to do when the webserver receives different types of 
    HTTP requests.
    """
    def set_default_headers(self):
      """ Instructs clients to close the connection after each response. """
      self.set_header('Connection', 'close')
  
    def unknown_request(self, app_id, http_request_data, pb_type):
      """ Function which handles unknown protocol buffers.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer
      Raises:
        Raises exception.
      """
      raise NotImplementedError("Unknown request of operation {0}" \
        .format(pb_type))
  
    @gen.coroutine
    def post(self):
      """ Function which handles POST requests. Data of the request is
          the request from the AppServer in an encoded protocol buffer
          format.
      """
      request = self.request
      http_request_data = request.body
      pb_type = request.headers['protocolbuffertype']
      app_data = request.headers['appdata']
      app_data = app_data.split(':')
  
      if len(app_data) == 4:
        app_id, user_email, nick_name, auth_domain = app_data
        os.environ['AUTH_DOMAIN'] = auth_domain
        os.environ['USER_EMAIL'] = user_email
        os.environ['USER_NICKNAME'] = nick_name
        os.environ['APPLICATION_ID'] = app_id
      elif len(app_data) == 1:
        app_id = app_data[0]
        os.environ['APPLICATION_ID'] = app_id
      else:
        return
  
      # If the application identifier has the HRD string prepened, remove it.
      app_id = clean_app_id(app_id)
  
      if pb_type == "Request":
        yield self.remote_request(app_id, http_request_data,
                                  service_id=request.headers.get('Module'),
                                  version_id=request.headers.get('Version'))
      else:
        self.unknown_request(app_id, http_request_data, pb_type)
  
    @tornado.web.asynchronous
    def get(self):
      """ Handles get request for the web server. Returns that it is currently
          up in json.
      """
      self.write(json.dumps(STATS))
      self.finish()
  
    @gen.coroutine
    def remote_request(self, app_id, http_request_data, service_id, version_id):
      """ Receives a remote request to which it should give the correct 
          response. The http_request_data holds an encoded protocol buffer
          of a certain type. Each type has a particular response type. 
      
      Args:
        app_id: The application ID that is sending this request.
        http_request_data: Encoded protocol buffer.
        service_id: A string specifying the client's service ID.
        version_id: A string specifying the client's version ID.
      """
      apirequest = remote_api_pb.Request()
      apirequest.ParseFromString(http_request_data)
      apiresponse = remote_api_pb.Response()
      response = None
      if not apirequest.has_method():
        apirequest.set_method("NOT_FOUND")
      if not apirequest.has_request():
        apirequest.set_method("NOT_FOUND")
        apirequest.clear_request()
      method = apirequest.method()
      http_request_data = apirequest.request()
      start = time.time()
  
      request_log = method
      if apirequest.has_request_id():
        request_log += ': {}'.format(apirequest.request_id())
      logger.debug(request_log)
  
      if method == "Put":
        response, errcode, errdetail = yield self.put_request(
          app_id, http_request_data)
      elif method == "Get":
        response, errcode, errdetail = yield self.get_request(
          app_id, http_request_data)
      elif method == "Delete":
        response, errcode, errdetail = yield self.delete_request(
          app_id, http_request_data)
      elif method == "RunQuery":
        response, errcode, errdetail = yield self.run_query(http_request_data)
      elif method == "BeginTransaction":
        response, errcode, errdetail = yield self.begin_transaction_request(
          app_id, http_request_data)
      elif method == "Commit":
        response, errcode, errdetail = yield self.commit_transaction_request(
          app_id, http_request_data)
      elif method == "Rollback":
        response, errcode, errdetail = self.rollback_transaction_request(
          app_id, http_request_data)
      elif method == "AllocateIds":
        response, errcode, errdetail = yield self.allocate_ids_request(
          app_id, http_request_data)
      elif method == "CreateIndex":
        response, errcode, errdetail = yield self.create_index_request(
          app_id, http_request_data)
      elif method == "GetIndices":
        response, errcode, errdetail = yield self.get_indices_request(app_id)
      elif method == "UpdateIndex":
        response, errcode, errdetail = self.update_index_request(
          app_id, http_request_data)
      elif method == "DeleteIndex":
        response, errcode, errdetail = yield self.delete_index_request(
          app_id, http_request_data)
      elif method == 'AddActions':
        response, errcode, errdetail = yield self.add_actions_request(
          app_id, http_request_data, service_id, version_id)
      elif method == 'datastore_v4.AllocateIds':
        response, errcode, errdetail = yield self.v4_allocate_ids_request(
          app_id, http_request_data)
      else:
        errcode = datastore_pb.Error.BAD_REQUEST
        errdetail = "Unknown datastore message"
  
      time_taken = time.time() - start
      if method in STATS:
        if errcode in STATS[method]:
          prev_req, pre_time = STATS[method][errcode]
          STATS[method][errcode] = prev_req + 1, pre_time + time_taken
        else:
          STATS[method][errcode] = (1, time_taken)
      else:
        STATS[method] = {}
        STATS[method][errcode] = (1, time_taken)
  
      apiresponse.set_response(response)
      if errcode != 0:
        apperror_pb = apiresponse.mutable_application_error()
        apperror_pb.set_code(errcode)
        apperror_pb.set_detail(errdetail)
  
      self.write(apiresponse.Encode())
  
    @gen.coroutine
    def begin_transaction_request(self, app_id, http_request_data):
      """ Handles the intial request to start a transaction. Replies with 
          a unique identifier to handle this transaction in future requests.
    
      Args:
        app_id: The application ID requesting the transaction.
        http_request_data: The encoded request.
      Returns:
        An encoded transaction protocol buffer with a unique handler.
      """
      global datastore_access
      begin_transaction_req_pb = datastore_pb.BeginTransactionRequest(
        http_request_data)
      multiple_eg = False
      if begin_transaction_req_pb.has_allow_multiple_eg():
        multiple_eg = bool(begin_transaction_req_pb.allow_multiple_eg())
  
      if READ_ONLY:
        logger.warning('Unable to begin transaction in read-only mode: {}'.
                       format(begin_transaction_req_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        handle = yield datastore_access.setup_transaction(app_id, multiple_eg)
      except dbconstants.InternalError as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except (zktransaction.ZKInternalException,
              dbconstants.AppScaleDBConnectionError) as error:
        logger.exception('Unable to begin transaction')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           str(error)))
  
      transaction_pb = datastore_pb.Transaction()
      transaction_pb.set_app(app_id)
      transaction_pb.set_handle(handle)
      raise gen.Return((transaction_pb.Encode(), 0, ''))
  
    @gen.coroutine
    def commit_transaction_request(self, app_id, http_request_data):
      """ Handles the commit phase of a transaction.
  
      Args:
        app_id: The application ID requesting the transaction commit.
        http_request_data: The encoded request of datastore_pb.Transaction.
      Returns:
        An encoded protocol buffer commit response.
      """
      global datastore_access
  
      if READ_ONLY:
        transaction_pb = datastore_pb.Transaction(http_request_data)
        logger.warning('Unable to commit in read-only mode: {}'.
          format(transaction_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      txid = datastore_pb.Transaction(http_request_data).handle()
  
      try:
        yield datastore_access.apply_txn_changes(app_id, txid)
      except (dbconstants.TxTimeoutException, dbconstants.Timeout) as timeout:
        raise gen.Return(('', datastore_pb.Error.TIMEOUT, str(timeout)))
      except dbconstants.AppScaleDBConnectionError:
        logger.exception('DB connection error during commit')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Datastore connection error on Commit request.'))
      except dbconstants.ConcurrentModificationException as error:
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
      except (dbconstants.TooManyGroupsException,
              dbconstants.BadRequest) as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
  
      commitres_pb = datastore_pb.CommitResponse()
      raise gen.Return((commitres_pb.Encode(), 0, ''))
  
    def rollback_transaction_request(self, app_id, http_request_data):
      """ Handles the rollback phase of a transaction.
  
      Args:
        app_id: The application ID requesting the rollback.
        http_request_data: The encoded request.
      Returns:
        An encoded protocol buffer void response.
      """
      global datastore_access
  
      if READ_ONLY:
        logger.warning('Unable to rollback in read-only mode: {}'.
          format(http_request_data))
        return ('', datastore_pb.Error.CAPABILITY_DISABLED,
                'Datastore is in read-only mode.')
  
      try:
        txn = datastore_pb.Transaction(http_request_data)
      except ProtocolBufferDecodeError as error:
        return '', datastore_pb.Error.BAD_REQUEST, str(error)
  
      try:
        datastore_access.rollback_transaction(app_id, txn.handle())
      except dbconstants.InternalError as error:
        logger.exception('Unable to rollback transaction')
        return '', datastore_pb.Error.INTERNAL_ERROR, str(error)
      except Exception as error:
        logger.exception('Unable to rollback transaction')
        return '', datastore_pb.Error.INTERNAL_ERROR, str(error)
  
      return api_base_pb.VoidProto().Encode(), 0, ''
  
    @gen.coroutine
    def run_query(self, http_request_data):
      """ High level function for running queries.
  
      Args:
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        Returns an encoded query response.
      """
      global datastore_access
      query = datastore_pb.Query(http_request_data)
      clone_qr_pb = UnprocessedQueryResult()
      try:
        yield datastore_access._dynamic_run_query(query, clone_qr_pb)
      except dbconstants.BadRequest as error:
        raise gen.Return( ('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKBadRequest as error:
        logger.exception(
          'Illegal arguments in transaction during {}'.format(query))
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKInternalException as error:
        logger.exception('ZKInternalException during {}'.format(query))
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except zktransaction.ZKTransactionException as error:
        logger.exception('Concurrent transaction during {}'.format(query))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
      except dbconstants.AppScaleDBConnectionError as error:
        logger.exception('DB connection error during query')
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.NeedsIndex as error:
        raise gen.Return(('', datastore_pb.Error.NEED_INDEX, str(error)))
  
      raise gen.Return((clone_qr_pb.Encode(), 0, ''))
  
    @gen.coroutine
    def create_index_request(self, app_id, http_request_data):
      """ High level function for creating composite indexes.
  
      Args:
         app_id: Name of the application.
         http_request_data: Stores the protocol buffer request from the 
                 AppServer.
      Returns: 
         Returns an encoded response.
      """
      global datastore_access
      request = entity_pb.CompositeIndex(http_request_data)
      response = api_base_pb.Integer64Proto()
  
      if READ_ONLY:
        logger.warning('Unable to create in read-only mode: {}'.format(request))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        index_id = yield datastore_access.create_composite_index(app_id, request)
        response.set_value(index_id)
      except dbconstants.AppScaleDBConnectionError as error:
        logger.exception('DB connection error during index creation')
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      raise gen.Return((response.Encode(), 0, ''))
  
    def update_index_request(self, app_id, http_request_data):
      """ High level function for updating a composite index.
  
      Args:
        app_id: A string containing the application ID.
        http_request_data: A string containing the protocol buffer request
          from the AppServer.
      Returns:
         A tuple containing an encoded response, error code, and error details.
      """
      global datastore_access
      index = entity_pb.CompositeIndex(http_request_data)
      response = api_base_pb.VoidProto()
  
      if READ_ONLY:
        logger.warning('Unable to update in read-only mode: {}'.format(index))
        return ('', datastore_pb.Error.CAPABILITY_DISABLED,
                'Datastore is in read-only mode.')
  
      state = index.state()
      if state not in [index.READ_WRITE, index.WRITE_ONLY]:
        state_name = entity_pb.CompositeIndex.State_Name(state)
        error_message = 'Unable to update index because state is {}. '\
          'Index: {}'.format(state_name, index)
        logger.error(error_message)
        return '', datastore_pb.Error.PERMISSION_DENIED, error_message
      else:
        # Updating index in background so we can return a response quickly.
        IOLoop.current().spawn_callback(
          datastore_access.update_composite_index, app_id, index)
  
      return response.Encode(), 0, ''
  
    @gen.coroutine
    def delete_index_request(self, app_id, http_request_data):
      """ Deletes a composite index for a given application.
    
      Args:
        app_id: Name of the application.
        http_request_data: A serialized CompositeIndices item
      Returns:
        A Tuple of an encoded entity_pb.VoidProto, error code, and 
        error explanation.
      """
      global datastore_access
      request = entity_pb.CompositeIndex(http_request_data)
      response = api_base_pb.VoidProto()
  
      if READ_ONLY:
        logger.warning('Unable to delete in read-only mode: {}'.format(request))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.delete_composite_index_metadata(app_id, request)
      except (dbconstants.AppScaleDBConnectionError,
              dbconstants.InternalError) as error:
        logger.exception('DB connection error during index deletion')
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
  
      raise gen.Return((response.Encode(), 0, ''))
  
    @gen.coroutine
    def get_indices_request(self, app_id):
      """ Gets the indices of the given application.
  
      Args:
        app_id: Name of the application.
      Returns: 
        A Tuple of an encoded response, error code, and error explanation.
      """
      global datastore_access
      response = datastore_pb.CompositeIndices()
      try:
        indices = datastore_access.get_indexes(app_id)
      except (dbconstants.AppScaleDBConnectionError,
              dbconstants.InternalError) as error:
        logger.exception(
          'Internal error while fetching indices for {}'.format(app_id))
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
  
      for index in indices:
        new_index = response.add_index()
        new_index.MergeFrom(index)
  
      raise gen.Return((response.Encode(), 0, ''))
  
    @gen.coroutine
    def allocate_ids_request(self, app_id, http_request_data):
      """ High level function for getting unique identifiers for entities.
  
      Args:
         app_id: Name of the application.
         http_request_data: Stores the protocol buffer request from the 
                 AppServer.
      Returns: 
         Returns an encoded response.
      Raises:
         NotImplementedError: when requesting a max id.
      """
      request = datastore_pb.AllocateIdsRequest(http_request_data)
  
      if request.has_max() and request.has_size():
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Both size and max cannot be set.'))
      if not (request.has_max() or request.has_size()):
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Either size or max must be set.'))
  
      if request.has_size():
        coroutine = datastore_access.allocate_size
        args = (app_id, request.size())
      else:
        coroutine = datastore_access.allocate_max
        args = (app_id, request.max())
  
      try:
        start, end = yield coroutine(*args)
      except dbconstants.AppScaleBadArg as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except dbconstants.AppScaleDBConnectionError as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
  
      response = datastore_pb.AllocateIdsResponse()
      response.set_start(start)
      response.set_end(end)
      raise gen.Return((response.Encode(), 0, ''))
  
    @staticmethod
    @gen.coroutine
    def v4_allocate_ids_request(app_id, http_request_data):
      """ Reserves entity IDs so that they will not be re-allocated.
  
      Args:
        app_id: Name of the application.
        http_request_data: The protocol buffer request from the AppServer.
      Returns:
         Returns an encoded response.
      """
      request = datastore_v4_pb.AllocateIdsRequest(http_request_data)
      response = datastore_v4_pb.AllocateIdsResponse()
  
      if not request.reserve_list():
        raise gen.Return((response.Encode(), datastore_v4_pb.Error.BAD_REQUEST,
                          'Request must include reserve list'))
  
      ids = [key.path_element_list()[-1].id() for key in request.reserve_list()]
      yield datastore_access.reserve_ids(app_id, ids)
  
      # Forward request to other datastore servers in order to adjust any blocks
      # they've already allocated.
      client = AsyncHTTPClient()
      headers = {'appdata': app_id}
  
      futures = []
      for server in datastore_servers:
        ip, port = server.split(':')
        port = int(port)
        if ip == options.private_ip and port == options.port:
          continue
  
        url = 'http://{}:{}/reserve-keys'.format(ip, port)
        future = client.fetch(url, method='POST', headers=headers,
                              body=http_request_data)
        futures.append(future)
  
      for future in futures:
        yield future
  
      raise gen.Return((response.Encode(), 0, ''))
  
    @gen.coroutine
    def put_request(self, app_id, http_request_data):
      """ High level function for doing puts.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        Returns an encoded put response.
      """
      global datastore_access
  
      putreq_pb = datastore_pb.PutRequest(http_request_data)
      putresp_pb = datastore_pb.PutResponse()
  
      if READ_ONLY:
        logger.warning(
          'Unable to put in read-only mode: {}'.format(putreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.dynamic_put(app_id, putreq_pb, putresp_pb)
        raise gen.Return((putresp_pb.Encode(), 0, ''))
      except (dbconstants.InternalError, zktransaction.ZKInternalException,
              dbconstants.AppScaleDBConnectionError) as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.Timeout as error:
        raise gen.Return(('', datastore_pb.Error.TIMEOUT, str(error)))
      except (dbconstants.BadRequest, zktransaction.ZKBadRequest) as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKTransactionException as error:
        logger.exception('Concurrent transaction during {}'.format(putreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
  
    @gen.coroutine
    def get_request(self, app_id, http_request_data):
      """ High level function for doing gets.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        An encoded get response.
      """
      global datastore_access
      getreq_pb = datastore_pb.GetRequest(http_request_data)
      getresp_pb = datastore_pb.GetResponse()
      try:
        yield datastore_access.dynamic_get(app_id, getreq_pb, getresp_pb)
      except zktransaction.ZKBadRequest as error:
        logger.exception('Illegal argument during {}'.format(getreq_pb))
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKInternalException as error:
        logger.exception('ZKInternalException during {}'.format(getreq_pb))
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except zktransaction.ZKTransactionException as error:
        logger.exception('Concurrent transaction during {}'.format(getreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
      except dbconstants.AppScaleDBConnectionError:
        logger.exception('DB connection error during get')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Datastore connection error on get.'))
  
      raise gen.Return((getresp_pb.Encode(), 0, ''))
  
    @gen.coroutine
    def delete_request(self, app_id, http_request_data):
      """ High level function for doing deletes.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        An encoded delete response.
      """
      global datastore_access
  
      delreq_pb = datastore_pb.DeleteRequest( http_request_data )
      delresp_pb = api_base_pb.VoidProto()
  
      if READ_ONLY:
        logger.warning(
          'Unable to delete in read-only mode: {}'.format(delreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.dynamic_delete(app_id, delreq_pb)
        raise gen.Return((delresp_pb.Encode(), 0, ''))
      except dbconstants.InternalError as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.Timeout as error:
        raise gen.Return(('', datastore_pb.Error.TIMEOUT, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKBadRequest as error:
        logger.exception('Illegal argument during {}'.format(delreq_pb))
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKInternalException:
        logger.exception('ZKInternalException during {}'.format(delreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Internal error with ZooKeeper connection.'))
      except zktransaction.ZKTransactionException:
        logger.exception('Concurrent transaction during {}'.format(delreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION,
           'Concurrent transaction exception on delete.'))
      except dbconstants.AppScaleDBConnectionError:
        logger.exception('DB connection error during delete')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Datastore connection error on delete.'))
  
    @gen.coroutine
    def add_actions_request(self, app_id, http_request_data, service_id,
                            version_id):
      """ High level function for adding transactional tasks.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
        service_id: A string specifying the client's service ID.
        version_id: A string specifying the client's version ID.
      Returns:
        An encoded AddActions response.
      """
      global datastore_access
  
      req_pb = taskqueue_service_pb.TaskQueueBulkAddRequest(http_request_data)
      resp_pb = taskqueue_service_pb.TaskQueueBulkAddResponse()
  
      if service_id is None:
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Module header must be defined'))
  
      if version_id is None:
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Version header must be defined'))
  
      if READ_ONLY:
        logger.warning('Unable to add transactional tasks in read-only mode')
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.dynamic_add_actions(
          app_id, req_pb, service_id, version_id)
        raise gen.Return((resp_pb.Encode(), 0, ''))
      except dbconstants.ExcessiveTasks as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except dbconstants.AppScaleDBConnectionError:
        logger.exception('DB connection error')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Datastore connection error when adding transaction tasks.'))
  
  
  def create_server_node():
    """ Creates a server registration entry in ZooKeeper. """
    try:
      zookeeper.handle.create(server_node, ephemeral=True)
    except NodeExistsError:
      # If the server gets restarted, the old node may exist for a short time.
      zookeeper.handle.delete(server_node)
      zookeeper.handle.create(server_node, ephemeral=True)
  
    logger.info('Datastore registered at {}'.format(server_node))
  
  
  def zk_state_listener(state):
    """ Handles changes to ZooKeeper connection state.
  
    Args:
      state: A string specifying the new ZooKeeper connection state.
    """
    if state == KazooState.CONNECTED:
      persistent_create_server_node = retry_data_watch_coroutine(
        server_node, create_server_node)
      IOLoop.instance().add_callback(persistent_create_server_node)
  
  
  def update_servers(new_servers):
    """ Updates the record of active datastore servers.
  
    Args:
      new_servers: A list of strings identifying server locations.
    """
    to_remove = [server for server in datastore_servers
                 if server not in new_servers]
    for old_server in to_remove:
      datastore_servers.remove(old_server)
  
    for new_server in new_servers:
      if new_server not in datastore_servers:
        datastore_servers.add(new_server)
  
  
  def update_servers_watch(new_servers):
    """ Updates the record of active datastore servers.
  
    Args:
      new_servers: A list of strings identifying server locations.
    """
    main_io_loop = IOLoop.instance()
    main_io_loop.add_callback(update_servers, new_servers)
  
  
  pb_application = tornado.web.Application([
    ('/clear', ClearHandler),
    ('/read-only', ReadOnlyHandler),
    ('/reserve-keys', ReserveKeysHandler),
    (r'/*', MainHandler),
  ])
  
  
  def main():
    """ Starts a web service for handing datastore requests. """
  
    global datastore_access
    global server_node
    global zookeeper
    zookeeper_locations = appscale_info.get_zk_locations_string()
  
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--type', choices=dbconstants.VALID_DATASTORES,
                        default=dbconstants.VALID_DATASTORES[0],
                        help='Database type')
    parser.add_argument('-p', '--port', type=int,
                        default=dbconstants.DEFAULT_PORT,
                        help='Datastore server port')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Output debug-level logging')
    args = parser.parse_args()
  
    if args.verbose:
      logging.getLogger('appscale').setLevel(logging.DEBUG)
  
    options.define('private_ip', appscale_info.get_private_ip())
    options.define('port', args.port)
    taskqueue_locations = get_load_balancer_ips()
  
    server_node = '{}/{}:{}'.format(DATASTORE_SERVERS_NODE, options.private_ip,
                                    options.port)
  
    datastore_batch = DatastoreFactory.getDatastore(
      args.type, log_level=logger.getEffectiveLevel())
    zookeeper = zktransaction.ZKTransaction(
      host=zookeeper_locations, db_access=datastore_batch,
      log_level=logger.getEffectiveLevel())
  
    zookeeper.handle.add_listener(zk_state_listener)
    zookeeper.handle.ensure_path(DATASTORE_SERVERS_NODE)
    # Since the client was started before adding the listener, make sure the
    # server node gets created.
    zk_state_listener(zookeeper.handle.state)
    zookeeper.handle.ChildrenWatch(DATASTORE_SERVERS_NODE, update_servers_watch)
  
    transaction_manager = TransactionManager(zookeeper.handle)
    datastore_access = DatastoreDistributed(
      datastore_batch, transaction_manager, zookeeper=zookeeper,
      log_level=logger.getEffectiveLevel(),
      taskqueue_locations=taskqueue_locations)
    index_manager = IndexManager(zookeeper.handle, datastore_access,
                                 perform_admin=True)
    datastore_access.index_manager = index_manager
  
    server = tornado.httpserver.HTTPServer(pb_application)
    server.listen(args.port)
  
    IOLoop.current().start()
  
  
  [... the rest is omitted ...]
files.new_path: AppDB/appscale/datastore/scripts/datastore.py
files.old_contents: # See LICENSE file
  #
  """
  This web service interfaces with the datastore. It takes protocol buffer
  requests from AppServers and responds according to the type of request its
  given (Put, Get, Delete, Query, etc).
  """
  import argparse
  import json
  import logging
  import os
  import sys
  import time
  import tornado.httpserver
  import tornado.web
  
  from appscale.common import appscale_info
  from appscale.common.appscale_info import get_load_balancer_ips
  from appscale.common.async_retrying import retry_data_watch_coroutine
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from kazoo.client import KazooState
  from kazoo.exceptions import NodeExistsError
  from tornado import gen
  from tornado.httpclient import AsyncHTTPClient
  from tornado.ioloop import IOLoop
  from tornado.options import options
  from .. import dbconstants
  from ..appscale_datastore_batch import DatastoreFactory
  from ..datastore_distributed import DatastoreDistributed
  from ..index_manager import IndexManager
  from ..utils import (clean_app_id,
                       logger,
                       UnprocessedQueryResult)
  from ..zkappscale import zktransaction
  from ..zkappscale.transaction_manager import TransactionManager
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.api import api_base_pb
  from google.appengine.api.taskqueue import taskqueue_service_pb
  from google.appengine.datastore import datastore_pb
  from google.appengine.datastore import datastore_v4_pb
  from google.appengine.datastore import entity_pb
  from google.appengine.ext.remote_api import remote_api_pb
  from google.net.proto.ProtocolBuffer import ProtocolBufferDecodeError
  
  # Global for accessing the datastore. An instance of DatastoreDistributed.
  datastore_access = None
  
  # A record of active datastore servers.
  datastore_servers = set()
  
  # The ZooKeeper path where this server registers its availability.
  server_node = None
  
  # ZooKeeper global variable for locking
  zookeeper = None
  
  # Determines whether or not to allow datastore writes. Note: After enabling,
  # datastore processes must be restarted and the groomer must be stopped.
  READ_ONLY = False
  
  # Global stats.
  STATS = {}
  
  # The ZooKeeper path where a list of active datastore servers is stored.
  DATASTORE_SERVERS_NODE = '/appscale/datastore/servers'
  
  
  class ClearHandler(tornado.web.RequestHandler):
    """ Defines what to do when the webserver receives a /clear HTTP request. """
    def set_default_headers(self):
      """ Instructs clients to close the connection after each response. """
      self.set_header('Connection', 'close')
  
    @tornado.web.asynchronous
    def post(self):
      """ Handles POST requests for clearing datastore server stats. """
      global STATS
      STATS = {}
      self.write({"message": "Statistics for this server cleared."})
      self.finish()
  
  
  class ReadOnlyHandler(tornado.web.RequestHandler):
    """ Handles requests to check or set read-only mode. """
    def set_default_headers(self):
      """ Instructs clients to close the connection after each response. """
      self.set_header('Connection', 'close')
  
    @tornado.web.asynchronous
    def post(self):
      """ Handle requests to turn read-only mode on or off. """
      global READ_ONLY
  
      payload = self.request.body
      data = json.loads(payload)
      if 'readOnly' not in data:
        self.set_status(dbconstants.HTTP_BAD_REQUEST)
  
      if data['readOnly']:
        READ_ONLY = True
        message = 'Write operations now disabled.'
      else:
        READ_ONLY = False
        message = 'Write operations now enabled.'
  
      logger.info(message)
      self.write({'message': message})
      self.finish()
  
  
  class ReserveKeysHandler(tornado.web.RequestHandler):
    """ Handles v4 AllocateIds requests from other servers. """
    @gen.coroutine
    def post(self):
      """ Prevents the provided IDs from being re-allocated. """
      project_id = self.request.headers['appdata']
      request = datastore_v4_pb.AllocateIdsRequest(self.request.body)
      ids = [key.path_element_list()[-1].id() for key in request.reserve_list()]
      yield datastore_access.reserve_ids(project_id, ids)
  
  
  class MainHandler(tornado.web.RequestHandler):
    """
    Defines what to do when the webserver receives different types of 
    HTTP requests.
    """
    def set_default_headers(self):
      """ Instructs clients to close the connection after each response. """
      self.set_header('Connection', 'close')
  
    def unknown_request(self, app_id, http_request_data, pb_type):
      """ Function which handles unknown protocol buffers.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer
      Raises:
        Raises exception.
      """
      raise NotImplementedError("Unknown request of operation {0}" \
        .format(pb_type))
  
    @gen.coroutine
    def post(self):
      """ Function which handles POST requests. Data of the request is
          the request from the AppServer in an encoded protocol buffer
          format.
      """
      request = self.request
      http_request_data = request.body
      pb_type = request.headers['protocolbuffertype']
      app_data = request.headers['appdata']
      app_data = app_data.split(':')
  
      if len(app_data) == 4:
        app_id, user_email, nick_name, auth_domain = app_data
        os.environ['AUTH_DOMAIN'] = auth_domain
        os.environ['USER_EMAIL'] = user_email
        os.environ['USER_NICKNAME'] = nick_name
        os.environ['APPLICATION_ID'] = app_id
      elif len(app_data) == 1:
        app_id = app_data[0]
        os.environ['APPLICATION_ID'] = app_id
      else:
        return
  
      # If the application identifier has the HRD string prepened, remove it.
      app_id = clean_app_id(app_id)
  
      if pb_type == "Request":
        yield self.remote_request(app_id, http_request_data,
                                  service_id=request.headers.get('Module'),
                                  version_id=request.headers.get('Version'))
      else:
        self.unknown_request(app_id, http_request_data, pb_type)
  
    @tornado.web.asynchronous
    def get(self):
      """ Handles get request for the web server. Returns that it is currently
          up in json.
      """
      self.write(json.dumps(STATS))
      self.finish()
  
    @gen.coroutine
    def remote_request(self, app_id, http_request_data, service_id, version_id):
      """ Receives a remote request to which it should give the correct 
          response. The http_request_data holds an encoded protocol buffer
          of a certain type. Each type has a particular response type. 
      
      Args:
        app_id: The application ID that is sending this request.
        http_request_data: Encoded protocol buffer.
        service_id: A string specifying the client's service ID.
        version_id: A string specifying the client's version ID.
      """
      apirequest = remote_api_pb.Request()
      apirequest.ParseFromString(http_request_data)
      apiresponse = remote_api_pb.Response()
      response = None
      if not apirequest.has_method():
        apirequest.set_method("NOT_FOUND")
      if not apirequest.has_request():
        apirequest.set_method("NOT_FOUND")
        apirequest.clear_request()
      method = apirequest.method()
      http_request_data = apirequest.request()
      start = time.time()
  
      request_log = method
      if apirequest.has_request_id():
        request_log += ': {}'.format(apirequest.request_id())
      logger.debug(request_log)
  
      if method == "Put":
        response, errcode, errdetail = yield self.put_request(
          app_id, http_request_data)
      elif method == "Get":
        response, errcode, errdetail = yield self.get_request(
          app_id, http_request_data)
      elif method == "Delete":
        response, errcode, errdetail = yield self.delete_request(
          app_id, http_request_data)
      elif method == "RunQuery":
        response, errcode, errdetail = yield self.run_query(http_request_data)
      elif method == "BeginTransaction":
        response, errcode, errdetail = yield self.begin_transaction_request(
          app_id, http_request_data)
      elif method == "Commit":
        response, errcode, errdetail = yield self.commit_transaction_request(
          app_id, http_request_data)
      elif method == "Rollback":
        response, errcode, errdetail = self.rollback_transaction_request(
          app_id, http_request_data)
      elif method == "AllocateIds":
        response, errcode, errdetail = yield self.allocate_ids_request(
          app_id, http_request_data)
      elif method == "CreateIndex":
        response, errcode, errdetail = yield self.create_index_request(
          app_id, http_request_data)
      elif method == "GetIndices":
        response, errcode, errdetail = yield self.get_indices_request(app_id)
      elif method == "UpdateIndex":
        response, errcode, errdetail = self.update_index_request(
          app_id, http_request_data)
      elif method == "DeleteIndex":
        response, errcode, errdetail = yield self.delete_index_request(
          app_id, http_request_data)
      elif method == 'AddActions':
        response, errcode, errdetail = yield self.add_actions_request(
          app_id, http_request_data, service_id, version_id)
      elif method == 'datastore_v4.AllocateIds':
        response, errcode, errdetail = yield self.v4_allocate_ids_request(
          app_id, http_request_data)
      else:
        errcode = datastore_pb.Error.BAD_REQUEST
        errdetail = "Unknown datastore message"
  
      time_taken = time.time() - start
      if method in STATS:
        if errcode in STATS[method]:
          prev_req, pre_time = STATS[method][errcode]
          STATS[method][errcode] = prev_req + 1, pre_time + time_taken
        else:
          STATS[method][errcode] = (1, time_taken)
      else:
        STATS[method] = {}
        STATS[method][errcode] = (1, time_taken)
  
      apiresponse.set_response(response)
      if errcode != 0:
        apperror_pb = apiresponse.mutable_application_error()
        apperror_pb.set_code(errcode)
        apperror_pb.set_detail(errdetail)
  
      self.write(apiresponse.Encode())
  
    @gen.coroutine
    def begin_transaction_request(self, app_id, http_request_data):
      """ Handles the intial request to start a transaction. Replies with 
          a unique identifier to handle this transaction in future requests.
    
      Args:
        app_id: The application ID requesting the transaction.
        http_request_data: The encoded request.
      Returns:
        An encoded transaction protocol buffer with a unique handler.
      """
      global datastore_access
      begin_transaction_req_pb = datastore_pb.BeginTransactionRequest(
        http_request_data)
      multiple_eg = False
      if begin_transaction_req_pb.has_allow_multiple_eg():
        multiple_eg = bool(begin_transaction_req_pb.allow_multiple_eg())
  
      if READ_ONLY:
        logger.warning('Unable to begin transaction in read-only mode: {}'.
                       format(begin_transaction_req_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        handle = yield datastore_access.setup_transaction(app_id, multiple_eg)
      except dbconstants.InternalError as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except (zktransaction.ZKInternalException,
              dbconstants.AppScaleDBConnectionError) as error:
        logger.exception('Unable to begin transaction')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           str(error)))
  
      transaction_pb = datastore_pb.Transaction()
      transaction_pb.set_app(app_id)
      transaction_pb.set_handle(handle)
      raise gen.Return((transaction_pb.Encode(), 0, ''))
  
    @gen.coroutine
    def commit_transaction_request(self, app_id, http_request_data):
      """ Handles the commit phase of a transaction.
  
      Args:
        app_id: The application ID requesting the transaction commit.
        http_request_data: The encoded request of datastore_pb.Transaction.
      Returns:
        An encoded protocol buffer commit response.
      """
      global datastore_access
  
      if READ_ONLY:
        transaction_pb = datastore_pb.Transaction(http_request_data)
        logger.warning('Unable to commit in read-only mode: {}'.
          format(transaction_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      result = yield datastore_access.commit_transaction(
        app_id, http_request_data)
      raise gen.Return(result)
  
    def rollback_transaction_request(self, app_id, http_request_data):
      """ Handles the rollback phase of a transaction.
  
      Args:
        app_id: The application ID requesting the rollback.
        http_request_data: The encoded request.
      Returns:
        An encoded protocol buffer void response.
      """
      global datastore_access
  
      if READ_ONLY:
        logger.warning('Unable to rollback in read-only mode: {}'.
          format(http_request_data))
        return ('', datastore_pb.Error.CAPABILITY_DISABLED,
                'Datastore is in read-only mode.')
  
      try:
        txn = datastore_pb.Transaction(http_request_data)
      except ProtocolBufferDecodeError as error:
        return '', datastore_pb.Error.BAD_REQUEST, str(error)
  
      try:
        datastore_access.rollback_transaction(app_id, txn.handle())
      except dbconstants.InternalError as error:
        logger.exception('Unable to rollback transaction')
        return '', datastore_pb.Error.INTERNAL_ERROR, str(error)
      except Exception as error:
        logger.exception('Unable to rollback transaction')
        return '', datastore_pb.Error.INTERNAL_ERROR, str(error)
  
      return api_base_pb.VoidProto().Encode(), 0, ''
  
    @gen.coroutine
    def run_query(self, http_request_data):
      """ High level function for running queries.
  
      Args:
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        Returns an encoded query response.
      """
      global datastore_access
      query = datastore_pb.Query(http_request_data)
      clone_qr_pb = UnprocessedQueryResult()
      try:
        yield datastore_access._dynamic_run_query(query, clone_qr_pb)
      except dbconstants.BadRequest as error:
        raise gen.Return( ('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKBadRequest as error:
        logger.exception(
          'Illegal arguments in transaction during {}'.format(query))
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKInternalException as error:
        logger.exception('ZKInternalException during {}'.format(query))
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except zktransaction.ZKTransactionException as error:
        logger.exception('Concurrent transaction during {}'.format(query))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
      except dbconstants.AppScaleDBConnectionError as error:
        logger.exception('DB connection error during query')
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.NeedsIndex as error:
        raise gen.Return(('', datastore_pb.Error.NEED_INDEX, str(error)))
  
      raise gen.Return((clone_qr_pb.Encode(), 0, ''))
  
    @gen.coroutine
    def create_index_request(self, app_id, http_request_data):
      """ High level function for creating composite indexes.
  
      Args:
         app_id: Name of the application.
         http_request_data: Stores the protocol buffer request from the 
                 AppServer.
      Returns: 
         Returns an encoded response.
      """
      global datastore_access
      request = entity_pb.CompositeIndex(http_request_data)
      response = api_base_pb.Integer64Proto()
  
      if READ_ONLY:
        logger.warning('Unable to create in read-only mode: {}'.format(request))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        index_id = yield datastore_access.create_composite_index(app_id, request)
        response.set_value(index_id)
      except dbconstants.AppScaleDBConnectionError as error:
        logger.exception('DB connection error during index creation')
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      raise gen.Return((response.Encode(), 0, ''))
  
    def update_index_request(self, app_id, http_request_data):
      """ High level function for updating a composite index.
  
      Args:
        app_id: A string containing the application ID.
        http_request_data: A string containing the protocol buffer request
          from the AppServer.
      Returns:
         A tuple containing an encoded response, error code, and error details.
      """
      global datastore_access
      index = entity_pb.CompositeIndex(http_request_data)
      response = api_base_pb.VoidProto()
  
      if READ_ONLY:
        logger.warning('Unable to update in read-only mode: {}'.format(index))
        return ('', datastore_pb.Error.CAPABILITY_DISABLED,
                'Datastore is in read-only mode.')
  
      state = index.state()
      if state not in [index.READ_WRITE, index.WRITE_ONLY]:
        state_name = entity_pb.CompositeIndex.State_Name(state)
        error_message = 'Unable to update index because state is {}. '\
          'Index: {}'.format(state_name, index)
        logger.error(error_message)
        return '', datastore_pb.Error.PERMISSION_DENIED, error_message
      else:
        # Updating index in background so we can return a response quickly.
        IOLoop.current().spawn_callback(
          datastore_access.update_composite_index, app_id, index)
  
      return response.Encode(), 0, ''
  
    @gen.coroutine
    def delete_index_request(self, app_id, http_request_data):
      """ Deletes a composite index for a given application.
    
      Args:
        app_id: Name of the application.
        http_request_data: A serialized CompositeIndices item
      Returns:
        A Tuple of an encoded entity_pb.VoidProto, error code, and 
        error explanation.
      """
      global datastore_access
      request = entity_pb.CompositeIndex(http_request_data)
      response = api_base_pb.VoidProto()
  
      if READ_ONLY:
        logger.warning('Unable to delete in read-only mode: {}'.format(request))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.delete_composite_index_metadata(app_id, request)
      except (dbconstants.AppScaleDBConnectionError,
              dbconstants.InternalError) as error:
        logger.exception('DB connection error during index deletion')
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
  
      raise gen.Return((response.Encode(), 0, ''))
  
    @gen.coroutine
    def get_indices_request(self, app_id):
      """ Gets the indices of the given application.
  
      Args:
        app_id: Name of the application.
      Returns: 
        A Tuple of an encoded response, error code, and error explanation.
      """
      global datastore_access
      response = datastore_pb.CompositeIndices()
      try:
        indices = datastore_access.get_indexes(app_id)
      except (dbconstants.AppScaleDBConnectionError,
              dbconstants.InternalError) as error:
        logger.exception(
          'Internal error while fetching indices for {}'.format(app_id))
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
  
      for index in indices:
        new_index = response.add_index()
        new_index.MergeFrom(index)
  
      raise gen.Return((response.Encode(), 0, ''))
  
    @gen.coroutine
    def allocate_ids_request(self, app_id, http_request_data):
      """ High level function for getting unique identifiers for entities.
  
      Args:
         app_id: Name of the application.
         http_request_data: Stores the protocol buffer request from the 
                 AppServer.
      Returns: 
         Returns an encoded response.
      Raises:
         NotImplementedError: when requesting a max id.
      """
      request = datastore_pb.AllocateIdsRequest(http_request_data)
  
      if request.has_max() and request.has_size():
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Both size and max cannot be set.'))
      if not (request.has_max() or request.has_size()):
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Either size or max must be set.'))
  
      if request.has_size():
        coroutine = datastore_access.allocate_size
        args = (app_id, request.size())
      else:
        coroutine = datastore_access.allocate_max
        args = (app_id, request.max())
  
      try:
        start, end = yield coroutine(*args)
      except dbconstants.AppScaleBadArg as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except dbconstants.AppScaleDBConnectionError as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
  
      response = datastore_pb.AllocateIdsResponse()
      response.set_start(start)
      response.set_end(end)
      raise gen.Return((response.Encode(), 0, ''))
  
    @staticmethod
    @gen.coroutine
    def v4_allocate_ids_request(app_id, http_request_data):
      """ Reserves entity IDs so that they will not be re-allocated.
  
      Args:
        app_id: Name of the application.
        http_request_data: The protocol buffer request from the AppServer.
      Returns:
         Returns an encoded response.
      """
      request = datastore_v4_pb.AllocateIdsRequest(http_request_data)
      response = datastore_v4_pb.AllocateIdsResponse()
  
      if not request.reserve_list():
        raise gen.Return((response.Encode(), datastore_v4_pb.Error.BAD_REQUEST,
                          'Request must include reserve list'))
  
      ids = [key.path_element_list()[-1].id() for key in request.reserve_list()]
      yield datastore_access.reserve_ids(app_id, ids)
  
      # Forward request to other datastore servers in order to adjust any blocks
      # they've already allocated.
      client = AsyncHTTPClient()
      headers = {'appdata': app_id}
  
      futures = []
      for server in datastore_servers:
        ip, port = server.split(':')
        port = int(port)
        if ip == options.private_ip and port == options.port:
          continue
  
        url = 'http://{}:{}/reserve-keys'.format(ip, port)
        future = client.fetch(url, method='POST', headers=headers,
                              body=http_request_data)
        futures.append(future)
  
      for future in futures:
        yield future
  
      raise gen.Return((response.Encode(), 0, ''))
  
    @gen.coroutine
    def put_request(self, app_id, http_request_data):
      """ High level function for doing puts.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        Returns an encoded put response.
      """
      global datastore_access
  
      putreq_pb = datastore_pb.PutRequest(http_request_data)
      putresp_pb = datastore_pb.PutResponse()
  
      if READ_ONLY:
        logger.warning(
          'Unable to put in read-only mode: {}'.format(putreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.dynamic_put(app_id, putreq_pb, putresp_pb)
        raise gen.Return((putresp_pb.Encode(), 0, ''))
      except (dbconstants.InternalError, zktransaction.ZKInternalException,
              dbconstants.AppScaleDBConnectionError) as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.Timeout as error:
        raise gen.Return(('', datastore_pb.Error.TIMEOUT, str(error)))
      except (dbconstants.BadRequest, zktransaction.ZKBadRequest) as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKTransactionException as error:
        logger.exception('Concurrent transaction during {}'.format(putreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
  
    @gen.coroutine
    def get_request(self, app_id, http_request_data):
      """ High level function for doing gets.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        An encoded get response.
      """
      global datastore_access
      getreq_pb = datastore_pb.GetRequest(http_request_data)
      getresp_pb = datastore_pb.GetResponse()
      try:
        yield datastore_access.dynamic_get(app_id, getreq_pb, getresp_pb)
      except zktransaction.ZKBadRequest as error:
        logger.exception('Illegal argument during {}'.format(getreq_pb))
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKInternalException as error:
        logger.exception('ZKInternalException during {}'.format(getreq_pb))
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except zktransaction.ZKTransactionException as error:
        logger.exception('Concurrent transaction during {}'.format(getreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION, str(error)))
      except dbconstants.AppScaleDBConnectionError:
        logger.exception('DB connection error during get')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Datastore connection error on get.'))
  
      raise gen.Return((getresp_pb.Encode(), 0, ''))
  
    @gen.coroutine
    def delete_request(self, app_id, http_request_data):
      """ High level function for doing deletes.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
      Returns:
        An encoded delete response.
      """
      global datastore_access
  
      delreq_pb = datastore_pb.DeleteRequest( http_request_data )
      delresp_pb = api_base_pb.VoidProto()
  
      if READ_ONLY:
        logger.warning(
          'Unable to delete in read-only mode: {}'.format(delreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.dynamic_delete(app_id, delreq_pb)
        raise gen.Return((delresp_pb.Encode(), 0, ''))
      except dbconstants.InternalError as error:
        raise gen.Return(('', datastore_pb.Error.INTERNAL_ERROR, str(error)))
      except dbconstants.Timeout as error:
        raise gen.Return(('', datastore_pb.Error.TIMEOUT, str(error)))
      except dbconstants.BadRequest as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKBadRequest as error:
        logger.exception('Illegal argument during {}'.format(delreq_pb))
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except zktransaction.ZKInternalException:
        logger.exception('ZKInternalException during {}'.format(delreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Internal error with ZooKeeper connection.'))
      except zktransaction.ZKTransactionException:
        logger.exception('Concurrent transaction during {}'.format(delreq_pb))
        raise gen.Return(
          ('', datastore_pb.Error.CONCURRENT_TRANSACTION,
           'Concurrent transaction exception on delete.'))
      except dbconstants.AppScaleDBConnectionError:
        logger.exception('DB connection error during delete')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Datastore connection error on delete.'))
  
    @gen.coroutine
    def add_actions_request(self, app_id, http_request_data, service_id,
                            version_id):
      """ High level function for adding transactional tasks.
  
      Args:
        app_id: Name of the application.
        http_request_data: Stores the protocol buffer request from the AppServer.
        service_id: A string specifying the client's service ID.
        version_id: A string specifying the client's version ID.
      Returns:
        An encoded AddActions response.
      """
      global datastore_access
  
      req_pb = taskqueue_service_pb.TaskQueueBulkAddRequest(http_request_data)
      resp_pb = taskqueue_service_pb.TaskQueueBulkAddResponse()
  
      if service_id is None:
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Module header must be defined'))
  
      if version_id is None:
        raise gen.Return(
          ('', datastore_pb.Error.BAD_REQUEST,
           'Version header must be defined'))
  
      if READ_ONLY:
        logger.warning('Unable to add transactional tasks in read-only mode')
        raise gen.Return(
          ('', datastore_pb.Error.CAPABILITY_DISABLED,
           'Datastore is in read-only mode.'))
  
      try:
        yield datastore_access.dynamic_add_actions(
          app_id, req_pb, service_id, version_id)
        raise gen.Return((resp_pb.Encode(), 0, ''))
      except dbconstants.ExcessiveTasks as error:
        raise gen.Return(('', datastore_pb.Error.BAD_REQUEST, str(error)))
      except dbconstants.AppScaleDBConnectionError:
        logger.exception('DB connection error')
        raise gen.Return(
          ('', datastore_pb.Error.INTERNAL_ERROR,
           'Datastore connection error when adding transaction tasks.'))
  
  
  def create_server_node():
    """ Creates a server registration entry in ZooKeeper. """
    try:
      zookeeper.handle.create(server_node, ephemeral=True)
    except NodeExistsError:
      # If the server gets restarted, the old node may exist for a short time.
      zookeeper.handle.delete(server_node)
      zookeeper.handle.create(server_node, ephemeral=True)
  
    logger.info('Datastore registered at {}'.format(server_node))
  
  
  def zk_state_listener(state):
    """ Handles changes to ZooKeeper connection state.
  
    Args:
      state: A string specifying the new ZooKeeper connection state.
    """
    if state == KazooState.CONNECTED:
      persistent_create_server_node = retry_data_watch_coroutine(
        server_node, create_server_node)
      IOLoop.instance().add_callback(persistent_create_server_node)
  
  
  def update_servers(new_servers):
    """ Updates the record of active datastore servers.
  
    Args:
      new_servers: A list of strings identifying server locations.
    """
    to_remove = [server for server in datastore_servers
                 if server not in new_servers]
    for old_server in to_remove:
      datastore_servers.remove(old_server)
  
    for new_server in new_servers:
      if new_server not in datastore_servers:
        datastore_servers.add(new_server)
  
  
  def update_servers_watch(new_servers):
    """ Updates the record of active datastore servers.
  
    Args:
      new_servers: A list of strings identifying server locations.
    """
    main_io_loop = IOLoop.instance()
    main_io_loop.add_callback(update_servers, new_servers)
  
  
  pb_application = tornado.web.Application([
    ('/clear', ClearHandler),
    ('/read-only', ReadOnlyHandler),
    ('/reserve-keys', ReserveKeysHandler),
    (r'/*', MainHandler),
  ])
  
  
  def main():
    """ Starts a web service for handing datastore requests. """
  
    global datastore_access
    global server_node
    global zookeeper
    zookeeper_locations = appscale_info.get_zk_locations_string()
  
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--type', choices=dbconstants.VALID_DATASTORES,
                        default=dbconstants.VALID_DATASTORES[0],
                        help='Database type')
    parser.add_argument('-p', '--port', type=int,
                        default=dbconstants.DEFAULT_PORT,
                        help='Datastore server port')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Output debug-level logging')
    args = parser.parse_args()
  
    if args.verbose:
      logging.getLogger('appscale').setLevel(logging.DEBUG)
  
    options.define('private_ip', appscale_info.get_private_ip())
    options.define('port', args.port)
    taskqueue_locations = get_load_balancer_ips()
  
    server_node = '{}/{}:{}'.format(DATASTORE_SERVERS_NODE, options.private_ip,
                                    options.port)
  
    datastore_batch = DatastoreFactory.getDatastore(
      args.type, log_level=logger.getEffectiveLevel())
    zookeeper = zktransaction.ZKTransaction(
      host=zookeeper_locations, db_access=datastore_batch,
      log_level=logger.getEffectiveLevel())
  
    zookeeper.handle.add_listener(zk_state_listener)
    zookeeper.handle.ensure_path(DATASTORE_SERVERS_NODE)
    # Since the client was started before adding the listener, make sure the
    # server node gets created.
    zk_state_listener(zookeeper.handle.state)
    zookeeper.handle.ChildrenWatch(DATASTORE_SERVERS_NODE, update_servers_watch)
  
    transaction_manager = TransactionManager(zookeeper.handle)
    datastore_access = DatastoreDistributed(
      datastore_batch, transaction_manager, zookeeper=zookeeper,
      log_level=logger.getEffectiveLevel(),
      taskqueue_locations=taskqueue_locations)
    index_manager = IndexManager(zookeeper.handle, datastore_access,
                                 perform_admin=True)
    datastore_access.index_manager = index_manager
  
    server = tornado.httpserver.HTTPServer(pb_application)
    server.listen(args.port)
  
    IOLoop.current().start()
  
  
  [... the rest is omitted ...]
files.old_path: AppDB/appscale/datastore/scripts/datastore.py
files.new_contents: #!/usr/bin/env python
  # Programmer: Navraj Chohan <nlake44@gmail.com>
  
  import datetime
  import random
  import sys
  import unittest
  
  from tornado import gen, testing
  
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from appscale.datastore import dbconstants
  from appscale.datastore import utils
  from appscale.datastore.datastore_distributed import DatastoreDistributed
  from appscale.datastore.dbconstants import APP_ENTITY_SCHEMA
  from appscale.datastore.dbconstants import JOURNAL_SCHEMA
  from appscale.datastore.dbconstants import TOMBSTONE
  from appscale.datastore.cassandra_env.entity_id_allocator import\
    ScatteredAllocator
  
  from appscale.datastore.cassandra_env.utils import deletions_for_entity
  from appscale.datastore.cassandra_env.utils import index_deletions
  from appscale.datastore.cassandra_env.utils import mutations_for_entity
  
  from appscale.datastore.utils import (
    encode_index_pb,
    get_entity_key,
    get_entity_kind,
    get_index_key_from_params,
    get_index_kv_from_tuple,
    get_kind_key
  )
  
  from appscale.datastore.zkappscale.entity_lock import EntityLock
  from appscale.datastore.zkappscale.zktransaction import ZKTransactionException
  from flexmock import flexmock
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.datastore import entity_pb
  from google.appengine.datastore import datastore_pb
  from google.appengine.ext import db
  
  
  class Item(db.Model):
    name = db.StringProperty(required = True)
  
  
  ASYNC_NONE = gen.Future()
  ASYNC_NONE.set_result(None)
  
  
  class TestDatastoreServer(testing.AsyncTestCase):
    """
    A set of test cases for the datastore server (datastore server v2)
    """
    BASIC_ENTITY = ['guestbook', 'Greeting', 'foo', 'content', 'hello world']
  
    def setUp(self, *args, **kwargs):
      super(TestDatastoreServer, self).setUp(*args, **kwargs)
  
  
  
    def get_zookeeper(self):
      zk_handle = flexmock(handler=flexmock(event_object=lambda: None,
                                            sleep_func=lambda: None,
                                            lock_object=lambda: None))
      zk_handle.should_receive('add_listener')
      zookeeper = flexmock(handle=zk_handle)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("increment_and_get_counter").and_return(0,1000)
      return zookeeper
  
    def test_get_entity_kind(self):
      item = Item(name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_entity_kind(key), "Item")
  
    def test_kind_key(self):
      item = Item(name="Dyan", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_kind_key("howdy", key.key().path()), "howdy\x00Item\x01Item:0000000000\x01")
  
      item1 = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item1)
      self.assertEquals(get_kind_key("howdy", key.key().path()), "howdy\x00Item\x01Item:Bob\x01")
  
      item2 = Item(key_name="Frank", name="Frank", _app="hello", parent = item1)
      key = db.model_to_protobuf(item2)
      self.assertEquals(get_kind_key("howdy", key.key().path()),
             "howdy\x00Item\x01Item:Bob\x01Item:Frank\x01")
  
    def test_get_entity_key(self):
      item = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(str(get_entity_key("howdy", key.key().path())), "howdy\x00Item:Bob\x01")
  
    def test_validate_key(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      item = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      dd.validate_key(key.key())
  
    def test_get_table_prefix(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      item = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(dd.get_table_prefix(key), "hello\x00")
  
    def test_get_index_key_from_params(self):
      params = ['a','b','c','d','e']
      self.assertEquals(get_index_key_from_params(params),
                        "a\x00b\x00c\x00d\x00e")
  
    def test_get_index_kv_from_tuple(self):
      item1 = Item(key_name="Bob", name="Bob", _app="hello")
      item2 = Item(key_name="Sally", name="Sally", _app="hello")
      key1 = db.model_to_protobuf(item1)
      key2 = db.model_to_protobuf(item2)
      tuples_list = [("a\x00b",key1),("a\x00b",key2)]
      self.assertEquals(get_index_kv_from_tuple(
        tuples_list), (['a\x00b\x00Item\x00name\x00\x9aBob\x01\x01\x00Item:Bob\x01',
        'a\x00b\x00Item:Bob\x01'],
        ['a\x00b\x00Item\x00name\x00\x9aSally\x01\x01\x00Item:Sally\x01',
        'a\x00b\x00Item:Sally\x01']))
  
    def test_get_composite_index_key(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd = flexmock(dd)
  
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id("appid")
  
      definition = composite_index.mutable_definition()
      definition.set_entity_type("kind")
  
      prop1 = definition.add_property()
      prop1.set_name("prop1")
      prop1.set_direction(1) # ascending
      prop2 = definition.add_property()
      prop2.set_name("prop2")
      prop1.set_direction(1) # ascending
  
      ent = self.get_new_entity_proto("appid", "kind", "entity_name", "prop1", "value", ns="")
  
      self.assertEquals(
        dd.get_composite_index_key(composite_index, ent),
        "appid\x00\x00123\x00\x9avalue\x01\x01\x00\x00kind:entity_name\x01")
  
    @testing.gen_test
    def test_delete_composite_index_metadata(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_delete").and_return(ASYNC_NONE)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd = flexmock(dd)
      dd.index_manager = flexmock(
        projects={'appid': flexmock(delete_index_definition=lambda id_: None)})
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(1)
      yield dd.delete_composite_index_metadata("appid", composite_index)
  
    @testing.gen_test
    def test_create_composite_index(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      zk_handle = flexmock(ensure_path=lambda path: None,
                           get=lambda path: (None, flexmock(version=None)),
                           set=lambda path, value, version: None)
      dd.zookeeper.handle = zk_handle
      dd = flexmock(dd)
      index = entity_pb.CompositeIndex()
      index.set_app_id("appid")
      index.set_state(2)
      definition = index.mutable_definition()
      definition.set_entity_type("kind")
      definition.set_ancestor(0)
      prop1 = definition.add_property()
      prop1.set_name("prop1")
      prop1.set_direction(1) # ascending
      prop2 = definition.add_property()
      prop2.set_name("prop2")
      prop1.set_direction(1) # ascending
  
      index_id = yield dd.create_composite_index("appid", index)
      assert index_id > 0
  
    @testing.gen_test
    def test_insert_composite_indexes(self):
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id("appid")
  
      definition = composite_index.mutable_definition()
      definition.set_entity_type("kind")
  
      prop1 = definition.add_property()
      prop1.set_name("prop1")
      prop1.set_direction(1) # ascending
      prop2 = definition.add_property()
      prop2.set_name("prop2")
      prop1.set_direction(1) # ascending
  
      ent = self.get_new_entity_proto(
        "appid", "kind", "entity_name", "prop1", "value", ns="")
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE).once()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      yield dd.insert_composite_indexes([ent], [composite_index])
  
    @testing.gen_test
    def test_fetch_keys(self):
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      async_get_1 = gen.Future()
      async_get_1.set_result({
        'test\x00blah\x00test_kind:bob\x01': {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
      async_get_2 = gen.Future()
      async_get_2.set_result({
        'test\x00blah\x00test_kind:bob\x01\x000000000002': {
          JOURNAL_SCHEMA[0]: entity_proto1.Encode()
      }})
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_delete").and_return(ASYNC_NONE)
      db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE)
      db_batch.should_receive("batch_get_entity").\
        and_return(async_get_1).\
        and_return(async_get_2)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      fetched = yield dd.fetch_keys([entity_proto1.key()])
      self.assertEquals(fetched[0], {
        'test\x00blah\x00test_kind:bob\x01': {
         'txnID': 1, 'entity': entity_proto1.Encode()
        }
      })
      self.assertEqual(fetched[1], ['test\x00blah\x00test_kind:bob\x01'])
  
    def test_rollback_transcation(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("notify_failed_transaction").and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      commit_request = datastore_pb.Transaction()
      commit_request.set_handle(123)
      commit_request.set_app("aaa")
      self.assertEquals(
        dd.rollback_transaction("app_id", commit_request.handle()), None)
  
    @staticmethod
    def get_new_entity_proto(app_id, kind, entity_name, prop_name, prop_value, ns=""):
      entity_proto = datastore_pb.EntityProto()
  
      reference = entity_proto.mutable_key()
      reference.set_app(app_id)
      reference.set_name_space(ns)
  
      path = reference.mutable_path()
      element = path.add_element()
      element.set_type(kind)
      element.set_name(entity_name)
  
      ent_group = entity_proto.mutable_entity_group()
      eg_element = ent_group.add_element()
      eg_element.set_type(kind)
      eg_element.set_name(entity_name)
  
      prop = entity_proto.add_property()
      prop.set_meaning(datastore_pb.Property.BYTESTRING)
      prop.set_name(prop_name)
      prop.set_multiple(1)
      val = prop.mutable_value()
      val.set_stringvalue(prop_value)
      return entity_proto
  
    @testing.gen_test
    def test_dynamic_put(self):
      db_batch = flexmock(session=flexmock())
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
      entity_proto2 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
      async_result = gen.Future()
      async_result.set_result({entity_key1: {}, entity_key2: {}})
  
      db_batch.should_receive('batch_get_entity').and_return(async_result)
      db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None,
        set_groups=lambda project, txid, groups: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'test': flexmock(indexes_pb=[])})
      putreq_pb = datastore_pb.PutRequest()
      putreq_pb.add_entity()
      putreq_pb.mutable_entity(0).MergeFrom(entity_proto1)
      putreq_pb.add_entity()
      putreq_pb.mutable_entity(1).MergeFrom(entity_proto2)
  
      putresp_pb = datastore_pb.PutResponse()
  
      async_true = gen.Future()
      async_true.set_result(True)
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire').and_return(async_true)
      entity_lock.should_receive('release')
  
      flexmock(ScatteredAllocator).should_receive('next').\
        and_return(random.randint(1, 500))
  
      yield dd.dynamic_put('test', putreq_pb, putresp_pb)
      self.assertEquals(len(putresp_pb.key_list()), 2)
  
    @testing.gen_test
    def test_put_entities(self):
      app_id = 'test'
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      entity_proto1 = self.get_new_entity_proto(
        app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
      entity_proto2 = self.get_new_entity_proto(
        app_id, "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
      entity_list = [entity_proto1, entity_proto2]
      async_result = gen.Future()
      async_result.set_result({entity_key1: {}, entity_key2: {}})
  
      db_batch.should_receive('batch_get_entity').and_return(async_result)
      db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None,
        set_groups=lambda project, txid, groups: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={app_id: flexmock(indexes_pb=[])})
  
      async_true = gen.Future()
      async_true.set_result(True)
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire').and_return(async_true)
      entity_lock.should_receive('release')
  
      yield dd.put_entities(app_id, entity_list)
  
    def test_acquire_locks_for_trans(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(False).once()
      self.assertRaises(TypeError, dd.acquire_locks_for_trans, [1], 1)
  
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(False).and_return(False)
      self.assertRaises(TypeError, dd.acquire_locks_for_trans, [1], 1)
  
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True)
  
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True).and_return(False)
      flexmock(dd).should_receive("get_table_prefix").and_return("prefix").never()
      flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").never()
      self.assertEquals({}, dd.acquire_locks_for_trans([], 1))
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("acquire_lock").once()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      entity = flexmock()
      entity.should_receive("app").and_return("appid")
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True).and_return(True)
      flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").once()
      self.assertEquals({'rootkey':1}, dd.acquire_locks_for_trans([entity], 1))
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("acquire_lock").once().and_raise(ZKTransactionException)
      zookeeper.should_receive("notify_failed_transaction").once()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      entity = flexmock()
      entity.should_receive("app").and_return("appid")
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True).and_return(True)
      flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").once()
      self.assertRaises(ZKTransactionException, dd.acquire_locks_for_trans, [entity], 1)
  
    @testing.gen_test
    def test_delete_entities(self):
      app_id = 'test'
      entity_proto1 = self.get_new_entity_proto(
        app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      row_key = "test\x00blah\x00test_kind:bob\x01"
      row_values = {row_key: {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                              APP_ENTITY_SCHEMA[1]: '1'}}
      async_result = gen.Future()
      async_result.set_result(row_values)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
      db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      row_keys = [entity_proto1.key()]
  
      yield dd.delete_entities(entity_proto1.key(), 1, row_keys)
  
    def test_release_put_locks_for_nontrans(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      entity_proto2 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entities = [entity_proto1, entity_proto2]
      dd.release_locks_for_nontrans(
        "test", entities, {
          'test\x00blah\x00test_kind:bob\x01': 1,
          'test\x00blah\x00test_kind:nancy\x01': 2
        })
  
    def test_root_key_from_entity_key(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      self.assertEquals(
        "test\x00blah\x00test_kind:bob\x01",
        dd.get_root_key_from_entity_key("test\x00blah\x00test_kind:bob\x01nancy\x01")
      )
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      self.assertEquals(
        "test\x00blah\x00test_kind:nancy\x01",
        dd.get_root_key_from_entity_key(entity_proto1.key())
      )
  
    @testing.gen_test
    def test_dynamic_get(self):
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      async_result = gen.Future()
      async_result.set_result({
        "test\x00blah\x00test_kind:nancy\x01": {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
      db_batch.should_receive('record_reads').and_return(ASYNC_NONE)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
      get_resp = datastore_pb.GetResponse()
  
      yield dd.dynamic_get("test", get_req, get_resp)
      self.assertEquals(get_resp.entity_size(), 1)
  
      # Now test while in a transaction
      get_resp = datastore_pb.GetResponse()
      get_req.mutable_transaction().set_handle(1)
      yield dd.dynamic_get("test", get_req, get_resp)
      self.assertEquals(get_resp.entity_size(), 1)
  
    @testing.gen_test
    def test_ancestor_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
      ancestor.MergeFrom(entity_key)
  
      async_result = gen.Future()
      async_result.set_result({
        "test\x00blah\x00test_kind:nancy\x01": {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
  
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('record_reads').and_return(ASYNC_NONE)
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
  
      entity_proto1 = {
        'test\x00blah\x00test_kind:nancy\x01': {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      }
      async_result_1 = gen.Future()
      async_result_1.set_result([entity_proto1, tombstone1])
      async_result_2 = gen.Future()
      async_result_2.set_result([])
  
      db_batch.should_receive("range_query").\
        and_return(async_result_1).\
        and_return(async_result_2)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      yield dd.ancestor_query(query, filter_info)
      # Now with a transaction
      transaction = query.mutable_transaction()
      transaction.set_handle(2)
      yield dd.ancestor_query(query, filter_info)
  
    @testing.gen_test
    def test_kindless_query(self):
      query = datastore_pb.Query()
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
  
      async_result = gen.Future()
      async_result.set_result({
        "test\x00blah\x00test_kind:nancy\x01": {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
  
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
  
      entity_proto1 = {
        'test\x00blah\x00test_kind:nancy\x01': {
          APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      }
      async_result_1 = gen.Future()
      async_result_1.set_result([entity_proto1, tombstone1])
      async_result_2 = gen.Future()
      async_result_2.set_result([])
      db_batch.should_receive("range_query").\
        and_return(async_result_1).\
        and_return(async_result_2)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      zookeeper.should_receive("acquire_lock").and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      filter_info = {
        '__key__' : [[0, 0]]
      }
      yield dd.kindless_query(query, filter_info)
  
    @testing.gen_test
    def test_dynamic_delete(self):
      async_true = gen.Future()
      async_true.set_result(True)
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire').and_return(async_true)
      entity_lock.should_receive('release')
  
      del_request = flexmock()
      del_request.should_receive("key_list")
      del_request.should_receive("has_transaction").never()
      del_request.should_receive("transaction").never()
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None,
        set_groups=lambda project_id, txid, groups: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'guestbook': flexmock(indexes_pb=[])})
      yield dd.dynamic_delete("appid", del_request)
  
      fake_key = entity_pb.Reference()
      fake_key.set_app('foo')
      path = fake_key.mutable_path()
      element = path.add_element()
      element.set_type('bar')
      element.set_id(1)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(True)
      transaction = flexmock()
      transaction.should_receive("handle").and_return(1)
      del_request.should_receive("transaction").and_return(transaction)
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'appid': flexmock(indexes_pb=[])})
      flexmock(utils).should_receive("get_entity_kind").and_return("kind")
      db_batch.should_receive('delete_entities_tx').and_return(ASYNC_NONE)
      yield dd.dynamic_delete("appid", del_request)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(False)
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'appid': flexmock(indexes_pb=[])})
      flexmock(dd).should_receive("delete_entities").and_return(ASYNC_NONE).once()
      yield dd.dynamic_delete("appid", del_request)
  
    def test_reverse_path(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      key = "Project:Synapse\x01Module:Core\x01"
      self.assertEquals(dd.reverse_path(key), "Module:Core\x01Project:Synapse\x01")
  
    def test_remove_exists_filters(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      self.assertEquals(dd.remove_exists_filters({}), {})
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]
      }
      self.assertEquals(dd.remove_exists_filters(filter_info), filter_info)
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EXISTS, "1")],
        "prop2": [(datastore_pb.Query_Filter.EXISTS, "2")]
      }
      self.assertEquals(dd.remove_exists_filters(filter_info), {})
  
    def test_is_zigzag_merge_join(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      query = datastore_pb.Query()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      self.assertEquals(dd.is_zigzag_merge_join(query, {}, {}), False)
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]
      }
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), True)
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")]
      }
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), False)
  
    @testing.gen_test
    def test_zigzag_merge_join(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      query = datastore_pb.Query()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_zigzag_merge_join").and_return(False)
      result = yield dd.zigzag_merge_join(None, None, None)
      self.assertEquals(result, None)
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]
      }
      result = yield dd.zigzag_merge_join(query, filter_info, [])
      self.assertEquals(result, None)
  
    def test_index_deletions(self):
      old_entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # No deletions should occur when the entity doesn't change.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      self.assertListEqual([], index_deletions(old_entity, old_entity))
  
      # When a property changes, the previous index entries should be deleted.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(old_entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
  
      deletions = index_deletions(old_entity, new_entity)
      self.assertEqual(len(deletions), 2)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = old_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # When given an index, an entry should be removed from the composite table.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 3)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.COMPOSITE_TABLE)
  
      # No composite deletions should occur when the entity type differs.
      definition.set_entity_type('TestEntity')
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 2)
  
    def test_deletions_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # Deleting an entity with one property should remove four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      deletions = deletions_for_entity(entity)
      self.assertEqual(len(deletions), 4)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.APP_KIND_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # Deleting an entity with two properties and one composite index should
      # remove seven entries.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = deletions_for_entity(entity, (composite_index,))
      self.assertEqual(len(deletions), 7)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[4]['table'], dbconstants.COMPOSITE_TABLE)
      self.assertEqual(deletions[5]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[6]['table'], dbconstants.APP_KIND_TABLE)
  
    def test_mutations_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
      txn = 1
  
      # Adding an entity with one property should add four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      mutations = mutations_for_entity(entity, txn)
      self.assertEqual(len(mutations), 4)
      self.assertEqual(mutations[0]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[1]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[2]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      # Updating an entity with one property should delete two entries and add
      # four more.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
      mutations = mutations_for_entity(entity, txn, new_entity)
      self.assertEqual(len(mutations), 6)
      self.assertEqual(mutations[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[0]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[1]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[4]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[5]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # Updating one property of an entity with two properties and one composite
      # index should remove three entries and add seven more.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
  
      mutations = mutations_for_entity(entity, txn, new_entity,
                                       (composite_index,))
      self.assertEqual(len(mutations), 10)
      self.assertEqual(mutations[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self
  
  [... the rest is omitted ...]
files.new_path: AppDB/test/unit/test_datastore_server.py
files.old_contents: #!/usr/bin/env python
  # Programmer: Navraj Chohan <nlake44@gmail.com>
  
  import datetime
  import random
  import sys
  import unittest
  
  from tornado import gen, testing
  
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from appscale.datastore import dbconstants
  from appscale.datastore import utils
  from appscale.datastore.datastore_distributed import DatastoreDistributed
  from appscale.datastore.dbconstants import APP_ENTITY_SCHEMA
  from appscale.datastore.dbconstants import JOURNAL_SCHEMA
  from appscale.datastore.dbconstants import TOMBSTONE
  from appscale.datastore.cassandra_env.entity_id_allocator import\
    ScatteredAllocator
  
  from appscale.datastore.cassandra_env.utils import deletions_for_entity
  from appscale.datastore.cassandra_env.utils import index_deletions
  from appscale.datastore.cassandra_env.utils import mutations_for_entity
  
  from appscale.datastore.utils import (
    encode_index_pb,
    get_entity_key,
    get_entity_kind,
    get_index_key_from_params,
    get_index_kv_from_tuple,
    get_kind_key
  )
  
  from appscale.datastore.zkappscale.entity_lock import EntityLock
  from appscale.datastore.zkappscale.zktransaction import ZKTransactionException
  from flexmock import flexmock
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.datastore import entity_pb
  from google.appengine.datastore import datastore_pb
  from google.appengine.ext import db
  
  
  class Item(db.Model):
    name = db.StringProperty(required = True)
  
  
  ASYNC_NONE = gen.Future()
  ASYNC_NONE.set_result(None)
  
  
  class TestDatastoreServer(testing.AsyncTestCase):
    """
    A set of test cases for the datastore server (datastore server v2)
    """
    BASIC_ENTITY = ['guestbook', 'Greeting', 'foo', 'content', 'hello world']
  
    def setUp(self, *args, **kwargs):
      super(TestDatastoreServer, self).setUp(*args, **kwargs)
  
  
  
    def get_zookeeper(self):
      zk_handle = flexmock(handler=flexmock(event_object=lambda: None,
                                            sleep_func=lambda: None,
                                            lock_object=lambda: None))
      zk_handle.should_receive('add_listener')
      zookeeper = flexmock(handle=zk_handle)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("increment_and_get_counter").and_return(0,1000)
      return zookeeper
  
    def test_get_entity_kind(self):
      item = Item(name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_entity_kind(key), "Item")
  
    def test_kind_key(self):
      item = Item(name="Dyan", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_kind_key("howdy", key.key().path()), "howdy\x00Item\x01Item:0000000000\x01")
  
      item1 = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item1)
      self.assertEquals(get_kind_key("howdy", key.key().path()), "howdy\x00Item\x01Item:Bob\x01")
  
      item2 = Item(key_name="Frank", name="Frank", _app="hello", parent = item1)
      key = db.model_to_protobuf(item2)
      self.assertEquals(get_kind_key("howdy", key.key().path()),
             "howdy\x00Item\x01Item:Bob\x01Item:Frank\x01")
  
    def test_get_entity_key(self):
      item = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(str(get_entity_key("howdy", key.key().path())), "howdy\x00Item:Bob\x01")
  
    def test_validate_key(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      item = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      dd.validate_key(key.key())
  
    def test_get_table_prefix(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      item = Item(key_name="Bob", name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(dd.get_table_prefix(key), "hello\x00")
  
    def test_get_index_key_from_params(self):
      params = ['a','b','c','d','e']
      self.assertEquals(get_index_key_from_params(params),
                        "a\x00b\x00c\x00d\x00e")
  
    def test_get_index_kv_from_tuple(self):
      item1 = Item(key_name="Bob", name="Bob", _app="hello")
      item2 = Item(key_name="Sally", name="Sally", _app="hello")
      key1 = db.model_to_protobuf(item1)
      key2 = db.model_to_protobuf(item2)
      tuples_list = [("a\x00b",key1),("a\x00b",key2)]
      self.assertEquals(get_index_kv_from_tuple(
        tuples_list), (['a\x00b\x00Item\x00name\x00\x9aBob\x01\x01\x00Item:Bob\x01',
        'a\x00b\x00Item:Bob\x01'],
        ['a\x00b\x00Item\x00name\x00\x9aSally\x01\x01\x00Item:Sally\x01',
        'a\x00b\x00Item:Sally\x01']))
  
    def test_get_composite_index_key(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd = flexmock(dd)
  
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id("appid")
  
      definition = composite_index.mutable_definition()
      definition.set_entity_type("kind")
  
      prop1 = definition.add_property()
      prop1.set_name("prop1")
      prop1.set_direction(1) # ascending
      prop2 = definition.add_property()
      prop2.set_name("prop2")
      prop1.set_direction(1) # ascending
  
      ent = self.get_new_entity_proto("appid", "kind", "entity_name", "prop1", "value", ns="")
  
      self.assertEquals(
        dd.get_composite_index_key(composite_index, ent),
        "appid\x00\x00123\x00\x9avalue\x01\x01\x00\x00kind:entity_name\x01")
  
    @testing.gen_test
    def test_delete_composite_index_metadata(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_delete").and_return(ASYNC_NONE)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd = flexmock(dd)
      dd.index_manager = flexmock(
        projects={'appid': flexmock(delete_index_definition=lambda id_: None)})
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(1)
      yield dd.delete_composite_index_metadata("appid", composite_index)
  
    @testing.gen_test
    def test_create_composite_index(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      zk_handle = flexmock(ensure_path=lambda path: None,
                           get=lambda path: (None, flexmock(version=None)),
                           set=lambda path, value, version: None)
      dd.zookeeper.handle = zk_handle
      dd = flexmock(dd)
      index = entity_pb.CompositeIndex()
      index.set_app_id("appid")
      index.set_state(2)
      definition = index.mutable_definition()
      definition.set_entity_type("kind")
      definition.set_ancestor(0)
      prop1 = definition.add_property()
      prop1.set_name("prop1")
      prop1.set_direction(1) # ascending
      prop2 = definition.add_property()
      prop2.set_name("prop2")
      prop1.set_direction(1) # ascending
  
      index_id = yield dd.create_composite_index("appid", index)
      assert index_id > 0
  
    @testing.gen_test
    def test_insert_composite_indexes(self):
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id("appid")
  
      definition = composite_index.mutable_definition()
      definition.set_entity_type("kind")
  
      prop1 = definition.add_property()
      prop1.set_name("prop1")
      prop1.set_direction(1) # ascending
      prop2 = definition.add_property()
      prop2.set_name("prop2")
      prop1.set_direction(1) # ascending
  
      ent = self.get_new_entity_proto(
        "appid", "kind", "entity_name", "prop1", "value", ns="")
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE).once()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      yield dd.insert_composite_indexes([ent], [composite_index])
  
    @testing.gen_test
    def test_fetch_keys(self):
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      async_get_1 = gen.Future()
      async_get_1.set_result({
        'test\x00blah\x00test_kind:bob\x01': {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
      async_get_2 = gen.Future()
      async_get_2.set_result({
        'test\x00blah\x00test_kind:bob\x01\x000000000002': {
          JOURNAL_SCHEMA[0]: entity_proto1.Encode()
      }})
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_delete").and_return(ASYNC_NONE)
      db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE)
      db_batch.should_receive("batch_get_entity").\
        and_return(async_get_1).\
        and_return(async_get_2)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      fetched = yield dd.fetch_keys([entity_proto1.key()])
      self.assertEquals(fetched[0], {
        'test\x00blah\x00test_kind:bob\x01': {
         'txnID': 1, 'entity': entity_proto1.Encode()
        }
      })
      self.assertEqual(fetched[1], ['test\x00blah\x00test_kind:bob\x01'])
  
    @testing.gen_test
    def test_commit_transaction(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      transaction_manager = flexmock(
        delete_transaction_id=lambda project, txid: None)
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive('apply_txn_changes').and_return(ASYNC_NONE)
      commit_request = datastore_pb.Transaction()
      commit_request.set_handle(123)
      commit_request.set_app("aaa")
      http_request = commit_request.Encode()
      result = yield dd.commit_transaction("app_id", http_request)
      self.assertEquals(result, (datastore_pb.CommitResponse().Encode(), 0, ""))
  
    def test_rollback_transcation(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("notify_failed_transaction").and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      commit_request = datastore_pb.Transaction()
      commit_request.set_handle(123)
      commit_request.set_app("aaa")
      self.assertEquals(
        dd.rollback_transaction("app_id", commit_request.handle()), None)
  
    @staticmethod
    def get_new_entity_proto(app_id, kind, entity_name, prop_name, prop_value, ns=""):
      entity_proto = datastore_pb.EntityProto()
  
      reference = entity_proto.mutable_key()
      reference.set_app(app_id)
      reference.set_name_space(ns)
  
      path = reference.mutable_path()
      element = path.add_element()
      element.set_type(kind)
      element.set_name(entity_name)
  
      ent_group = entity_proto.mutable_entity_group()
      eg_element = ent_group.add_element()
      eg_element.set_type(kind)
      eg_element.set_name(entity_name)
  
      prop = entity_proto.add_property()
      prop.set_meaning(datastore_pb.Property.BYTESTRING)
      prop.set_name(prop_name)
      prop.set_multiple(1)
      val = prop.mutable_value()
      val.set_stringvalue(prop_value)
      return entity_proto
  
    @testing.gen_test
    def test_dynamic_put(self):
      db_batch = flexmock(session=flexmock())
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
      entity_proto2 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
      async_result = gen.Future()
      async_result.set_result({entity_key1: {}, entity_key2: {}})
  
      db_batch.should_receive('batch_get_entity').and_return(async_result)
      db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None,
        set_groups=lambda project, txid, groups: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'test': flexmock(indexes_pb=[])})
      putreq_pb = datastore_pb.PutRequest()
      putreq_pb.add_entity()
      putreq_pb.mutable_entity(0).MergeFrom(entity_proto1)
      putreq_pb.add_entity()
      putreq_pb.mutable_entity(1).MergeFrom(entity_proto2)
  
      putresp_pb = datastore_pb.PutResponse()
  
      async_true = gen.Future()
      async_true.set_result(True)
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire').and_return(async_true)
      entity_lock.should_receive('release')
  
      flexmock(ScatteredAllocator).should_receive('next').\
        and_return(random.randint(1, 500))
  
      yield dd.dynamic_put('test', putreq_pb, putresp_pb)
      self.assertEquals(len(putresp_pb.key_list()), 2)
  
    @testing.gen_test
    def test_put_entities(self):
      app_id = 'test'
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      entity_proto1 = self.get_new_entity_proto(
        app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
      entity_proto2 = self.get_new_entity_proto(
        app_id, "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
      entity_list = [entity_proto1, entity_proto2]
      async_result = gen.Future()
      async_result.set_result({entity_key1: {}, entity_key2: {}})
  
      db_batch.should_receive('batch_get_entity').and_return(async_result)
      db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None,
        set_groups=lambda project, txid, groups: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={app_id: flexmock(indexes_pb=[])})
  
      async_true = gen.Future()
      async_true.set_result(True)
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire').and_return(async_true)
      entity_lock.should_receive('release')
  
      yield dd.put_entities(app_id, entity_list)
  
    def test_acquire_locks_for_trans(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(False).once()
      self.assertRaises(TypeError, dd.acquire_locks_for_trans, [1], 1)
  
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(False).and_return(False)
      self.assertRaises(TypeError, dd.acquire_locks_for_trans, [1], 1)
  
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True)
  
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True).and_return(False)
      flexmock(dd).should_receive("get_table_prefix").and_return("prefix").never()
      flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").never()
      self.assertEquals({}, dd.acquire_locks_for_trans([], 1))
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("acquire_lock").once()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      entity = flexmock()
      entity.should_receive("app").and_return("appid")
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True).and_return(True)
      flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").once()
      self.assertEquals({'rootkey':1}, dd.acquire_locks_for_trans([entity], 1))
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("acquire_lock").once().and_raise(ZKTransactionException)
      zookeeper.should_receive("notify_failed_transaction").once()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      entity = flexmock()
      entity.should_receive("app").and_return("appid")
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True).and_return(True)
      flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").once()
      self.assertRaises(ZKTransactionException, dd.acquire_locks_for_trans, [entity], 1)
  
    @testing.gen_test
    def test_delete_entities(self):
      app_id = 'test'
      entity_proto1 = self.get_new_entity_proto(
        app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      row_key = "test\x00blah\x00test_kind:bob\x01"
      row_values = {row_key: {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                              APP_ENTITY_SCHEMA[1]: '1'}}
      async_result = gen.Future()
      async_result.set_result(row_values)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
      db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      row_keys = [entity_proto1.key()]
  
      yield dd.delete_entities(entity_proto1.key(), 1, row_keys)
  
    def test_release_put_locks_for_nontrans(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      entity_proto2 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entities = [entity_proto1, entity_proto2]
      dd.release_locks_for_nontrans(
        "test", entities, {
          'test\x00blah\x00test_kind:bob\x01': 1,
          'test\x00blah\x00test_kind:nancy\x01': 2
        })
  
    def test_root_key_from_entity_key(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      self.assertEquals(
        "test\x00blah\x00test_kind:bob\x01",
        dd.get_root_key_from_entity_key("test\x00blah\x00test_kind:bob\x01nancy\x01")
      )
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      self.assertEquals(
        "test\x00blah\x00test_kind:nancy\x01",
        dd.get_root_key_from_entity_key(entity_proto1.key())
      )
  
    @testing.gen_test
    def test_dynamic_get(self):
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      async_result = gen.Future()
      async_result.set_result({
        "test\x00blah\x00test_kind:nancy\x01": {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
      db_batch.should_receive('record_reads').and_return(ASYNC_NONE)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
      get_resp = datastore_pb.GetResponse()
  
      yield dd.dynamic_get("test", get_req, get_resp)
      self.assertEquals(get_resp.entity_size(), 1)
  
      # Now test while in a transaction
      get_resp = datastore_pb.GetResponse()
      get_req.mutable_transaction().set_handle(1)
      yield dd.dynamic_get("test", get_req, get_resp)
      self.assertEquals(get_resp.entity_size(), 1)
  
    @testing.gen_test
    def test_ancestor_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
      ancestor.MergeFrom(entity_key)
  
      async_result = gen.Future()
      async_result.set_result({
        "test\x00blah\x00test_kind:nancy\x01": {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
  
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('record_reads').and_return(ASYNC_NONE)
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
  
      entity_proto1 = {
        'test\x00blah\x00test_kind:nancy\x01': {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      }
      async_result_1 = gen.Future()
      async_result_1.set_result([entity_proto1, tombstone1])
      async_result_2 = gen.Future()
      async_result_2.set_result([])
  
      db_batch.should_receive("range_query").\
        and_return(async_result_1).\
        and_return(async_result_2)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      yield dd.ancestor_query(query, filter_info)
      # Now with a transaction
      transaction = query.mutable_transaction()
      transaction.set_handle(2)
      yield dd.ancestor_query(query, filter_info)
  
    @testing.gen_test
    def test_kindless_query(self):
      query = datastore_pb.Query()
      entity_proto1 = self.get_new_entity_proto(
        "test", "test_kind", "nancy", "prop1name", "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
  
      async_result = gen.Future()
      async_result.set_result({
        "test\x00blah\x00test_kind:nancy\x01": {
          APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      })
  
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(async_result)
  
      entity_proto1 = {
        'test\x00blah\x00test_kind:nancy\x01': {
          APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
          APP_ENTITY_SCHEMA[1]: 1
        }
      }
      async_result_1 = gen.Future()
      async_result_1.set_result([entity_proto1, tombstone1])
      async_result_2 = gen.Future()
      async_result_2.set_result([])
      db_batch.should_receive("range_query").\
        and_return(async_result_1).\
        and_return(async_result_2)
  
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      zookeeper.should_receive("acquire_lock").and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      filter_info = {
        '__key__' : [[0, 0]]
      }
      yield dd.kindless_query(query, filter_info)
  
    @testing.gen_test
    def test_dynamic_delete(self):
      async_true = gen.Future()
      async_true.set_result(True)
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire').and_return(async_true)
      entity_lock.should_receive('release')
  
      del_request = flexmock()
      del_request.should_receive("key_list")
      del_request.should_receive("has_transaction").never()
      del_request.should_receive("transaction").never()
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None,
        set_groups=lambda project_id, txid, groups: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'guestbook': flexmock(indexes_pb=[])})
      yield dd.dynamic_delete("appid", del_request)
  
      fake_key = entity_pb.Reference()
      fake_key.set_app('foo')
      path = fake_key.mutable_path()
      element = path.add_element()
      element.set_type('bar')
      element.set_id(1)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(True)
      transaction = flexmock()
      transaction.should_receive("handle").and_return(1)
      del_request.should_receive("transaction").and_return(transaction)
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'appid': flexmock(indexes_pb=[])})
      flexmock(utils).should_receive("get_entity_kind").and_return("kind")
      db_batch.should_receive('delete_entities_tx').and_return(ASYNC_NONE)
      yield dd.dynamic_delete("appid", del_request)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(False)
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.index_manager = flexmock(
        projects={'appid': flexmock(indexes_pb=[])})
      flexmock(dd).should_receive("delete_entities").and_return(ASYNC_NONE).once()
      yield dd.dynamic_delete("appid", del_request)
  
    def test_reverse_path(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      key = "Project:Synapse\x01Module:Core\x01"
      self.assertEquals(dd.reverse_path(key), "Module:Core\x01Project:Synapse\x01")
  
    def test_remove_exists_filters(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      self.assertEquals(dd.remove_exists_filters({}), {})
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]
      }
      self.assertEquals(dd.remove_exists_filters(filter_info), filter_info)
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EXISTS, "1")],
        "prop2": [(datastore_pb.Query_Filter.EXISTS, "2")]
      }
      self.assertEquals(dd.remove_exists_filters(filter_info), {})
  
    def test_is_zigzag_merge_join(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      query = datastore_pb.Query()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      self.assertEquals(dd.is_zigzag_merge_join(query, {}, {}), False)
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]
      }
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), True)
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")]
      }
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), False)
  
    @testing.gen_test
    def test_zigzag_merge_join(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
  
      query = datastore_pb.Query()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      flexmock(dd).should_receive("is_zigzag_merge_join").and_return(False)
      result = yield dd.zigzag_merge_join(None, None, None)
      self.assertEquals(result, None)
  
      filter_info = {
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]
      }
      result = yield dd.zigzag_merge_join(query, filter_info, [])
      self.assertEquals(result, None)
  
    def test_index_deletions(self):
      old_entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # No deletions should occur when the entity doesn't change.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      self.assertListEqual([], index_deletions(old_entity, old_entity))
  
      # When a property changes, the previous index entries should be deleted.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(old_entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
  
      deletions = index_deletions(old_entity, new_entity)
      self.assertEqual(len(deletions), 2)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = old_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # When given an index, an entry should be removed from the composite table.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 3)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.COMPOSITE_TABLE)
  
      # No composite deletions should occur when the entity type differs.
      definition.set_entity_type('TestEntity')
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 2)
  
    def test_deletions_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # Deleting an entity with one property should remove four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      deletions = deletions_for_entity(entity)
      self.assertEqual(len(deletions), 4)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.APP_KIND_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # Deleting an entity with two properties and one composite index should
      # remove seven entries.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = deletions_for_entity(entity, (composite_index,))
      self.assertEqual(len(deletions), 7)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[4]['table'], dbconstants.COMPOSITE_TABLE)
      self.assertEqual(deletions[5]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[6]['table'], dbconstants.APP_KIND_TABLE)
  
    def test_mutations_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
      txn = 1
  
      # Adding an entity with one property should add four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version_sync').and_return(True)
      mutations = mutations_for_entity(entity, txn)
      self.assertEqual(len(mutations), 4)
      self.assertEqual(mutations[0]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[1]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[2]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      # Updating an entity with one property should delete two entries and add
      # four more.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
      mutations = mutations_for_entity(entity, txn, new_entity)
      self.assertEqual(len(mutations), 6)
      self.assertEqual(mutations[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[0]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[1]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[4]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[5]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop
  
  [... the rest is omitted ...]
files.old_path: AppDB/test/unit/test_datastore_server.py