files: 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(
         
  
  [... the rest is omitted ...]
files: AppDB/appscale/datastore/datastore_distributed.py
files: 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
  
  [... the rest is omitted ...]
files: AppDB/appscale/datastore/datastore_distributed.py
files: """ 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: AppDB/appscale/datastore/zkappscale/transaction_manager.py
files: """ 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: AppDB/appscale/datastore/zkappscale/transaction_manager.py
files: #!/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\x00
  
  [... the rest is omitted ...]
files: AppDB/test/unit/test_datastore_server.py
files: #!/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[
  
  [... the rest is omitted ...]
files: AppDB/test/unit/test_datastore_server.py