files.new_contents: import array
  import datetime
  import itertools
  import logging
  import md5
  import random
  import sys
  import threading
  import time
  
  import dbconstants
  import helper_functions
  
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from kazoo.client import KazooState
  from .dbconstants import APP_ENTITY_SCHEMA
  from .dbconstants import ID_KEY_LENGTH
  from .dbconstants import MAX_TX_DURATION
  from .cassandra_env import cassandra_interface
  from .cassandra_env.entity_id_allocator import EntityIDAllocator
  from .cassandra_env.entity_id_allocator import ScatteredAllocator
  from .cassandra_env.utils import deletions_for_entity
  from .cassandra_env.utils import mutations_for_entity
  from .utils import clean_app_id
  from .utils import encode_entity_table_key
  from .utils import encode_index_pb
  from .utils import encode_path_from_filter
  from .utils import get_composite_index_keys
  from .utils import get_entity_key
  from .utils import get_entity_kind
  from .utils import get_index_key_from_params
  from .utils import get_kind_key
  from .utils import group_for_key
  from .utils import reference_property_to_reference
  from .utils import UnprocessedQueryCursor
  from .zkappscale import entity_lock
  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 import datastore_errors
  from google.appengine.api.datastore_distributed import _MAX_ACTIONS_PER_TXN
  from google.appengine.api.taskqueue import taskqueue_service_pb
  from google.appengine.api.taskqueue.taskqueue_distributed import\
    TaskQueueServiceStub
  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
  
  
  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
  
    # The number of entries looked at when doing a composite query
    # It will keep looking at this size window when getting the result
    _MAX_COMPOSITE_WINDOW = 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):
      """
         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()
  
      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 stub object for each project using transactional tasks.
      self.taskqueue_stubs = {}
  
      # Maintain a scattered allocator for each project.
      self.scattered_allocators = {}
  
      # Maintain a sequential allocator for each project.
      self.sequential_allocators = {}
  
      self.transaction_manager = transaction_manager
      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:
          logging.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
    
    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}
  
      self.datastore_batch.batch_put_entity(dbconstants.COMPOSITE_TABLE, 
                                            row_keys, 
                                            dbconstants.COMPOSITE_SCHEMA,
                                            row_values)
  
    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))
      index_keys = []
      composite_id = str(index.id())
      index_keys.append(self._SEPARATOR.join([app_id, 'index', composite_id]))
      self.datastore_batch.batch_delete(dbconstants.METADATA_TABLE,
                                        index_keys, 
                                        column_names=dbconstants.METADATA_TABLE)
  
    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.
      """
      # Generate a random number based on time of creation.
      rand = int(str(int(time.time())) + str(random.randint(0, 999999)))
      index.set_id(rand)
      encoded_entity = index.Encode()
      row_key = self._SEPARATOR.join([app_id, 'index', str(rand)])
      row_keys = [row_key]
      row_values = {}
      row_values[row_key] = {dbconstants.METADATA_SCHEMA[0]: encoded_entity}
      self.datastore_batch.batch_put_entity(dbconstants.METADATA_TABLE, 
                                            row_keys, 
                                            dbconstants.METADATA_SCHEMA, 
                                            row_values)    
      return rand 
  
    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 = 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 = self.__fetch_entities(references)
        entities = [entity_pb.EntityProto(entity) for entity in pb_entities]
  
        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))
  
    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]
      return allocator.allocate_size(size)
  
    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]
      return allocator.allocate_max(max_id)
  
    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]
  
        allocator.set_min_counter(counter)
  
    def put_entities(self, app, entities, composite_indexes=()):
      """ Updates indexes of existing entities, inserts new entities and 
          indexes for them.
  
      Args:
        app: A string containing the application ID.
        entities: List of entities.
        composite_indexes: A list or tuple of CompositeIndex objects.
      """
      self.logger.debug('Inserting {} entities'.format(len(entities)))
      entity_keys = []
      for entity in entities:
        prefix = self.get_table_prefix(entity)
        entity_keys.append(get_entity_key(prefix, entity.key().path()))
  
      current_values = self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, entity_keys, APP_ENTITY_SCHEMA)
  
      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])
        lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
        try:
          with lock:
            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})
            self.datastore_batch.batch_mutate(app, batch, entity_changes, txid)
        finally:
          self.transaction_manager.delete_transaction_id(app, txid)
  
    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 = 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})
  
        self.datastore_batch._normal_batch(batch, txid)
  
    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():
          last_path.set_id(allocator.next())
          group = entity.mutable_entity_group()
          root = entity.key().path().element(0)
          group.add_element().CopyFrom(root)
  
      if put_request.has_transaction():
        self.datastore_batch.put_entities_tx(
          app_id, put_request.transaction().handle(), entities)
      else:
        try:
          self.put_entities(app_id, entities, put_request.composite_index_list())
        except entity_lock.LockTimeout as timeout_error:
          raise dbconstants.AppScaleDBConnectionError(str(timeout_error))
        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)
  
    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 = self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE,
        row_keys,
        APP_ENTITY_SCHEMA)
      return result, row_keys
  
    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 = self.fetch_keys(keys)
        fetched_groups = {group_for_key(key).Encode() for key in keys}
        self.datastore_batch.record_reads(
          app_id, get_request.transaction().handle(), fetched_groups)
      else:
        results, row_keys = 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))
  
    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())
  
      # We use the marked changes field to signify if we should 
      # look up composite indexes because delete request do not
      # include that information.
      composite_indexes = []
      filtered_indexes = []
      if delete_request.has_mark_changes():
        all_composite_indexes = self.datastore_batch.get_indices(app_id)
        for index in all_composite_indexes:
          new_index = entity_pb.CompositeIndex()
          new_index.ParseFromString(index)
          composite_indexes.append(new_index)
        # Only get composites of the correct kinds.
        for index in composite_indexes:
          if index.definition().entity_type() in ent_kinds:
            filtered_indexes.append(index)
  
      if delete_request.has_transaction():
        txid = delete_request.transaction().handle()
        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])
          lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
          try:
            with lock:
              self.delete_entities(
                group_key,
                txid,
                key_list,
                composite_indexes=filtered_indexes
              )
            self.logger.debug('Removed {} entities'.format(len(key_list)))
          except entity_lock.LockTimeout as timeout_error:
            raise dbconstants.AppScaleDBConnectionError(str(timeout_error))
          finally:
            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
  
    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:
        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 = 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]
      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
  
    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 = 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]])
      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
  
    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)
      return self.__fetch_entities_from_row_list(rowkeys)
  
    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)
      return self.__fetch_entities_dict_from_row_list(rowkeys)
  
    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 = 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]]
  
      return clean_results
  
    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:
          return results[:limit]
  
        entities = 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:
              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
          added_padding = True
  
    def __extract_entities(self, kv):
      """ Given a result from a range query on the Entity table return a 
          list of encoded entities.
  
      Args:
        kv: Key and values from a range query on the entity table.
      Returns:
        The extracted entities.
      """
      keys = [item.keys()[0] for item in kv]
      results = []    
      for index, entity in enumerate(kv):
        key = keys[index]
        entity = entity[key][APP_ENTITY_SCHEMA[0]]
        results.append(entity)
  
      return results
  
    def ordered_ancestor_query(self, query, filter_info, order_info):
      """ Performs an ordered ancestor query. It grabs all entities of a 
          given ancestor and then orders in memory.
      
      Args:
        query: The query to run.
        filter_info: Tuple with filter operators and values
        order_info: Tuple with property name and the sort order.
      Returns:
        A list of entities.
      Raises:
        ZKTransactionException: If a lock could not be acquired.
      """ 
      ancestor = query.ancestor()
      prefix = self.get_table_prefix(query)
      path = buffer(prefix + self._SEPARATOR) + encode_index_pb(ancestor.path())
      txn_id = 0
      if query.has_transaction():
        txn_id = query.transaction().handle()   
  
      startrow = path
      endrow = path + self._TERM_STRING
      end_inclusive = self._ENABLE_INCLUSIVITY
      start_inclusive = self._ENABLE_INCLUSIVITY
      if query.has_compiled_cursor() and query.compiled_cursor().position_size():
        cursor = appscale_stub_util.ListCursor(query)
        last_result = cursor._GetLastResult()
        startrow = self.__get_start_key(prefix, None, None, last_result)
        start_inclusive = self._DISABLE_INCLUSIVITY
        if query.compiled_cursor().position_list()[0].start_inclusive() == 1:
          start_inclusive = self._ENABLE_INCLUSIVITY
  
      limit =
  
  [... 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 random
  import sys
  import threading
  import time
  
  import dbconstants
  import helper_functions
  
  from appscale.common.unpackaged import APPSCALE_PYTHON_APPSERVER
  from kazoo.client import KazooState
  from .dbconstants import APP_ENTITY_SCHEMA
  from .dbconstants import ID_KEY_LENGTH
  from .dbconstants import MAX_TX_DURATION
  from .cassandra_env import cassandra_interface
  from .cassandra_env.entity_id_allocator import EntityIDAllocator
  from .cassandra_env.entity_id_allocator import ScatteredAllocator
  from .cassandra_env.utils import deletions_for_entity
  from .cassandra_env.utils import mutations_for_entity
  from .utils import clean_app_id
  from .utils import encode_entity_table_key
  from .utils import encode_index_pb
  from .utils import encode_path_from_filter
  from .utils import get_composite_index_keys
  from .utils import get_entity_key
  from .utils import get_entity_kind
  from .utils import get_index_key_from_params
  from .utils import get_kind_key
  from .utils import group_for_key
  from .utils import reference_property_to_reference
  from .utils import UnprocessedQueryCursor
  from .zkappscale import entity_lock
  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 import datastore_errors
  from google.appengine.api.datastore_distributed import _MAX_ACTIONS_PER_TXN
  from google.appengine.api.taskqueue import taskqueue_service_pb
  from google.appengine.api.taskqueue.taskqueue_distributed import\
    TaskQueueServiceStub
  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
  
  
  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
  
    # The number of entries looked at when doing a composite query
    # It will keep looking at this size window when getting the result
    _MAX_COMPOSITE_WINDOW = 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):
      """
         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()
  
      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 stub object for each project using transactional tasks.
      self.taskqueue_stubs = {}
  
      # Maintain a scattered allocator for each project.
      self.scattered_allocators = {}
  
      # Maintain a sequential allocator for each project.
      self.sequential_allocators = {}
  
      self.transaction_manager = transaction_manager
      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:
          logging.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
    
    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}
  
      self.datastore_batch.batch_put_entity(dbconstants.COMPOSITE_TABLE, 
                                            row_keys, 
                                            dbconstants.COMPOSITE_SCHEMA,
                                            row_values)
  
    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))
      index_keys = []
      composite_id = str(index.id())
      index_keys.append(self._SEPARATOR.join([app_id, 'index', composite_id]))
      self.datastore_batch.batch_delete(dbconstants.METADATA_TABLE,
                                        index_keys, 
                                        column_names=dbconstants.METADATA_TABLE)
  
    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.
      """
      # Generate a random number based on time of creation.
      rand = int(str(int(time.time())) + str(random.randint(0, 999999)))
      index.set_id(rand)
      encoded_entity = index.Encode()
      row_key = self._SEPARATOR.join([app_id, 'index', str(rand)])
      row_keys = [row_key]
      row_values = {}
      row_values[row_key] = {dbconstants.METADATA_SCHEMA[0]: encoded_entity}
      self.datastore_batch.batch_put_entity(dbconstants.METADATA_TABLE, 
                                            row_keys, 
                                            dbconstants.METADATA_SCHEMA, 
                                            row_values)    
      return rand 
  
    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 = 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 = self.__fetch_entities(references)
        entities = [entity_pb.EntityProto(entity) for entity in pb_entities]
  
        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))
  
    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]
      return allocator.allocate_size(size)
  
    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]
      return allocator.allocate_max(max_id)
  
    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]
  
        allocator.set_min_counter(counter)
  
    def put_entities(self, app, entities, composite_indexes=()):
      """ Updates indexes of existing entities, inserts new entities and 
          indexes for them.
  
      Args:
        app: A string containing the application ID.
        entities: List of entities.
        composite_indexes: A list or tuple of CompositeIndex objects.
      """
      self.logger.debug('Inserting {} entities'.format(len(entities)))
      entity_keys = []
      for entity in entities:
        prefix = self.get_table_prefix(entity)
        entity_keys.append(get_entity_key(prefix, entity.key().path()))
  
      current_values = self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE, entity_keys, APP_ENTITY_SCHEMA)
  
      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)
        lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
        try:
          with lock:
            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})
            self.datastore_batch.batch_mutate(app, batch, entity_changes, txid)
        finally:
          self.transaction_manager.delete_transaction_id(app, txid)
  
    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 = 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})
  
        self.datastore_batch._normal_batch(batch, txid)
  
    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():
          last_path.set_id(allocator.next())
          group = entity.mutable_entity_group()
          root = entity.key().path().element(0)
          group.add_element().CopyFrom(root)
  
      if put_request.has_transaction():
        self.datastore_batch.put_entities_tx(
          app_id, put_request.transaction().handle(), entities)
      else:
        try:
          self.put_entities(app_id, entities, put_request.composite_index_list())
        except entity_lock.LockTimeout as timeout_error:
          raise dbconstants.AppScaleDBConnectionError(str(timeout_error))
        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)
  
    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 = self.datastore_batch.batch_get_entity(
        dbconstants.APP_ENTITY_TABLE,
        row_keys,
        APP_ENTITY_SCHEMA)
      return result, row_keys
  
    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 = self.fetch_keys(keys)
        fetched_groups = {group_for_key(key).Encode() for key in keys}
        self.datastore_batch.record_reads(
          app_id, get_request.transaction().handle(), fetched_groups)
      else:
        results, row_keys = 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))
  
    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())
  
      # We use the marked changes field to signify if we should 
      # look up composite indexes because delete request do not
      # include that information.
      composite_indexes = []
      filtered_indexes = []
      if delete_request.has_mark_changes():
        all_composite_indexes = self.datastore_batch.get_indices(app_id)
        for index in all_composite_indexes:
          new_index = entity_pb.CompositeIndex()
          new_index.ParseFromString(index)
          composite_indexes.append(new_index)
        # Only get composites of the correct kinds.
        for index in composite_indexes:
          if index.definition().entity_type() in ent_kinds:
            filtered_indexes.append(index)
  
      if delete_request.has_transaction():
        txid = delete_request.transaction().handle()
        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)
          lock = entity_lock.EntityLock(self.zookeeper.handle, [group_key], txid)
          try:
            with lock:
              self.delete_entities(
                group_key,
                txid,
                key_list,
                composite_indexes=filtered_indexes
              )
            self.logger.debug('Removed {} entities'.format(len(key_list)))
          except entity_lock.LockTimeout as timeout_error:
            raise dbconstants.AppScaleDBConnectionError(str(timeout_error))
          finally:
            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
  
    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:
        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 = 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]
      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
  
    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 = 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]])
      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
  
    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)
      return self.__fetch_entities_from_row_list(rowkeys)
  
    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)
      return self.__fetch_entities_dict_from_row_list(rowkeys)
  
    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 = 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]]
  
      return clean_results
  
    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:
          return results[:limit]
  
        entities = 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:
              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
          added_padding = True
  
    def __extract_entities(self, kv):
      """ Given a result from a range query on the Entity table return a 
          list of encoded entities.
  
      Args:
        kv: Key and values from a range query on the entity table.
      Returns:
        The extracted entities.
      """
      keys = [item.keys()[0] for item in kv]
      results = []    
      for index, entity in enumerate(kv):
        key = keys[index]
        entity = entity[key][APP_ENTITY_SCHEMA[0]]
        results.append(entity)
  
      return results
  
    def ordered_ancestor_query(self, query, filter_info, order_info):
      """ Performs an ordered ancestor query. It grabs all entities of a 
          given ancestor and then orders in memory.
      
      Args:
        query: The query to run.
        filter_info: Tuple with filter operators and values
        order_info: Tuple with property name and the sort order.
      Returns:
        A list of entities.
      Raises:
        ZKTransactionException: If a lock could not be acquired.
      """ 
      ancestor = query.ancestor()
      prefix = self.get_table_prefix(query)
      path = buffer(prefix + self._SEPARATOR) + encode_index_pb(ancestor.path())
      txn_id = 0
      if query.has_transaction():
        txn_id = query.transaction().handle()   
  
      startrow = path
      endrow = path + self._TERM_STRING
      end_inclusive = self._ENABLE_INCLUSIVITY
      start_inclusive = self._ENABLE_INCLUSIVITY
      if query.has_compiled_cursor() and query.compiled_cursor().position_size():
        cursor = appscale_stub_util.ListCursor(query)
        last_result = cursor._GetLastResult()
        startrow = self.__get_start_key(prefix, None, None, last_result)
        start_inclusive = self._DISABLE_INCLUSIVITY
        if query.compiled_cursor().position_list()[0].start_inclusive() == 1:
          start_inclusive = self._ENABLE_INCLUSIVITY
  
      limit = self._MAXIMUM_RESULTS
  
      unordered = self.fetch_from_entity_table(
        startrow, endrow, limit, 0, start_inclusive, end
  
  [... the rest is omitted ...]
files.old_path: AppDB/appscale/datastore/datastore_distributed.py
files.new_contents: """ Generates and keeps track of transaction IDs. """
  from __future__ import division
  
  import json
  import logging
  import time
  
  from kazoo.exceptions import KazooException
  from kazoo.exceptions import NodeExistsError
  from kazoo.exceptions import NotEmptyError
  from tornado.ioloop import IOLoop
  
  from appscale.common.async_retrying import retry_children_watch_coroutine
  from .constants import CONTAINER_PREFIX
  from .constants import COUNTER_NODE_PREFIX
  from .constants import MAX_SEQUENCE_COUNTER
  from .constants import OFFSET_NODE
  from .entity_lock import zk_group_path
  from ..dbconstants import BadRequest
  from ..dbconstants import InternalError
  
  logger = logging.getLogger('appscale-datastore')
  
  # Nodes that indicate a cross-group transaction start with this string.
  XG_PREFIX = 'xg'
  
  
  class ProjectTransactionManager(object):
    """ Generates and keeps track of transaction IDs for a project. """
    def __init__(self, project_id, zk_client):
      """ Creates a new ProjectTransactionManager.
  
      Args:
        project_id: A string specifying a project ID.
      """
      self.project_id = project_id
      self.zk_client = zk_client
  
      self._project_node = '/appscale/apps/{}'.format(self.project_id)
  
      # Allows users to manually modify transaction IDs after a binary migration.
      self._txid_manual_offset = 0
      self._offset_node = '/'.join([self._project_node, OFFSET_NODE])
      self.zk_client.DataWatch(self._offset_node, self._update_offset)
  
      # Allows the manager to use a new container after one is exhausted.
      self._txid_automatic_offset = 0
      self._counter_path = '/'.join([self._project_node, CONTAINER_PREFIX])
      self.zk_client.ensure_path(self._project_node)
      self.zk_client.ensure_path(self._counter_path)
      self.zk_client.ChildrenWatch(self._project_node, self._update_project)
  
      # Containers that do not need to be checked for open transactions.
      self._inactive_containers = set()
  
    def create_transaction_id(self, xg):
      """ Generates a new transaction ID.
  
      Args:
        xg: A boolean indicating a cross-group transaction.
      Returns:
        An integer specifying the created transaction ID.
      Raises:
        InternalError if unable to create a new transaction ID.
      """
      current_time = time.time()
      counter_path_prefix = '/'.join([self._counter_path, COUNTER_NODE_PREFIX])
      try:
        new_path = self.zk_client.create(
          counter_path_prefix, value=str(current_time), sequence=True)
      except KazooException:
        message = 'Unable to create new transaction ID'
        logger.exception(message)
        raise InternalError(message)
  
      counter = int(new_path.split('/')[-1].lstrip(COUNTER_NODE_PREFIX))
  
      if counter < 0:
        logger.debug('Removing invalid counter')
        self._delete_counter(new_path)
        self._update_auto_offset()
        return self.create_transaction_id(xg)
  
      txid = self._txid_manual_offset + self._txid_automatic_offset + counter
  
      if txid == 0:
        self._delete_counter(new_path)
        return self.create_transaction_id(xg)
  
      if xg:
        xg_path = '/'.join([new_path, XG_PREFIX])
        try:
          self.zk_client.create(xg_path, value=str(current_time))
        except KazooException:
          message = 'Unable to create new cross-group transaction ID'
          logger.exception(message)
          raise InternalError(message)
  
      self._last_txid_created = txid
      return txid
  
    def delete_transaction_id(self, txid):
      """ Removes a transaction ID from the list of active transactions.
  
      Args:
        txid: An integer specifying a transaction ID.
      """
      path = self._txid_to_path(txid)
      self._delete_counter(path)
  
    def get_open_transactions(self):
      """ Fetches a list of active transactions.
  
      Returns:
        A list of integers specifying transaction IDs.
      Raises:
        InternalError if unable to fetch list of transaction IDs.
      """
      txids = []
      active_containers = self._active_containers()
      for index, container in enumerate(active_containers):
        container_name = container.split('/')[-1]
        container_count = int(container_name[len(CONTAINER_PREFIX):] or 1)
        container_size = MAX_SEQUENCE_COUNTER + 1
        auto_offset = (container_count - 1) * container_size
        offset = self._txid_manual_offset + auto_offset
  
        try:
          paths = self.zk_client.get_children(container)
        except KazooException:
          message = 'Unable to fetch list of counters'
          logger.exception(message)
          raise InternalError(message)
  
        counter_nodes = [path.split('/')[-1] for path in paths]
        txids.extend([offset + int(node.lstrip(COUNTER_NODE_PREFIX))
                      for node in counter_nodes])
  
        # If there are no counters left in an old container, mark it inactive.
        if not counter_nodes and index < len(active_containers) - 1:
          self._inactive_containers.add(container_name)
  
      return txids
  
    def set_groups(self, txid, groups):
      """ Defines which groups will be involved in a transaction.
  
      Args:
        txid: An integer specifying a transaction ID.
        groups: An iterable of entity group Reference objects.
      """
      txid_path = self._txid_to_path(txid)
      groups_path = '/'.join([txid_path, 'groups'])
      encoded_groups = [zk_group_path(group) for group in groups]
      try:
        self.zk_client.create(groups_path, value=json.dumps(encoded_groups))
      except KazooException:
        message = 'Unable to set lock list for transaction'
        logger.exception(message)
        raise InternalError(message)
  
    def _delete_counter(self, path):
      """ Removes a counter node.
  
      Args:
        path: A string specifying a ZooKeeper path.
      """
      try:
        self.zk_client.delete(path, recursive=True)
      except KazooException:
        # Let the transaction groomer clean it up.
        logger.exception('Unable to delete counter')
  
    def _active_containers(self):
      """ Determines the containers that need to be checked for transactions.
  
      Returns:
        A tuple of strings specifying ZooKeeper paths.
      """
      container_name = self._counter_path.split('/')[-1]
      container_count = int(container_name[len(CONTAINER_PREFIX):] or 1)
  
      all_containers = [CONTAINER_PREFIX + str(index + 1)
                        for index in range(container_count)]
      all_containers[0] = CONTAINER_PREFIX
  
      return tuple('/'.join([self._project_node, container])
                   for container in all_containers
                   if container not in self._inactive_containers)
  
    def _txid_to_path(self, txid):
      """ Determines the ZooKeeper path for a given transaction ID.
  
      Args:
        txid: An integer specifying a transaction ID.
      Returns:
        A strings specifying the transaction's ZooKeeper path.
      """
      corrected_counter = txid - self._txid_manual_offset
  
      # The number of counters a container can store (including 0).
      container_size = MAX_SEQUENCE_COUNTER + 1
  
      container_count = int(corrected_counter / container_size) + 1
      container_suffix = '' if container_count == 1 else str(container_count)
      container_name = CONTAINER_PREFIX + container_suffix
      container_path = '/'.join([self._project_node, container_name])
  
      counter_value = corrected_counter % container_size
      node_name = COUNTER_NODE_PREFIX + str(counter_value).zfill(10)
      return '/'.join([container_path, node_name])
  
    def _update_auto_offset(self):
      """ Ensures there is a usable sequence container. """
      container_name = self._counter_path.split('/')[-1]
      container_count = int(container_name[len(CONTAINER_PREFIX):] or 1)
      next_node = CONTAINER_PREFIX + str(container_count + 1)
      next_path = '/'.join([self._project_node, next_node])
  
      try:
        self.zk_client.create(next_path)
      except NodeExistsError:
        # Another process may have already created the new counter.
        pass
      except KazooException:
        message = 'Unable to create transaction ID counter'
        logger.exception(message)
        raise InternalError(message)
  
      try:
        node_list = self.zk_client.get_children(self._project_node)
      except KazooException:
        message = 'Unable to find transaction ID counter'
        logger.exception(message)
        raise InternalError(message)
  
      self._update_project_sync(node_list)
  
    def _update_offset(self, new_offset, _):
      """ Watches for updates to the manual offset node. """
      # This assignment is atomic, so it does not need to happen in the IOLoop.
      self._txid_manual_offset = int(new_offset or 0)
  
    def _update_project_sync(self, node_list):
      """ Updates the record of usable sequence containers. """
      counters = [int(node[len(CONTAINER_PREFIX):] or 1)
                  for node in node_list if node.startswith(CONTAINER_PREFIX)]
      counters.sort()
  
      container_suffix = '' if len(counters) == 1 else str(counters[-1])
      latest_node = CONTAINER_PREFIX + container_suffix
  
      self._counter_path = '/'.join([self._project_node, latest_node])
  
      # The number of counters a container can store (including 0).
      container_size = MAX_SEQUENCE_COUNTER + 1
      self._txid_automatic_offset = (len(counters) - 1) * container_size
  
    def _update_project(self, node_list):
      """ Watches for updates to the list of containers. """
      IOLoop.instance().add_callback(self._update_project_sync, node_list)
  
  
  class TransactionManager(object):
    """ Generates and keeps track of transaction IDs. """
    def __init__(self, zk_client):
      """ Creates a new TransactionManager.
  
      Args:
        zk_client: A KazooClient.
      """
      self.zk_client = zk_client
      self.zk_client.ensure_path('/appscale/projects')
      self.projects = {}
  
      # Since this manager can be used synchronously, ensure that the projects
      # are populated for this IOLoop iteration.
      project_ids = self.zk_client.get_children('/appscale/projects')
      self._update_projects_sync(project_ids)
  
      self.zk_client.ChildrenWatch('/appscale/projects', self._update_projects)
  
    def create_transaction_id(self, project_id, xg=False):
      """ Generates a new transaction ID.
  
      Args:
        project_id: A string specifying a project ID.
        xg: A boolean indicating a cross-group transaction.
      Returns:
        An integer specifying the created transaction ID.
      Raises:
        BadRequest if the project does not exist.
        InternalError if unable to create a new transaction ID.
      """
      try:
        project_tx_manager = self.projects[project_id]
      except KeyError:
        raise BadRequest('The project {} was not found'.format(project_id))
  
      return project_tx_manager.create_transaction_id(xg)
  
    def delete_transaction_id(self, project_id, txid):
      """ Removes a transaction ID from the list of active transactions.
  
      Args:
        project_id: A string specifying a project ID.
        txid: An integer specifying a transaction ID.
      Raises:
        BadRequest if the project does not exist.
      """
      try:
        project_tx_manager = self.projects[project_id]
      except KeyError:
        raise BadRequest('The project {} was not found'.format(project_id))
  
      return project_tx_manager.delete_transaction_id(txid)
  
    def get_open_transactions(self, project_id):
      """ Fetch a list of open transactions for a given project.
  
      Args:
        project_id: A string specifying a project ID.
      Returns:
        A list of integers specifying transaction IDs.
      Raises:
        BadRequest if the project does not exist.
        InternalError if unable to fetch list of open transactions.
      """
      try:
        project_tx_manager = self.projects[project_id]
      except KeyError:
        raise BadRequest('The project {} was not found'.format(project_id))
  
      return project_tx_manager.get_open_transactions()
  
    def set_groups(self, project_id, txid, groups):
      """ Defines which groups will be involved in a transaction.
  
      Args:
        project_id: A string specifying a project ID.
        txid: An integer specifying a transaction ID.
        groups: An iterable of entity group Reference objects.
      """
      try:
        project_tx_manager = self.projects[project_id]
      except KeyError:
        raise BadRequest('The project {} was not found'.format(project_id))
  
      return project_tx_manager.set_groups(txid, groups)
  
    def _update_projects_sync(self, new_project_ids):
      """ Updates the available projects for starting transactions.
  
      Args:
        new_project_ids: A list of strings specifying current project IDs.
      """
      for project_id in new_project_ids:
        if project_id not in self.projects:
          self.projects[project_id] = ProjectTransactionManager(project_id,
                                                                self.zk_client)
  
      for project_id in self.projects.keys():
        if project_id not in new_project_ids:
          del self.projects[project_id]
  
    def _update_projects(self, project_ids):
      """ Watches for changes to list of existing projects.
  
      Args:
        project_ids: A list of strings specifying current project IDs.
      """
      persistent_setup_projects = retry_children_watch_coroutine(
        '/appscale/projects', self._update_projects_sync)
      IOLoop.instance().add_callback(persistent_setup_projects, project_ids)
  
  
  [... the rest is omitted ...]
files.new_path: AppDB/appscale/datastore/zkappscale/transaction_manager.py
files.old_contents: """ Generates and keeps track of transaction IDs. """
  from __future__ import division
  
  import logging
  import time
  
  from kazoo.exceptions import KazooException
  from kazoo.exceptions import NodeExistsError
  from kazoo.exceptions import NotEmptyError
  from tornado.ioloop import IOLoop
  
  from appscale.common.async_retrying import retry_children_watch_coroutine
  from .constants import CONTAINER_PREFIX
  from .constants import COUNTER_NODE_PREFIX
  from .constants import MAX_SEQUENCE_COUNTER
  from .constants import OFFSET_NODE
  from ..dbconstants import BadRequest
  from ..dbconstants import InternalError
  
  logger = logging.getLogger('appscale-datastore')
  
  # Nodes that indicate a cross-group transaction start with this string.
  XG_PREFIX = 'xg'
  
  
  class ProjectTransactionManager(object):
    """ Generates and keeps track of transaction IDs for a project. """
    def __init__(self, project_id, zk_client):
      """ Creates a new ProjectTransactionManager.
  
      Args:
        project_id: A string specifying a project ID.
      """
      self.project_id = project_id
      self.zk_client = zk_client
  
      self._project_node = '/appscale/apps/{}'.format(self.project_id)
  
      # Allows users to manually modify transaction IDs after a binary migration.
      self._txid_manual_offset = 0
      self._offset_node = '/'.join([self._project_node, OFFSET_NODE])
      self.zk_client.DataWatch(self._offset_node, self._update_offset)
  
      # Allows the manager to use a new container after one is exhausted.
      self._txid_automatic_offset = 0
      self._counter_path = '/'.join([self._project_node, CONTAINER_PREFIX])
      self.zk_client.ensure_path(self._project_node)
      self.zk_client.ensure_path(self._counter_path)
      self.zk_client.ChildrenWatch(self._project_node, self._update_project)
  
      # Containers that do not need to be checked for open transactions.
      self._inactive_containers = set()
  
    def create_transaction_id(self, xg):
      """ Generates a new transaction ID.
  
      Args:
        xg: A boolean indicating a cross-group transaction.
      Returns:
        An integer specifying the created transaction ID.
      Raises:
        InternalError if unable to create a new transaction ID.
      """
      current_time = time.time()
      counter_path_prefix = '/'.join([self._counter_path, COUNTER_NODE_PREFIX])
      try:
        new_path = self.zk_client.create(
          counter_path_prefix, value=str(current_time), sequence=True)
      except KazooException:
        message = 'Unable to create new transaction ID'
        logger.exception(message)
        raise InternalError(message)
  
      counter = int(new_path.split('/')[-1].lstrip(COUNTER_NODE_PREFIX))
  
      if counter < 0:
        logger.debug('Removing invalid counter')
        self._delete_counter(new_path)
        self._update_auto_offset()
        return self.create_transaction_id(xg)
  
      txid = self._txid_manual_offset + self._txid_automatic_offset + counter
  
      if txid == 0:
        self._delete_counter(new_path)
        return self.create_transaction_id(xg)
  
      if xg:
        xg_path = '/'.join([new_path, XG_PREFIX])
        try:
          self.zk_client.create(xg_path, value=str(current_time))
        except KazooException:
          message = 'Unable to create new cross-group transaction ID'
          logger.exception(message)
          raise InternalError(message)
  
      self._last_txid_created = txid
      return txid
  
    def delete_transaction_id(self, txid):
      """ Removes a transaction ID from the list of active transactions.
  
      Args:
        txid: An integer specifying a transaction ID.
      """
      corrected_counter = txid - self._txid_manual_offset
  
      # The number of counters a container can store (including 0).
      container_size = MAX_SEQUENCE_COUNTER + 1
  
      container_count = int(corrected_counter / container_size) + 1
      container_suffix = '' if container_count == 1 else str(container_count)
      container_name = CONTAINER_PREFIX + container_suffix
      container_path = '/'.join([self._project_node, container_name])
  
      counter_value = corrected_counter % container_size
      node_name = COUNTER_NODE_PREFIX + str(counter_value).zfill(10)
      full_path = '/'.join([container_path, node_name])
      self._delete_counter(full_path)
  
    def get_open_transactions(self):
      """ Fetches a list of active transactions.
  
      Returns:
        A list of integers specifying transaction IDs.
      Raises:
        InternalError if unable to fetch list of transaction IDs.
      """
      txids = []
      active_containers = self._active_containers()
      for index, container in enumerate(active_containers):
        container_name = container.split('/')[-1]
        container_count = int(container_name[len(CONTAINER_PREFIX):] or 1)
        container_size = MAX_SEQUENCE_COUNTER + 1
        auto_offset = (container_count - 1) * container_size
        offset = self._txid_manual_offset + auto_offset
  
        try:
          paths = self.zk_client.get_children(container)
        except KazooException:
          message = 'Unable to fetch list of counters'
          logger.exception(message)
          raise InternalError(message)
  
        counter_nodes = [path.split('/')[-1] for path in paths]
        txids.extend([offset + int(node.lstrip(COUNTER_NODE_PREFIX))
                      for node in counter_nodes])
  
        # If there are no counters left in an old container, mark it inactive.
        if not counter_nodes and index < len(active_containers) - 1:
          self._inactive_containers.add(container_name)
  
      return txids
  
    def _delete_counter(self, path):
      """ Removes a counter node.
  
      Args:
        path: A string specifying a ZooKeeper path.
      """
      try:
        try:
          self.zk_client.delete(path)
        except NotEmptyError:
          # Cross-group transaction nodes have a child node.
          self.zk_client.delete(path, recursive=True)
      except KazooException:
        # Let the transaction groomer clean it up.
        logger.exception('Unable to delete counter')
  
    def _active_containers(self):
      """ Determines the containers that need to be checked for transactions.
  
      Returns:
        A tuple of strings specifying ZooKeeper paths.
      """
      container_name = self._counter_path.split('/')[-1]
      container_count = int(container_name[len(CONTAINER_PREFIX):] or 1)
  
      all_containers = [CONTAINER_PREFIX + str(index + 1)
                        for index in range(container_count)]
      all_containers[0] = CONTAINER_PREFIX
  
      return tuple('/'.join([self._project_node, container])
                   for container in all_containers
                   if container not in self._inactive_containers)
  
    def _update_auto_offset(self):
      """ Ensures there is a usable sequence container. """
      container_name = self._counter_path.split('/')[-1]
      container_count = int(container_name[len(CONTAINER_PREFIX):] or 1)
      next_node = CONTAINER_PREFIX + str(container_count + 1)
      next_path = '/'.join([self._project_node, next_node])
  
      try:
        self.zk_client.create(next_path)
      except NodeExistsError:
        # Another process may have already created the new counter.
        pass
      except KazooException:
        message = 'Unable to create transaction ID counter'
        logger.exception(message)
        raise InternalError(message)
  
      try:
        node_list = self.zk_client.get_children(self._project_node)
      except KazooException:
        message = 'Unable to find transaction ID counter'
        logger.exception(message)
        raise InternalError(message)
  
      self._update_project_sync(node_list)
  
    def _update_offset(self, new_offset, _):
      """ Watches for updates to the manual offset node. """
      # This assignment is atomic, so it does not need to happen in the IOLoop.
      self._txid_manual_offset = int(new_offset or 0)
  
    def _update_project_sync(self, node_list):
      """ Updates the record of usable sequence containers. """
      counters = [int(node[len(CONTAINER_PREFIX):] or 1)
                  for node in node_list if node.startswith(CONTAINER_PREFIX)]
      counters.sort()
  
      container_suffix = '' if len(counters) == 1 else str(counters[-1])
      latest_node = CONTAINER_PREFIX + container_suffix
  
      self._counter_path = '/'.join([self._project_node, latest_node])
  
      # The number of counters a container can store (including 0).
      container_size = MAX_SEQUENCE_COUNTER + 1
      self._txid_automatic_offset = (len(counters) - 1) * container_size
  
    def _update_project(self, node_list):
      """ Watches for updates to the list of containers. """
      IOLoop.instance().add_callback(self._update_project_sync, node_list)
  
  
  class TransactionManager(object):
    """ Generates and keeps track of transaction IDs. """
    def __init__(self, zk_client):
      """ Creates a new TransactionManager.
  
      Args:
        zk_client: A KazooClient.
      """
      self.zk_client = zk_client
      self.zk_client.ensure_path('/appscale/projects')
      self.projects = {}
  
      # Since this manager can be used synchronously, ensure that the projects
      # are populated for this IOLoop iteration.
      project_ids = self.zk_client.get_children('/appscale/projects')
      self._update_projects_sync(project_ids)
  
      self.zk_client.ChildrenWatch('/appscale/projects', self._update_projects)
  
    def create_transaction_id(self, project_id, xg=False):
      """ Generates a new transaction ID.
  
      Args:
        project_id: A string specifying a project ID.
        xg: A boolean indicating a cross-group transaction.
      Returns:
        An integer specifying the created transaction ID.
      Raises:
        BadRequest if the project does not exist.
        InternalError if unable to create a new transaction ID.
      """
      try:
        project_tx_manager = self.projects[project_id]
      except KeyError:
        raise BadRequest('The project {} was not found'.format(project_id))
  
      return project_tx_manager.create_transaction_id(xg)
  
    def delete_transaction_id(self, project_id, txid):
      """ Removes a transaction ID from the list of active transactions.
  
      Args:
        project_id: A string specifying a project ID.
        txid: An integer specifying a transaction ID.
      Raises:
        BadRequest if the project does not exist.
      """
      try:
        project_tx_manager = self.projects[project_id]
      except KeyError:
        raise BadRequest('The project {} was not found'.format(project_id))
  
      return project_tx_manager.delete_transaction_id(txid)
  
    def get_open_transactions(self, project_id):
      """ Fetch a list of open transactions for a given project.
  
      Args:
        project_id: A string specifying a project ID.
      Returns:
        A list of integers specifying transaction IDs.
      Raises:
        BadRequest if the project does not exist.
        InternalError if unable to fetch list of open transactions.
      """
      try:
        project_tx_manager = self.projects[project_id]
      except KeyError:
        raise BadRequest('The project {} was not found'.format(project_id))
  
      return project_tx_manager.get_open_transactions()
  
    def _update_projects_sync(self, new_project_ids):
      """ Updates the available projects for starting transactions.
  
      Args:
        new_project_ids: A list of strings specifying current project IDs.
      """
      for project_id in new_project_ids:
        if project_id not in self.projects:
          self.projects[project_id] = ProjectTransactionManager(project_id,
                                                                self.zk_client)
  
      for project_id in self.projects.keys():
        if project_id not in new_project_ids:
          del self.projects[project_id]
  
    def _update_projects(self, project_ids):
      """ Watches for changes to list of existing projects.
  
      Args:
        project_ids: A list of strings specifying current project IDs.
      """
      persistent_setup_projects = retry_children_watch_coroutine(
        '/appscale/projects', self._update_projects_sync)
      IOLoop.instance().add_callback(persistent_setup_projects, project_ids)
  
  
  [... the rest is omitted ...]
files.old_path: AppDB/appscale/datastore/zkappscale/transaction_manager.py
files.new_contents: #!/usr/bin/env python
  # Programmer: Navraj Chohan <nlake44@gmail.com>
  
  import datetime
  import random
  import sys
  import unittest
  
  from appscale.common import appscale_info
  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.cassandra_interface import (
    DatastoreProxy)
  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 (
    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 cassandra.cluster import Cluster
  from flexmock import flexmock
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.api import api_base_pb
  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)
  
  
  class TestDatastoreServer(unittest.TestCase):
    """
    A set of test cases for the datastore server (datastore server v2)
    """
    BASIC_ENTITY = ['guestbook', 'Greeting', 'foo', 'content', 'hello world']
  
    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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      item = Item(name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_entity_kind(key), "Item")
  
    def test_kind_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').and_return(True)
      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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      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').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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
  
      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)
      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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      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').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")
  
    def test_get_indices(self):
      session = flexmock(default_consistency_level=None)
      cluster = flexmock(connect=lambda keyspace: session)
      flexmock(appscale_info).should_receive('get_db_ips')
      flexmock(Cluster).new_instances(cluster)
      flexmock(DatastoreProxy).should_receive('range_query').and_return({})
      db_batch = DatastoreProxy()
  
      self.assertEquals(db_batch.get_indices("appid"), [])
  
    def test_delete_composite_index_metadata(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd = flexmock(dd)
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(1)
      dd.delete_composite_index_metadata("appid", composite_index)
  
    def test_create_composite_index(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      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
  
      dd.create_composite_index("appid", index)
      assert index.id() > 0 
  
    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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None).once()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.insert_composite_indexes([ent], [composite_index])
  
    def testFetchKeys(self):
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "bob", "prop1name", 
                                                "prop1val", ns="blah")
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return({'test\x00blah\x00test_kind:bob\x01':
                  {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(), 
                   APP_ENTITY_SCHEMA[1]: 1}}).and_return({'test\x00blah\x00test_kind:bob\x01\x000000000002':
                  {JOURNAL_SCHEMA[0]: entity_proto1.Encode()}})
  
      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)
  
      self.assertEquals(({'test\x00blah\x00test_kind:bob\x01':
                           {'txnID': 1, 'entity': entity_proto1.Encode()}
                         },
                         ['test\x00blah\x00test_kind:bob\x01']),
                         dd.fetch_keys([entity_proto1.key()]))
  
    def test_commit_transaction(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').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')
      commit_request = datastore_pb.Transaction()
      commit_request.set_handle(123)
      commit_request.set_app("aaa")
      http_request = commit_request.Encode()
      self.assertEquals(dd.commit_transaction("app_id", http_request),
                        (datastore_pb.CommitResponse().Encode(), 0, ""))
  
    def test_rollback_transcation(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').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")
      http_request = commit_request.Encode()
      self.assertEquals(dd.rollback_transaction("app_id", http_request),
                        (api_base_pb.VoidProto().Encode(), 0, ""))
  
    def get_new_entity_proto(self, 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
  
    def test_dynamic_put(self):
      PREFIX = "x\x01"
      db_batch = flexmock(session=flexmock())
      db_batch.should_receive('valid_data_version').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'
  
      db_batch.should_receive('batch_get_entity').and_return(
        {entity_key1: {}, entity_key2: {}})
      db_batch.should_receive('batch_mutate')
      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())
      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()
  
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire')
      entity_lock.should_receive('release')
  
      flexmock(ScatteredAllocator).should_receive('next').\
        and_return(random.randint(1, 500))
  
      dd.dynamic_put('test', putreq_pb, putresp_pb)
      self.assertEquals(len(putresp_pb.key_list()), 2)
  
    def test_put_entities(self):
      app_id = 'test'
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').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]
  
      db_batch.should_receive('batch_get_entity').and_return(
        {entity_key1: {}, entity_key2: {}})
      db_batch.should_receive('batch_mutate')
      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())
  
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire')
      entity_lock.should_receive('release')
  
      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').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)
  
    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'}}
  
      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').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(row_values)
      db_batch.should_receive('batch_mutate')
      db_batch.should_receive('_normal_batch')
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      row_keys = [entity_proto1.key()]
  
      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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
      db_batch.should_receive("batch_delete").and_return(None)
  
      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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
  
      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()))
  
    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')
  
      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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
      db_batch.should_receive('record_reads')
  
      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()
      
      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)
      dd.dynamic_get("test", get_req, get_resp)     
      self.assertEquals(get_resp.entity_size(), 1)
  
    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)
      
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('record_reads')
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
  
      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)
      dd.ancestor_query(query, filter_info, None)
      # Now with a transaction
      transaction = query.mutable_transaction() 
      transaction.set_handle(2)
      dd.ancestor_query(query, filter_info, None)
  
    def test_ordered_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)
  
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01":
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
      db_batch.should_receive('record_reads')
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      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)
      dd.ordered_ancestor_query(query, filter_info, None)
  
      # Now with a transaction
      transaction = query.mutable_transaction()
      transaction.set_handle(2)
      dd.ordered_ancestor_query(query, filter_info, None) 
  
    
    def test_kindless_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)
      
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
  
      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]]
      }
      dd.kindless_query(query, filter_info)
  
    def test_dynamic_delete(self):
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire')
      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').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.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())
      flexmock(utils).should_receive("get_entity_kind").and_return("kind")
      db_batch.should_receive('delete_entities_tx')
      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())
      flexmock(dd).should_receive("delete_entities").once()
      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)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      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)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      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)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      query = datastore_pb.Query()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      db_batch.should_receive("remove_exists_filters").and_return({})
      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")]}
      db_batch.should_receive("remove_exists_filters").and_return(filter_info)
           
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), True)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), False)
  
    def test_zigzag_merge_join(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("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      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)
      self.assertEquals(dd.zigzag_merge_join(None, None, None), None)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      flexmock(query).should_receive("kind").and_return("kind")
      flexmock(dd).should_receive("get_table_prefix").and_return("prefix")
      flexmock(dd).should_receive("__apply_filters").and_return([])
      flexmock(query).should_receive("limit").and_return(1)
      self.assertEquals(dd.zigzag_merge_join(query, filter_info, []), None)
  
    def test_index_deletions(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      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').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      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
  
  [... 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 appscale.common import appscale_info
  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.cassandra_interface import (
    DatastoreProxy)
  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 (
    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 cassandra.cluster import Cluster
  from flexmock import flexmock
  
  sys.path.append(APPSCALE_PYTHON_APPSERVER)
  from google.appengine.api import api_base_pb
  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)
  
  
  class TestDatastoreServer(unittest.TestCase):
    """
    A set of test cases for the datastore server (datastore server v2)
    """
    BASIC_ENTITY = ['guestbook', 'Greeting', 'foo', 'content', 'hello world']
  
    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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      item = Item(name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_entity_kind(key), "Item")
  
    def test_kind_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').and_return(True)
      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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      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').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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
  
      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)
      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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      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):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      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').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")
  
    def test_get_indices(self):
      session = flexmock(default_consistency_level=None)
      cluster = flexmock(connect=lambda keyspace: session)
      flexmock(appscale_info).should_receive('get_db_ips')
      flexmock(Cluster).new_instances(cluster)
      flexmock(DatastoreProxy).should_receive('range_query').and_return({})
      db_batch = DatastoreProxy()
  
      self.assertEquals(db_batch.get_indices("appid"), [])
  
    def test_delete_composite_index_metadata(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd = flexmock(dd)
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(1)
      dd.delete_composite_index_metadata("appid", composite_index)
  
    def test_create_composite_index(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      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
  
      dd.create_composite_index("appid", index)
      assert index.id() > 0 
  
    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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None).once()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      dd.insert_composite_indexes([ent], [composite_index])
  
    def testFetchKeys(self):
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "bob", "prop1name", 
                                                "prop1val", ns="blah")
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return({'test\x00blah\x00test_kind:bob\x01':
                  {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(), 
                   APP_ENTITY_SCHEMA[1]: 1}}).and_return({'test\x00blah\x00test_kind:bob\x01\x000000000002':
                  {JOURNAL_SCHEMA[0]: entity_proto1.Encode()}})
  
      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)
  
      self.assertEquals(({'test\x00blah\x00test_kind:bob\x01':
                           {'txnID': 1, 'entity': entity_proto1.Encode()}
                         },
                         ['test\x00blah\x00test_kind:bob\x01']),
                         dd.fetch_keys([entity_proto1.key()]))
  
    def test_commit_transaction(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').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')
      commit_request = datastore_pb.Transaction()
      commit_request.set_handle(123)
      commit_request.set_app("aaa")
      http_request = commit_request.Encode()
      self.assertEquals(dd.commit_transaction("app_id", http_request),
                        (datastore_pb.CommitResponse().Encode(), 0, ""))
  
    def test_rollback_transcation(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').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")
      http_request = commit_request.Encode()
      self.assertEquals(dd.rollback_transaction("app_id", http_request),
                        (api_base_pb.VoidProto().Encode(), 0, ""))
  
    def get_new_entity_proto(self, 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
  
    def test_dynamic_put(self):
      PREFIX = "x\x01"
      db_batch = flexmock(session=flexmock())
      db_batch.should_receive('valid_data_version').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'
  
      db_batch.should_receive('batch_get_entity').and_return(
        {entity_key1: {}, entity_key2: {}})
      db_batch.should_receive('batch_mutate')
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      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()
  
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire')
      entity_lock.should_receive('release')
  
      flexmock(ScatteredAllocator).should_receive('next').\
        and_return(random.randint(1, 500))
  
      dd.dynamic_put('test', putreq_pb, putresp_pb)
      self.assertEquals(len(putresp_pb.key_list()), 2)
  
    def test_put_entities(self):
      app_id = 'test'
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').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]
  
      db_batch.should_receive('batch_get_entity').and_return(
        {entity_key1: {}, entity_key2: {}})
      db_batch.should_receive('batch_mutate')
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
  
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire')
      entity_lock.should_receive('release')
  
      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').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)
  
    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'}}
  
      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').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(row_values)
      db_batch.should_receive('batch_mutate')
      db_batch.should_receive('_normal_batch')
  
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
  
      row_keys = [entity_proto1.key()]
  
      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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
      db_batch.should_receive("batch_delete").and_return(None)
  
      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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
  
      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()))
  
    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')
  
      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').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
      db_batch.should_receive('record_reads')
  
      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()
      
      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)
      dd.dynamic_get("test", get_req, get_resp)     
      self.assertEquals(get_resp.entity_size(), 1)
  
    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)
      
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('record_reads')
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
  
      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)
      dd.ancestor_query(query, filter_info, None)
      # Now with a transaction
      transaction = query.mutable_transaction() 
      transaction.set_handle(2)
      dd.ancestor_query(query, filter_info, None)
  
    def test_ordered_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)
  
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01":
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
      db_batch.should_receive('record_reads')
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      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)
      dd.ordered_ancestor_query(query, filter_info, None)
  
      # Now with a transaction
      transaction = query.mutable_transaction()
      transaction.set_handle(2)
      dd.ordered_ancestor_query(query, filter_info, None) 
  
    
    def test_kindless_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)
      
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
  
      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]]
      }
      dd.kindless_query(query, filter_info)
  
    def test_dynamic_delete(self):
      entity_lock = flexmock(EntityLock)
      entity_lock.should_receive('acquire')
      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').and_return(True)
      transaction_manager = flexmock(
        create_transaction_id=lambda project, xg: 1,
        delete_transaction_id=lambda project, txid: None)
      dd = DatastoreDistributed(db_batch, transaction_manager,
                                self.get_zookeeper())
      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())
      flexmock(utils).should_receive("get_entity_kind").and_return("kind")
      db_batch.should_receive('delete_entities_tx')
      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())
      flexmock(dd).should_receive("delete_entities").once()
      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)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      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)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      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)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      query = datastore_pb.Query()
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      db_batch.should_receive("remove_exists_filters").and_return({})
      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")]}
      db_batch.should_receive("remove_exists_filters").and_return(filter_info)
           
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), True)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), False)
  
    def test_zigzag_merge_join(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("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      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)
      self.assertEquals(dd.zigzag_merge_join(None, None, None), None)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      flexmock(query).should_receive("kind").and_return("kind")
      flexmock(dd).should_receive("get_table_prefix").and_return("prefix")
      flexmock(dd).should_receive("__apply_filters").and_return([])
      flexmock(query).should_receive("limit").and_return(1)
      self.assertEquals(dd.zigzag_merge_join(query, filter_info, []), None)
  
    def test_index_deletions(self):
      zk_client = flexmock()
      zk_client.should_receive('add_listener')
  
      zookeeper = flexmock(handle=zk_client)
      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').and_return(True)
      transaction_manager = flexmock()
      dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
      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
  
  [... the rest is omitted ...]
files.old_path: AppDB/test/unit/test_datastore_server.py