files: # Programmer: Navraj Chohan <nlake44@gmail.com>
  
  """
   Cassandra Interface for AppScale
  """
  import cassandra
  import logging
  import sys
  import time
  
  from cassandra.cluster import Cluster
  from cassandra.concurrent import execute_concurrent
  from cassandra.policies import FallthroughRetryPolicy
  from cassandra.policies import RetryPolicy
  from cassandra.query import BatchStatement
  from cassandra.query import ConsistencyLevel
  from cassandra.query import SimpleStatement
  from cassandra.query import ValueSequence
  from .. import dbconstants
  from .. import helper_functions
  from ..dbconstants import AppScaleDBConnectionError
  from ..dbconstants import Operations
  from ..dbconstants import TxnActions
  from ..dbinterface import AppDBInterface
  from ..unpackaged import APPSCALE_LIB_DIR
  from ..utils import clean_app_id
  from ..utils import encode_index_pb
  from ..utils import get_composite_index_keys
  from ..utils import get_composite_indexes_rows
  from ..utils import get_entity_key
  from ..utils import get_entity_kind
  from ..utils import get_index_kv_from_tuple
  from ..utils import get_kind_key
  
  sys.path.append(APPSCALE_LIB_DIR)
  import appscale_info
  
  # The directory Cassandra is installed to.
  CASSANDRA_INSTALL_DIR = '/opt/cassandra'
  
  # Full path for the nodetool binary.
  NODE_TOOL = '{}/cassandra/bin/nodetool'.format(CASSANDRA_INSTALL_DIR)
  
  # The keyspace used for all tables
  KEYSPACE = "Keyspace1"
  
  # Cassandra watch name.
  CASSANDRA_MONIT_WATCH_NAME = "cassandra-9999"
  
  # The number of times to retry connecting to Cassandra.
  INITIAL_CONNECT_RETRIES = 20
  
  # The data layout version that the datastore expects.
  EXPECTED_DATA_VERSION = 1.0
  
  # The metadata key for the data layout version.
  VERSION_INFO_KEY = 'version'
  
  # The metadata key indicating that the database has been primed.
  PRIMED_KEY = 'primed'
  
  # The size in bytes that a batch must be to use the batches table.
  LARGE_BATCH_THRESHOLD = 5 << 10
  
  
  def batch_size(batch):
    """ Calculates the size of a batch.
  
    Args:
      batch: A list of dictionaries representing mutations.
    Returns:
      An integer specifying the size in bytes of the batch.
    """
    size = 0
    for mutation in batch:
      size += len(mutation['key'])
      if 'values' in mutation:
        for value in mutation['values'].values():
          size += len(value)
    return size
  
  
  def deletions_for_entity(entity, composite_indices=()):
    """ Get a list of deletions needed across tables for deleting an entity.
  
    Args:
      entity: An entity object.
      composite_indices: A list or tuple of composite indices.
    Returns:
      A list of dictionaries representing mutation operations.
    """
    deletions = []
    app_id = clean_app_id(entity.key().app())
    namespace = entity.key().name_space()
    prefix = dbconstants.KEY_DELIMITER.join([app_id, namespace])
  
    asc_rows = get_index_kv_from_tuple([(prefix, entity)])
    for entry in asc_rows:
      deletions.append({'table': dbconstants.ASC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': Operations.DELETE})
  
    dsc_rows = get_index_kv_from_tuple(
      [(prefix, entity)], reverse=True)
    for entry in dsc_rows:
      deletions.append({'table': dbconstants.DSC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': Operations.DELETE})
  
    for key in get_composite_indexes_rows([entity], composite_indices):
      deletions.append({'table': dbconstants.COMPOSITE_TABLE,
                        'key': key,
                        'operation': Operations.DELETE})
  
    entity_key = get_entity_key(prefix, entity.key().path())
    deletions.append({'table': dbconstants.APP_ENTITY_TABLE,
                      'key': entity_key,
                      'operation': Operations.DELETE})
  
    kind_key = get_kind_key(prefix, entity.key().path())
    deletions.append({'table': dbconstants.APP_KIND_TABLE,
                      'key': kind_key,
                      'operation': Operations.DELETE})
  
    return deletions
  
  
  def index_deletions(old_entity, new_entity, composite_indices=()):
    """ Get a list of index deletions needed for updating an entity. For changing
    an existing entity, this involves examining the property list of both
    entities to see which index entries need to be removed.
  
    Args:
      old_entity: An entity object.
      new_entity: An entity object.
      composite_indices: A list or tuple of composite indices.
    Returns:
      A list of dictionaries representing mutation operations.
    """
    deletions = []
    app_id = clean_app_id(old_entity.key().app())
    namespace = old_entity.key().name_space()
    kind = get_entity_kind(old_entity.key())
    entity_key = str(encode_index_pb(old_entity.key().path()))
  
    new_props = {}
    for prop in new_entity.property_list():
      if prop.name() not in new_props:
        new_props[prop.name()] = []
      new_props[prop.name()].append(prop)
  
    changed_props = {}
    for prop in old_entity.property_list():
      if prop.name() in new_props and prop in new_props[prop.name()]:
        continue
  
      if prop.name() not in changed_props:
        changed_props[prop.name()] = []
      changed_props[prop.name()].append(prop)
  
      value = str(encode_index_pb(prop.value()))
  
      key = dbconstants.KEY_DELIMITER.join(
        [app_id, namespace, kind, prop.name(), value, entity_key])
      deletions.append({'table': dbconstants.ASC_PROPERTY_TABLE,
                        'key': key,
                        'operation': Operations.DELETE})
  
      reverse_key = dbconstants.KEY_DELIMITER.join(
        [app_id, namespace, kind, prop.name(),
         helper_functions.reverse_lex(value), entity_key])
      deletions.append({'table': dbconstants.DSC_PROPERTY_TABLE,
                        'key': reverse_key,
                        'operation': Operations.DELETE})
  
    changed_prop_names = set(changed_props.keys())
    for index in composite_indices:
      if index.definition().entity_type() != kind:
        continue
  
      index_props = set(prop.name() for prop
                        in index.definition().property_list())
      if index_props.isdisjoint(changed_prop_names):
        continue
  
      old_entries = set(get_composite_index_keys(index, old_entity))
      new_entries = set(get_composite_index_keys(index, new_entity))
      for entry in (old_entries - new_entries):
        deletions.append({'table': dbconstants.COMPOSITE_TABLE,
                          'key': entry,
                          'operation': Operations.DELETE})
  
    return deletions
  
  
  def mutations_for_entity(entity, txn, current_value=None,
                           composite_indices=()):
    """ Get a list of mutations needed across tables for an entity change.
  
    Args:
      entity: An entity object.
      txn: A transaction ID handler.
      current_value: The entity object currently stored.
      composite_indices: A list of composite indices for the entity kind.
    Returns:
      A list of dictionaries representing mutations.
    """
    mutations = []
    if current_value is not None:
      mutations.extend(
        index_deletions(current_value, entity, composite_indices))
  
    app_id = clean_app_id(entity.key().app())
    namespace = entity.key().name_space()
    encoded_path = str(encode_index_pb(entity.key().path()))
    prefix = dbconstants.KEY_DELIMITER.join([app_id, namespace])
    entity_key = dbconstants.KEY_DELIMITER.join([prefix, encoded_path])
    entity_value = {dbconstants.APP_ENTITY_SCHEMA[0]: entity.Encode(),
                    dbconstants.APP_ENTITY_SCHEMA[1]: str(txn)}
    mutations.append({'table': dbconstants.APP_ENTITY_TABLE,
                      'key': entity_key,
                      'operation': Operations.PUT,
                      'values': entity_value})
  
    reference_value = {'reference': entity_key}
  
    kind_key = get_kind_key(prefix, entity.key().path())
    mutations.append({'table': dbconstants.APP_KIND_TABLE,
                      'key': kind_key,
                      'operation': Operations.PUT,
                      'values': reference_value})
  
    asc_rows = get_index_kv_from_tuple([(prefix, entity)])
    for entry in asc_rows:
      mutations.append({'table': dbconstants.ASC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': Operations.PUT,
                        'values': reference_value})
  
    dsc_rows = get_index_kv_from_tuple([(prefix, entity)], reverse=True)
    for entry in dsc_rows:
      mutations.append({'table': dbconstants.DSC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': Operations.PUT,
                        'values': reference_value})
  
    for key in get_composite_indexes_rows([entity], composite_indices):
      mutations.append({'table': dbconstants.COMPOSITE_TABLE,
                        'key': key,
                        'operation': Operations.PUT,
                        'values': reference_value})
  
    return mutations
  
  
  class FailedBatch(Exception):
    pass
  
  
  class IdempotentRetryPolicy(RetryPolicy):
    """ A policy used for retrying idempotent statements. """
    def on_read_timeout(self, query, consistency, required_responses,
                        received_responses, data_retrieved, retry_num):
      """ This is called when a ReadTimeout occurs.
  
      Args:
        query: A statement that timed out.
        consistency: The consistency level of the statement.
        required_responses: The number of responses required.
        received_responses: The number of responses received.
        data_retrieved: Indicates whether any responses contained data.
        retry_num: The number of times the statement has been tried.
      """
      if retry_num >= 5:
        return self.RETHROW, None
      else:
        return self.RETRY, consistency
  
    def on_write_timeout(self, query, consistency, write_type,
                         required_responses, received_responses, retry_num):
      """ This is called when a WriteTimeout occurs.
  
      Args:
        query: A statement that timed out.
        consistency: The consistency level of the statement.
        required_responses: The number of responses required.
        received_responses: The number of responses received.
        data_retrieved: Indicates whether any responses contained data.
        retry_num: The number of times the statement has been tried.
        """
      if retry_num >= 5:
        return self.RETHROW, None
      else:
        return self.RETRY, consistency
  
  
  class ThriftColumn(object):
    """ Columns created by default with thrift interface. """
    KEY = 'key'
    COLUMN_NAME = 'column1'
    VALUE = 'value'
  
  
  class DatastoreProxy(AppDBInterface):
    """ 
      Cassandra implementation of the AppDBInterface
    """
    def __init__(self, log_level=logging.INFO):
      """
      Constructor.
      """
      class_name = self.__class__.__name__
      self.logger = logging.getLogger(class_name)
      self.logger.setLevel(log_level)
      self.logger.info('Starting {}'.format(class_name))
  
      self.hosts = appscale_info.get_db_ips()
      self.retry_policy = IdempotentRetryPolicy()
      self.no_retries = FallthroughRetryPolicy()
  
      remaining_retries = INITIAL_CONNECT_RETRIES
      while True:
        try:
          self.cluster = Cluster(self.hosts,
                                 default_retry_policy=self.retry_policy)
          self.session = self.cluster.connect(KEYSPACE)
          break
        except cassandra.cluster.NoHostAvailable as connection_error:
          remaining_retries -= 1
          if remaining_retries < 0:
            raise connection_error
          time.sleep(3)
  
      self.session.default_consistency_level = ConsistencyLevel.QUORUM
  
    def close(self):
      """ Close all sessions and connections to Cassandra. """
      self.cluster.shutdown()
  
    def batch_get_entity(self, table_name, row_keys, column_names):
      """
      Takes in batches of keys and retrieves their corresponding rows.
      
      Args:
        table_name: The table to access
        row_keys: A list of keys to access
        column_names: A list of columns to access
      Returns:
        A dictionary of rows and columns/values of those rows. The format 
        looks like such: {key:{column_name:value,...}}
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the batch_get could not be performed due to
          an error with Cassandra.
      """
      if not isinstance(table_name, str): raise TypeError("Expected a str")
      if not isinstance(column_names, list): raise TypeError("Expected a list")
      if not isinstance(row_keys, list): raise TypeError("Expected a list")
  
      row_keys_bytes = [bytearray(row_key) for row_key in row_keys]
  
      statement = 'SELECT * FROM "{table}" '\
                  'WHERE {key} IN %s and {column} IN %s'.format(
                    table=table_name,
                    key=ThriftColumn.KEY,
                    column=ThriftColumn.COLUMN_NAME,
                  )
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
      parameters = (ValueSequence(row_keys_bytes), ValueSequence(column_names))
  
      try:
        results = self.session.execute(query, parameters=parameters)
  
        results_dict = {row_key: {} for row_key in row_keys}
        for (key, column, value) in results:
          if key not in results_dict:
            results_dict[key] = {}
          results_dict[key][column] = value
  
        return results_dict
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_get_entity'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def batch_put_entity(self, table_name, row_keys, column_names, cell_values,
                         ttl=None):
      """
      Allows callers to store multiple rows with a single call. A row can 
      have multiple columns and values with them. We refer to each row as 
      an entity.
     
      Args: 
        table_name: The table to mutate
        row_keys: A list of keys to store on
        column_names: A list of columns to mutate
        cell_values: A dict of key/value pairs
        ttl: The number of seconds to keep the row.
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the batch_put could not be performed due to
          an error with Cassandra.
      """
      if not isinstance(table_name, str):
        raise TypeError("Expected a str")
      if not isinstance(column_names, list):
        raise TypeError("Expected a list")
      if not isinstance(row_keys, list):
        raise TypeError("Expected a list")
      if not isinstance(cell_values, dict):
        raise TypeError("Expected a dict")
  
      insert_str = """
        INSERT INTO "{table}" ({key}, {column}, {value})
        VALUES (?, ?, ?)
      """.format(table=table_name,
                 key=ThriftColumn.KEY,
                 column=ThriftColumn.COLUMN_NAME,
                 value=ThriftColumn.VALUE)
  
      if ttl is not None:
        insert_str += 'USING TTL {}'.format(ttl)
  
      statement = self.session.prepare(insert_str)
  
      statements_and_params = []
      for row_key in row_keys:
        for column in column_names:
          params = (bytearray(row_key), column,
                    bytearray(cell_values[row_key][column]))
          statements_and_params.append((statement, params))
  
      try:
        execute_concurrent(self.session, statements_and_params,
                           raise_on_first_error=True)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_put_entity'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def prepare_insert(self, table):
      """ Prepare an insert statement.
  
      Args:
        table: A string containing the table name.
      Returns:
        A PreparedStatement object.
      """
      statement = """
        INSERT INTO "{table}" ({key}, {column}, {value})
        VALUES (?, ?, ?)
      """.format(table=table,
                 key=ThriftColumn.KEY,
                 column=ThriftColumn.COLUMN_NAME,
                 value=ThriftColumn.VALUE)
      return self.session.prepare(statement)
  
    def prepare_delete(self, table):
      """ Prepare a delete statement.
  
      Args:
        table: A string containing the table name.
      Returns:
        A PreparedStatement object.
      """
      statement = """
        DELETE FROM "{table}" WHERE {key} = ?
      """.format(table=table, key=ThriftColumn.KEY)
      return self.session.prepare(statement)
  
    def _normal_batch(self, mutations):
      """ Use Cassandra's native batch statement to apply mutations atomically.
  
      Args:
        mutations: A list of dictionaries representing mutations.
      """
      self.logger.debug('Normal batch: {} mutations'.format(len(mutations)))
      batch = BatchStatement(consistency_level=ConsistencyLevel.QUORUM,
                             retry_policy=self.retry_policy)
      prepared_statements = {'insert': {}, 'delete': {}}
      for mutation in mutations:
        table = mutation['table']
        if mutation['operation'] == Operations.PUT:
          if table not in prepared_statements['insert']:
            prepared_statements['insert'][table] = self.prepare_insert(table)
          values = mutation['values']
          for column in values:
            batch.add(
              prepared_statements['insert'][table],
              (bytearray(mutation['key']), column, bytearray(values[column]))
            )
        elif mutation['operation'] == Operations.DELETE:
          if table not in prepared_statements['delete']:
            prepared_statements['delete'][table] = self.prepare_delete(table)
          batch.add(
            prepared_statements['delete'][table],
            (bytearray(mutation['key']),)
          )
  
      try:
        self.session.execute(batch)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_mutate'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def apply_mutations(self, mutations):
      """ Apply mutations across tables.
  
      Args:
        mutations: A list of dictionaries representing mutations.
      """
      prepared_statements = {'insert': {}, 'delete': {}}
      statements_and_params = []
      for mutation in mutations:
        table = mutation['table']
        if mutation['operation'] == Operations.PUT:
          if table not in prepared_statements['insert']:
            prepared_statements['insert'][table] = self.prepare_insert(table)
          values = mutation['values']
          for column in values:
            params = (bytearray(mutation['key']), column,
                      bytearray(values[column]))
            statements_and_params.append(
              (prepared_statements['insert'][table], params))
        elif mutation['operation'] == Operations.DELETE:
          if table not in prepared_statements['delete']:
            prepared_statements['delete'][table] = self.prepare_delete(table)
          params = (bytearray(mutation['key']),)
          statements_and_params.append(
            (prepared_statements['delete'][table], params))
  
      execute_concurrent(self.session, statements_and_params,
                         raise_on_first_error=True)
  
    def _large_batch(self, app, mutations, entity_changes, txn):
      """ Insert or delete multiple rows across tables in an atomic statement.
  
      Args:
        app: A string containing the application ID.
        mutations: A list of dictionaries representing mutations.
        entity_changes: A list of changes at the entity level.
        txn: A transaction ID handler.
      Raises:
        FailedBatch if a concurrent process modifies the batch status.
        AppScaleDBConnectionError if a database connection error was encountered.
      """
      self.logger.debug('Large batch: transaction {}, {} mutations'.
                        format(txn, len(mutations)))
      set_status = """
        INSERT INTO batch_status (app, transaction, applied)
        VALUES (%(app)s, %(transaction)s, False)
        IF NOT EXISTS
      """
      parameters = {'app': app, 'transaction': txn}
      result = self.session.execute(set_status, parameters)
      if not result.was_applied:
        raise FailedBatch('A batch for transaction {} already exists'.
                          format(txn))
  
      insert_item = """
        INSERT INTO batches (app, transaction, namespace, path,
                             old_value, new_value)
        VALUES (?, ?, ?, ?, ?, ?)
      """
      insert_statement = self.session.prepare(insert_item)
  
      statements_and_params = []
      for entity_change in entity_changes:
        old_value = None
        if entity_change['old'] is not None:
          old_value = bytearray(entity_change['old'].Encode())
        new_value = None
        if entity_change['new'] is not None:
          new_value = bytearray(entity_change['new'].Encode())
  
        parameters = (app, txn, entity_change['key'].name_space(),
                      bytearray(entity_change['key'].path().Encode()), old_value,
                      new_value)
        statements_and_params.append((insert_statement, parameters))
  
      try:
        execute_concurrent(self.session, statements_and_params,
                           raise_on_first_error=True)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during large batch'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
      update_status = """
        UPDATE batch_status
        SET applied = True
        WHERE app = %(app)s
        AND transaction = %(transaction)s
        IF applied = False
      """
      parameters = {'app': app, 'transaction': txn}
      result = self.session.execute(update_status, parameters)
      if not result.was_applied:
        raise FailedBatch('Another process modified batch for transaction {}'.
                          format(txn))
  
      try:
        self.apply_mutations(mutations)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during large batch'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
      clear_batch = """
        DELETE FROM batches
        WHERE app = %(app)s AND transaction = %(transaction)s
      """
      parameters = {'app': app, 'transaction': txn}
      self.session.execute(clear_batch, parameters)
  
      clear_status = """
        DELETE FROM batch_status
        WHERE app = %(app)s and transaction = %(transaction)s
      """
      parameters = {'app': app, 'transaction': txn}
      self.session.execute(clear_status, parameters)
  
    def batch_mutate(self, app, mutations, entity_changes, txn):
      """ Insert or delete multiple rows across tables in an atomic statement.
  
      Args:
        app: A string containing the application ID.
        mutations: A list of dictionaries representing mutations.
        entity_changes: A list of changes at the entity level.
        txn: A transaction ID handler.
      """
      size = batch_size(mutations)
      self.logger.debug('batch_size: {}'.format(size))
      if size > LARGE_BATCH_THRESHOLD:
        self._large_batch(app, mutations, entity_changes, txn)
      else:
        self._normal_batch(mutations)
  
    def batch_delete(self, table_name, row_keys, column_names=()):
      """
      Remove a set of rows corresponding to a set of keys.
       
      Args:
        table_name: Table to delete rows from
        row_keys: A list of keys to remove
        column_names: Not used
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the batch_delete could not be performed due
          to an error with Cassandra.
      """ 
      if not isinstance(table_name, str): raise TypeError("Expected a str")
      if not isinstance(row_keys, list): raise TypeError("Expected a list")
  
      row_keys_bytes = [bytearray(row_key) for row_key in row_keys]
  
      statement = 'DELETE FROM "{table}" WHERE {key} IN %s'.\
        format(
          table=table_name,
          key=ThriftColumn.KEY
        )
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
      parameters = (ValueSequence(row_keys_bytes),)
  
      try:
        self.session.execute(query, parameters=parameters)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_delete'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def delete_table(self, table_name):
      """ 
      Drops a given table (aka column family in Cassandra)
    
      Args:
        table_name: A string name of the table to drop
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the delete_table could not be performed due
          to an error with Cassandra.
      """
      if not isinstance(table_name, str): raise TypeError("Expected a str")
  
      statement = 'DROP TABLE IF EXISTS "{table}"'.format(table=table_name)
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
  
      try:
        self.session.execute(query)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during delete_table'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def create_table(self, table_name, column_names):
      """ 
      Creates a table if it doesn't already exist.
      
      Args:
        table_name: The column family name
        column_names: Not used but here to match the interface
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the create_table could not be performed due
          to an error with Cassandra.
      """
      if not isinstance(table_name, str): raise TypeError("Expected a str")
      if not isinstance(column_names, list): raise TypeError("Expected a list")
  
      statement = 'CREATE TABLE IF NOT EXISTS "{table}" ('\
          '{key} blob,'\
          '{column} text,'\
          '{value} blob,'\
          'PRIMARY KEY ({key}, {column})'\
        ') WITH COMPACT STORAGE'.format(
          table=table_name,
          key=ThriftColumn.KEY,
          column=ThriftColumn.COLUMN_NAME,
          value=ThriftColumn.VALUE
        )
      query = SimpleStatement(statement, retry_policy=self.no_retries)
  
      try:
        self.session.execute(query)
      except cassandra.OperationTimedOut:
        logging.warning('Encountered an operation timeout while creating a '
                        'table. Waiting 1 minute for schema to settle.')
        time.sleep(60)
        raise AppScaleDBConnectionError('Exception during create_table')
      except (error for error in dbconstants.TRANSIENT_CASSANDRA_ERRORS
              if error != cassandra.OperationTimedOut):
        message = 'Exception during create_table'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def range_query(self,
                    table_name,
                    column_names, 
                    start_key, 
                    end_key, 
                    limit, 
                    offset=0, 
                    start_inclusive=True, 
                    end_inclusive=True,
                    keys_only=False):
      """ 
      Gets a dense range ordered by keys. Returns an ordered list of 
      a dictionary of [key:{column1:value1, column2:value2},...]
      or a list of keys if keys only.
       
      Args:
        table_name: Name of table to access
        column_names: Columns which get returned within the key range
        start_key: String for which the query starts at
        end_key: String for which the query ends at
        limit: Maximum number of results to return
        offset: Cuts off these many from the results [offset:]
        start_inclusive: Boolean if results should include the start_key
        end_inclusive: Boolean if results should include the end_key
        keys_only: Boolean if to only keys and not values
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the range_query could not be performed due
          to an error with Cassandra.
      Returns:
        An ordered list of dictionaries of key=>columns/values
      """
      if not isinstance(table_name, str):
        raise TypeError('table_name must be a string')
      if not isinstance(column_names, list):
        raise TypeError('column_names must be a list')
      if not isinstance(start_key, str):
        raise TypeError('start_key must be a string')
      if not isinstance(end_key, str):
        raise TypeError('end_key must be a string')
      if not isinstance(limit, (int, long)) and limit is not None:
        raise TypeError('limit must be int, long, or NoneType')
      if not isinstance(offset, (int, long)):
        raise TypeError('offset must be int or long')
  
      if start_inclusive:
        gt_compare = '>='
      else:
        gt_compare = '>'
  
      if end_inclusive:
        lt_compare = '<='
      else:
        lt_compare = '<'
  
      query_limit = ''
      if limit is not None:
        query_limit = 'LIMIT {}'.format(len(column_names) * limit)
  
      statement = """
        SELECT * FROM "{table}" WHERE
        token({key}) {gt_compare} %s AND
        token({key}) {lt_compare} %s AND
        {column} IN %s
        {limit}
        ALLOW FILTERING
      """.format(table=table_name,
                 key=ThriftColumn.KEY,
                 gt_compare=gt_compare,
                 lt_compare=lt_compare,
                 column=ThriftColumn.COLUMN_NAME,
                 limit=query_limit)
  
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
      parameters = (bytearray(start_key), bytearray(end_key),
                    ValueSequence(column_names))
  
      try:
        results = self.session.execute(query, parameters=parameters)
  
        results_list = []
        current_item = {}
        current_key = None
        for (key, column, value) in results:
          if keys_only:
            results_list.append(key)
            continue
  
          if key != current_key:
            if current_item:
              results_list.append({current_key: current_item})
            current_item = {}
            current_key = key
  
          current_item[column] = value
        if current_item:
          results_list.append({current_key: current_item})
        return results_list[offset:]
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during range_query'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def get_metadata(self, key):
      """ Retrieve a value from the datastore metadata table.
  
      Args:
        key: A string containing the key to fetch.
      Returns:
        A string containing the value or None if the key is not present.
      """
      statement = """
        SELECT {value} FROM "{table}"
        WHERE {key} = %s
        AND {column} = %s
      """.format(
        value=ThriftColumn.VALUE,
        table=dbconstants.DATASTORE_METADATA_TABLE,
        key=ThriftColumn.KEY,
        column=ThriftColumn.COLUMN_NAME
      )
      try:
        results = self.session.execute(statement, (bytearray(key), key))
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Unable to fetch {} from datastore metadata'.format(key)
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
      try:
        return results[0].value
      except IndexError:
        return None
  
    def set_metadata(self, key, value):
      """ Set a datastore metadata value.
  
      Args:
        key: A string containing the key to set.
        value: A string containing the value to set.
      """
      if not isinstance(key, str):
        raise TypeError('key should be a string')
  
      if not isinstance(value, str):
        raise TypeError('value should be a string')
  
      statement = """
        INSERT INTO "{table}" ({key}, {column}, {value})
        VALUES (%(key)s, %(column)s, %(value)s)
      """.format(
        table=dbconstants.DATASTORE_METADATA_TABLE,
        key=ThriftColumn.KEY,
        column=ThriftColumn.COLUMN_NAME,
        value=ThriftColumn.VALUE
      )
      parameters = {'key': bytearray(key),
                    'column': key,
                    'value': bytearray(value)}
      try:
        self.session.execute(statement, parameters)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Unable to set datastore metadata for {}'.format(key)
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
      except cassandra.InvalidRequest:
        self.create_table(dbconstants.DATASTORE_METADATA_TABLE,
                          dbconstants.DATASTORE_METADATA_SCHEMA)
        self.session.execute(statement, parameters)
  
    def get_indices(self, app_id):
      """ Gets the indices of the given application.
  
      Args:
        app_id: Name of the application.
      Returns:
        Returns a list of encoded entity_pb.CompositeIndex objects.
      """
      start_key = dbconstants.KEY_DELIMITER.join([app_id, 'index', ''])
      end_key = dbconstants.KEY_DELIMITER.join(
        [app_id, 'index', dbconstants.TERMINATING_STRING])
      result = self.range_query(
        dbconstants.METADATA_TABLE,
        dbconstants.METADATA_SCHEMA,
        start_key,
        end_key,
        dbconstants.MAX_NUMBER_OF_COMPOSITE_INDEXES,
        offset=0,
        start_inclusive=True,
        end_inclusive=True)
      list_result = []
      for list_item in result:
        for key, value in list_item.iteritems():
          list_result.append(value['data'])
      return list_result
  
    def valid_data_version(self):
      """ Checks whether or not the data layout can be used.
  
      Returns:
        A boolean.
      """
      try:
        version = self.get_metadata(VERSION_INFO_KEY)
      except cassandra.InvalidRequest:
        return False
  
      return version is not None and float(version) == EXPECTED_DATA_VERSION
  
files: AppDB/appscale/datastore/cassandra_env/cassandra_interface.py
files: # Programmer: Navraj Chohan <nlake44@gmail.com>
  
  """
   Cassandra Interface for AppScale
  """
  import cassandra
  import logging
  import sys
  import time
  
  from cassandra.cluster import Cluster
  from cassandra.concurrent import execute_concurrent
  from cassandra.policies import FallthroughRetryPolicy
  from cassandra.policies import RetryPolicy
  from cassandra.query import BatchStatement
  from cassandra.query import ConsistencyLevel
  from cassandra.query import SimpleStatement
  from cassandra.query import ValueSequence
  from .. import dbconstants
  from .. import helper_functions
  from ..dbconstants import AppScaleDBConnectionError
  from ..dbconstants import TxnActions
  from ..dbinterface import AppDBInterface
  from ..unpackaged import APPSCALE_LIB_DIR
  from ..utils import clean_app_id
  from ..utils import encode_index_pb
  from ..utils import get_composite_index_keys
  from ..utils import get_composite_indexes_rows
  from ..utils import get_entity_key
  from ..utils import get_entity_kind
  from ..utils import get_index_kv_from_tuple
  from ..utils import get_kind_key
  
  sys.path.append(APPSCALE_LIB_DIR)
  import appscale_info
  
  # The directory Cassandra is installed to.
  CASSANDRA_INSTALL_DIR = '/opt/cassandra'
  
  # Full path for the nodetool binary.
  NODE_TOOL = '{}/cassandra/bin/nodetool'.format(CASSANDRA_INSTALL_DIR)
  
  # The keyspace used for all tables
  KEYSPACE = "Keyspace1"
  
  # Cassandra watch name.
  CASSANDRA_MONIT_WATCH_NAME = "cassandra-9999"
  
  # The number of times to retry connecting to Cassandra.
  INITIAL_CONNECT_RETRIES = 20
  
  # The data layout version that the datastore expects.
  EXPECTED_DATA_VERSION = 1.0
  
  # The metadata key for the data layout version.
  VERSION_INFO_KEY = 'version'
  
  # The metadata key indicating that the database has been primed.
  PRIMED_KEY = 'primed'
  
  # The size in bytes that a batch must be to use the batches table.
  LARGE_BATCH_THRESHOLD = 5 << 10
  
  
  def batch_size(batch):
    """ Calculates the size of a batch.
  
    Args:
      batch: A list of dictionaries representing mutations.
    Returns:
      An integer specifying the size in bytes of the batch.
    """
    size = 0
    for mutation in batch:
      size += len(mutation['key'])
      if 'values' in mutation:
        for value in mutation['values'].values():
          size += len(value)
    return size
  
  
  def deletions_for_entity(entity, composite_indices=()):
    """ Get a list of deletions needed across tables for deleting an entity.
  
    Args:
      entity: An entity object.
      composite_indices: A list or tuple of composite indices.
    Returns:
      A list of dictionaries representing mutation operations.
    """
    deletions = []
    app_id = clean_app_id(entity.key().app())
    namespace = entity.key().name_space()
    prefix = dbconstants.KEY_DELIMITER.join([app_id, namespace])
  
    asc_rows = get_index_kv_from_tuple([(prefix, entity)])
    for entry in asc_rows:
      deletions.append({'table': dbconstants.ASC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': TxnActions.DELETE})
  
    dsc_rows = get_index_kv_from_tuple(
      [(prefix, entity)], reverse=True)
    for entry in dsc_rows:
      deletions.append({'table': dbconstants.DSC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': TxnActions.DELETE})
  
    for key in get_composite_indexes_rows([entity], composite_indices):
      deletions.append({'table': dbconstants.COMPOSITE_TABLE,
                        'key': key,
                        'operation': TxnActions.DELETE})
  
    entity_key = get_entity_key(prefix, entity.key().path())
    deletions.append({'table': dbconstants.APP_ENTITY_TABLE,
                      'key': entity_key,
                      'operation': TxnActions.DELETE})
  
    kind_key = get_kind_key(prefix, entity.key().path())
    deletions.append({'table': dbconstants.APP_KIND_TABLE,
                      'key': kind_key,
                      'operation': TxnActions.DELETE})
  
    return deletions
  
  
  def index_deletions(old_entity, new_entity, composite_indices=()):
    """ Get a list of index deletions needed for updating an entity. For changing
    an existing entity, this involves examining the property list of both
    entities to see which index entries need to be removed.
  
    Args:
      old_entity: An entity object.
      new_entity: An entity object.
      composite_indices: A list or tuple of composite indices.
    Returns:
      A list of dictionaries representing mutation operations.
    """
    deletions = []
    app_id = clean_app_id(old_entity.key().app())
    namespace = old_entity.key().name_space()
    kind = get_entity_kind(old_entity.key())
    entity_key = str(encode_index_pb(old_entity.key().path()))
  
    new_props = {}
    for prop in new_entity.property_list():
      if prop.name() not in new_props:
        new_props[prop.name()] = []
      new_props[prop.name()].append(prop)
  
    changed_props = {}
    for prop in old_entity.property_list():
      if prop.name() in new_props and prop in new_props[prop.name()]:
        continue
  
      if prop.name() not in changed_props:
        changed_props[prop.name()] = []
      changed_props[prop.name()].append(prop)
  
      value = str(encode_index_pb(prop.value()))
  
      key = dbconstants.KEY_DELIMITER.join(
        [app_id, namespace, kind, prop.name(), value, entity_key])
      deletions.append({'table': dbconstants.ASC_PROPERTY_TABLE,
                        'key': key,
                        'operation': TxnActions.DELETE})
  
      reverse_key = dbconstants.KEY_DELIMITER.join(
        [app_id, namespace, kind, prop.name(),
         helper_functions.reverse_lex(value), entity_key])
      deletions.append({'table': dbconstants.DSC_PROPERTY_TABLE,
                        'key': reverse_key,
                        'operation': TxnActions.DELETE})
  
    changed_prop_names = set(changed_props.keys())
    for index in composite_indices:
      if index.definition().entity_type() != kind:
        continue
  
      index_props = set(prop.name() for prop
                        in index.definition().property_list())
      if index_props.isdisjoint(changed_prop_names):
        continue
  
      old_entries = set(get_composite_index_keys(index, old_entity))
      new_entries = set(get_composite_index_keys(index, new_entity))
      for entry in (old_entries - new_entries):
        deletions.append({'table': dbconstants.COMPOSITE_TABLE,
                          'key': entry,
                          'operation': TxnActions.DELETE})
  
    return deletions
  
  
  def mutations_for_entity(entity, txn, current_value=None,
                           composite_indices=()):
    """ Get a list of mutations needed across tables for an entity change.
  
    Args:
      entity: An entity object.
      txn: A transaction ID handler.
      current_value: The entity object currently stored.
      composite_indices: A list of composite indices for the entity kind.
    Returns:
      A list of dictionaries representing mutations.
    """
    mutations = []
    if current_value is not None:
      mutations.extend(
        index_deletions(current_value, entity, composite_indices))
  
    app_id = clean_app_id(entity.key().app())
    namespace = entity.key().name_space()
    encoded_path = str(encode_index_pb(entity.key().path()))
    prefix = dbconstants.KEY_DELIMITER.join([app_id, namespace])
    entity_key = dbconstants.KEY_DELIMITER.join([prefix, encoded_path])
    entity_value = {dbconstants.APP_ENTITY_SCHEMA[0]: entity.Encode(),
                    dbconstants.APP_ENTITY_SCHEMA[1]: str(txn)}
    mutations.append({'table': dbconstants.APP_ENTITY_TABLE,
                      'key': entity_key,
                      'operation': TxnActions.PUT,
                      'values': entity_value})
  
    reference_value = {'reference': entity_key}
  
    kind_key = get_kind_key(prefix, entity.key().path())
    mutations.append({'table': dbconstants.APP_KIND_TABLE,
                      'key': kind_key,
                      'operation': TxnActions.PUT,
                      'values': reference_value})
  
    asc_rows = get_index_kv_from_tuple([(prefix, entity)])
    for entry in asc_rows:
      mutations.append({'table': dbconstants.ASC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': TxnActions.PUT,
                        'values': reference_value})
  
    dsc_rows = get_index_kv_from_tuple([(prefix, entity)], reverse=True)
    for entry in dsc_rows:
      mutations.append({'table': dbconstants.DSC_PROPERTY_TABLE,
                        'key': entry[0],
                        'operation': TxnActions.PUT,
                        'values': reference_value})
  
    for key in get_composite_indexes_rows([entity], composite_indices):
      mutations.append({'table': dbconstants.COMPOSITE_TABLE,
                        'key': key,
                        'operation': TxnActions.PUT,
                        'values': reference_value})
  
    return mutations
  
  
  class FailedBatch(Exception):
    pass
  
  
  class IdempotentRetryPolicy(RetryPolicy):
    """ A policy used for retrying idempotent statements. """
    def on_read_timeout(self, query, consistency, required_responses,
                        received_responses, data_retrieved, retry_num):
      """ This is called when a ReadTimeout occurs.
  
      Args:
        query: A statement that timed out.
        consistency: The consistency level of the statement.
        required_responses: The number of responses required.
        received_responses: The number of responses received.
        data_retrieved: Indicates whether any responses contained data.
        retry_num: The number of times the statement has been tried.
      """
      if retry_num >= 5:
        return self.RETHROW, None
      else:
        return self.RETRY, consistency
  
    def on_write_timeout(self, query, consistency, write_type,
                         required_responses, received_responses, retry_num):
      """ This is called when a WriteTimeout occurs.
  
      Args:
        query: A statement that timed out.
        consistency: The consistency level of the statement.
        required_responses: The number of responses required.
        received_responses: The number of responses received.
        data_retrieved: Indicates whether any responses contained data.
        retry_num: The number of times the statement has been tried.
        """
      if retry_num >= 5:
        return self.RETHROW, None
      else:
        return self.RETRY, consistency
  
  
  class ThriftColumn(object):
    """ Columns created by default with thrift interface. """
    KEY = 'key'
    COLUMN_NAME = 'column1'
    VALUE = 'value'
  
  
  class DatastoreProxy(AppDBInterface):
    """ 
      Cassandra implementation of the AppDBInterface
    """
    def __init__(self, log_level=logging.INFO):
      """
      Constructor.
      """
      class_name = self.__class__.__name__
      self.logger = logging.getLogger(class_name)
      self.logger.setLevel(log_level)
      self.logger.info('Starting {}'.format(class_name))
  
      self.hosts = appscale_info.get_db_ips()
      self.retry_policy = IdempotentRetryPolicy()
      self.no_retries = FallthroughRetryPolicy()
  
      remaining_retries = INITIAL_CONNECT_RETRIES
      while True:
        try:
          self.cluster = Cluster(self.hosts,
                                 default_retry_policy=self.retry_policy)
          self.session = self.cluster.connect(KEYSPACE)
          break
        except cassandra.cluster.NoHostAvailable as connection_error:
          remaining_retries -= 1
          if remaining_retries < 0:
            raise connection_error
          time.sleep(3)
  
      self.session.default_consistency_level = ConsistencyLevel.QUORUM
  
    def close(self):
      """ Close all sessions and connections to Cassandra. """
      self.cluster.shutdown()
  
    def batch_get_entity(self, table_name, row_keys, column_names):
      """
      Takes in batches of keys and retrieves their corresponding rows.
      
      Args:
        table_name: The table to access
        row_keys: A list of keys to access
        column_names: A list of columns to access
      Returns:
        A dictionary of rows and columns/values of those rows. The format 
        looks like such: {key:{column_name:value,...}}
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the batch_get could not be performed due to
          an error with Cassandra.
      """
      if not isinstance(table_name, str): raise TypeError("Expected a str")
      if not isinstance(column_names, list): raise TypeError("Expected a list")
      if not isinstance(row_keys, list): raise TypeError("Expected a list")
  
      row_keys_bytes = [bytearray(row_key) for row_key in row_keys]
  
      statement = 'SELECT * FROM "{table}" '\
                  'WHERE {key} IN %s and {column} IN %s'.format(
                    table=table_name,
                    key=ThriftColumn.KEY,
                    column=ThriftColumn.COLUMN_NAME,
                  )
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
      parameters = (ValueSequence(row_keys_bytes), ValueSequence(column_names))
  
      try:
        results = self.session.execute(query, parameters=parameters)
  
        results_dict = {row_key: {} for row_key in row_keys}
        for (key, column, value) in results:
          if key not in results_dict:
            results_dict[key] = {}
          results_dict[key][column] = value
  
        return results_dict
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_get_entity'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def batch_put_entity(self, table_name, row_keys, column_names, cell_values,
                         ttl=None):
      """
      Allows callers to store multiple rows with a single call. A row can 
      have multiple columns and values with them. We refer to each row as 
      an entity.
     
      Args: 
        table_name: The table to mutate
        row_keys: A list of keys to store on
        column_names: A list of columns to mutate
        cell_values: A dict of key/value pairs
        ttl: The number of seconds to keep the row.
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the batch_put could not be performed due to
          an error with Cassandra.
      """
      if not isinstance(table_name, str):
        raise TypeError("Expected a str")
      if not isinstance(column_names, list):
        raise TypeError("Expected a list")
      if not isinstance(row_keys, list):
        raise TypeError("Expected a list")
      if not isinstance(cell_values, dict):
        raise TypeError("Expected a dict")
  
      insert_str = """
        INSERT INTO "{table}" ({key}, {column}, {value})
        VALUES (?, ?, ?)
      """.format(table=table_name,
                 key=ThriftColumn.KEY,
                 column=ThriftColumn.COLUMN_NAME,
                 value=ThriftColumn.VALUE)
  
      if ttl is not None:
        insert_str += 'USING TTL {}'.format(ttl)
  
      statement = self.session.prepare(insert_str)
  
      statements_and_params = []
      for row_key in row_keys:
        for column in column_names:
          params = (bytearray(row_key), column,
                    bytearray(cell_values[row_key][column]))
          statements_and_params.append((statement, params))
  
      try:
        execute_concurrent(self.session, statements_and_params,
                           raise_on_first_error=True)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_put_entity'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def prepare_insert(self, table):
      """ Prepare an insert statement.
  
      Args:
        table: A string containing the table name.
      Returns:
        A PreparedStatement object.
      """
      statement = """
        INSERT INTO "{table}" ({key}, {column}, {value})
        VALUES (?, ?, ?)
      """.format(table=table,
                 key=ThriftColumn.KEY,
                 column=ThriftColumn.COLUMN_NAME,
                 value=ThriftColumn.VALUE)
      return self.session.prepare(statement)
  
    def prepare_delete(self, table):
      """ Prepare a delete statement.
  
      Args:
        table: A string containing the table name.
      Returns:
        A PreparedStatement object.
      """
      statement = """
        DELETE FROM "{table}" WHERE {key} = ?
      """.format(table=table, key=ThriftColumn.KEY)
      return self.session.prepare(statement)
  
    def _normal_batch(self, mutations):
      """ Use Cassandra's native batch statement to apply mutations atomically.
  
      Args:
        mutations: A list of dictionaries representing mutations.
      """
      self.logger.debug('Normal batch: {} mutations'.format(len(mutations)))
      batch = BatchStatement(consistency_level=ConsistencyLevel.QUORUM,
                             retry_policy=self.retry_policy)
      prepared_statements = {'insert': {}, 'delete': {}}
      for mutation in mutations:
        table = mutation['table']
        if mutation['operation'] == TxnActions.PUT:
          if table not in prepared_statements['insert']:
            prepared_statements['insert'][table] = self.prepare_insert(table)
          values = mutation['values']
          for column in values:
            batch.add(
              prepared_statements['insert'][table],
              (bytearray(mutation['key']), column, bytearray(values[column]))
            )
        elif mutation['operation'] == TxnActions.DELETE:
          if table not in prepared_statements['delete']:
            prepared_statements['delete'][table] = self.prepare_delete(table)
          batch.add(
            prepared_statements['delete'][table],
            (bytearray(mutation['key']),)
          )
  
      try:
        self.session.execute(batch)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_mutate'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def apply_mutations(self, mutations):
      """ Apply mutations across tables.
  
      Args:
        mutations: A list of dictionaries representing mutations.
      """
      prepared_statements = {'insert': {}, 'delete': {}}
      statements_and_params = []
      for mutation in mutations:
        table = mutation['table']
        if mutation['operation'] == TxnActions.PUT:
          if table not in prepared_statements['insert']:
            prepared_statements['insert'][table] = self.prepare_insert(table)
          values = mutation['values']
          for column in values:
            params = (bytearray(mutation['key']), column,
                      bytearray(values[column]))
            statements_and_params.append(
              (prepared_statements['insert'][table], params))
        elif mutation['operation'] == TxnActions.DELETE:
          if table not in prepared_statements['delete']:
            prepared_statements['delete'][table] = self.prepare_delete(table)
          params = (bytearray(mutation['key']),)
          statements_and_params.append(
            (prepared_statements['delete'][table], params))
  
      execute_concurrent(self.session, statements_and_params,
                         raise_on_first_error=True)
  
    def _large_batch(self, app, mutations, entity_changes, txn):
      """ Insert or delete multiple rows across tables in an atomic statement.
  
      Args:
        app: A string containing the application ID.
        mutations: A list of dictionaries representing mutations.
        entity_changes: A list of changes at the entity level.
        txn: A transaction ID handler.
      Raises:
        FailedBatch if a concurrent process modifies the batch status.
        AppScaleDBConnectionError if a database connection error was encountered.
      """
      self.logger.debug('Large batch: transaction {}, {} mutations'.
                        format(txn, len(mutations)))
      set_status = """
        INSERT INTO batch_status (app, transaction, applied)
        VALUES (%(app)s, %(transaction)s, False)
        IF NOT EXISTS
      """
      parameters = {'app': app, 'transaction': txn}
      result = self.session.execute(set_status, parameters)
      if not result.was_applied:
        raise FailedBatch('A batch for transaction {} already exists'.
                          format(txn))
  
      insert_item = """
        INSERT INTO batches (app, transaction, namespace, path,
                             old_value, new_value)
        VALUES (?, ?, ?, ?, ?, ?)
      """
      insert_statement = self.session.prepare(insert_item)
  
      statements_and_params = []
      for entity_change in entity_changes:
        old_value = None
        if entity_change['old'] is not None:
          old_value = bytearray(entity_change['old'].Encode())
        new_value = None
        if entity_change['new'] is not None:
          new_value = bytearray(entity_change['new'].Encode())
  
        parameters = (app, txn, entity_change['key'].name_space(),
                      bytearray(entity_change['key'].path().Encode()), old_value,
                      new_value)
        statements_and_params.append((insert_statement, parameters))
  
      try:
        execute_concurrent(self.session, statements_and_params,
                           raise_on_first_error=True)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during large batch'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
      update_status = """
        UPDATE batch_status
        SET applied = True
        WHERE app = %(app)s
        AND transaction = %(transaction)s
        IF applied = False
      """
      parameters = {'app': app, 'transaction': txn}
      result = self.session.execute(update_status, parameters)
      if not result.was_applied:
        raise FailedBatch('Another process modified batch for transaction {}'.
                          format(txn))
  
      try:
        self.apply_mutations(mutations)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during large batch'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
      clear_batch = """
        DELETE FROM batches
        WHERE app = %(app)s AND transaction = %(transaction)s
      """
      parameters = {'app': app, 'transaction': txn}
      self.session.execute(clear_batch, parameters)
  
      clear_status = """
        DELETE FROM batch_status
        WHERE app = %(app)s and transaction = %(transaction)s
      """
      parameters = {'app': app, 'transaction': txn}
      self.session.execute(clear_status, parameters)
  
    def batch_mutate(self, app, mutations, entity_changes, txn):
      """ Insert or delete multiple rows across tables in an atomic statement.
  
      Args:
        app: A string containing the application ID.
        mutations: A list of dictionaries representing mutations.
        entity_changes: A list of changes at the entity level.
        txn: A transaction ID handler.
      """
      size = batch_size(mutations)
      self.logger.debug('batch_size: {}'.format(size))
      if size > LARGE_BATCH_THRESHOLD:
        self._large_batch(app, mutations, entity_changes, txn)
      else:
        self._normal_batch(mutations)
  
    def batch_delete(self, table_name, row_keys, column_names=()):
      """
      Remove a set of rows corresponding to a set of keys.
       
      Args:
        table_name: Table to delete rows from
        row_keys: A list of keys to remove
        column_names: Not used
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the batch_delete could not be performed due
          to an error with Cassandra.
      """ 
      if not isinstance(table_name, str): raise TypeError("Expected a str")
      if not isinstance(row_keys, list): raise TypeError("Expected a list")
  
      row_keys_bytes = [bytearray(row_key) for row_key in row_keys]
  
      statement = 'DELETE FROM "{table}" WHERE {key} IN %s'.\
        format(
          table=table_name,
          key=ThriftColumn.KEY
        )
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
      parameters = (ValueSequence(row_keys_bytes),)
  
      try:
        self.session.execute(query, parameters=parameters)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during batch_delete'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def delete_table(self, table_name):
      """ 
      Drops a given table (aka column family in Cassandra)
    
      Args:
        table_name: A string name of the table to drop
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the delete_table could not be performed due
          to an error with Cassandra.
      """
      if not isinstance(table_name, str): raise TypeError("Expected a str")
  
      statement = 'DROP TABLE IF EXISTS "{table}"'.format(table=table_name)
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
  
      try:
        self.session.execute(query)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during delete_table'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def create_table(self, table_name, column_names):
      """ 
      Creates a table if it doesn't already exist.
      
      Args:
        table_name: The column family name
        column_names: Not used but here to match the interface
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the create_table could not be performed due
          to an error with Cassandra.
      """
      if not isinstance(table_name, str): raise TypeError("Expected a str")
      if not isinstance(column_names, list): raise TypeError("Expected a list")
  
      statement = 'CREATE TABLE IF NOT EXISTS "{table}" ('\
          '{key} blob,'\
          '{column} text,'\
          '{value} blob,'\
          'PRIMARY KEY ({key}, {column})'\
        ') WITH COMPACT STORAGE'.format(
          table=table_name,
          key=ThriftColumn.KEY,
          column=ThriftColumn.COLUMN_NAME,
          value=ThriftColumn.VALUE
        )
      query = SimpleStatement(statement, retry_policy=self.no_retries)
  
      try:
        self.session.execute(query)
      except cassandra.OperationTimedOut:
        logging.warning('Encountered an operation timeout while creating a '
                        'table. Waiting 1 minute for schema to settle.')
        time.sleep(60)
        raise AppScaleDBConnectionError('Exception during create_table')
      except (error for error in dbconstants.TRANSIENT_CASSANDRA_ERRORS
              if error != cassandra.OperationTimedOut):
        message = 'Exception during create_table'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def range_query(self,
                    table_name,
                    column_names, 
                    start_key, 
                    end_key, 
                    limit, 
                    offset=0, 
                    start_inclusive=True, 
                    end_inclusive=True,
                    keys_only=False):
      """ 
      Gets a dense range ordered by keys. Returns an ordered list of 
      a dictionary of [key:{column1:value1, column2:value2},...]
      or a list of keys if keys only.
       
      Args:
        table_name: Name of table to access
        column_names: Columns which get returned within the key range
        start_key: String for which the query starts at
        end_key: String for which the query ends at
        limit: Maximum number of results to return
        offset: Cuts off these many from the results [offset:]
        start_inclusive: Boolean if results should include the start_key
        end_inclusive: Boolean if results should include the end_key
        keys_only: Boolean if to only keys and not values
      Raises:
        TypeError: If an argument passed in was not of the expected type.
        AppScaleDBConnectionError: If the range_query could not be performed due
          to an error with Cassandra.
      Returns:
        An ordered list of dictionaries of key=>columns/values
      """
      if not isinstance(table_name, str):
        raise TypeError('table_name must be a string')
      if not isinstance(column_names, list):
        raise TypeError('column_names must be a list')
      if not isinstance(start_key, str):
        raise TypeError('start_key must be a string')
      if not isinstance(end_key, str):
        raise TypeError('end_key must be a string')
      if not isinstance(limit, (int, long)) and limit is not None:
        raise TypeError('limit must be int, long, or NoneType')
      if not isinstance(offset, (int, long)):
        raise TypeError('offset must be int or long')
  
      if start_inclusive:
        gt_compare = '>='
      else:
        gt_compare = '>'
  
      if end_inclusive:
        lt_compare = '<='
      else:
        lt_compare = '<'
  
      query_limit = ''
      if limit is not None:
        query_limit = 'LIMIT {}'.format(len(column_names) * limit)
  
      statement = """
        SELECT * FROM "{table}" WHERE
        token({key}) {gt_compare} %s AND
        token({key}) {lt_compare} %s AND
        {column} IN %s
        {limit}
        ALLOW FILTERING
      """.format(table=table_name,
                 key=ThriftColumn.KEY,
                 gt_compare=gt_compare,
                 lt_compare=lt_compare,
                 column=ThriftColumn.COLUMN_NAME,
                 limit=query_limit)
  
      query = SimpleStatement(statement, retry_policy=self.retry_policy)
      parameters = (bytearray(start_key), bytearray(end_key),
                    ValueSequence(column_names))
  
      try:
        results = self.session.execute(query, parameters=parameters)
  
        results_list = []
        current_item = {}
        current_key = None
        for (key, column, value) in results:
          if keys_only:
            results_list.append(key)
            continue
  
          if key != current_key:
            if current_item:
              results_list.append({current_key: current_item})
            current_item = {}
            current_key = key
  
          current_item[column] = value
        if current_item:
          results_list.append({current_key: current_item})
        return results_list[offset:]
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Exception during range_query'
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
    def get_metadata(self, key):
      """ Retrieve a value from the datastore metadata table.
  
      Args:
        key: A string containing the key to fetch.
      Returns:
        A string containing the value or None if the key is not present.
      """
      statement = """
        SELECT {value} FROM "{table}"
        WHERE {key} = %s
        AND {column} = %s
      """.format(
        value=ThriftColumn.VALUE,
        table=dbconstants.DATASTORE_METADATA_TABLE,
        key=ThriftColumn.KEY,
        column=ThriftColumn.COLUMN_NAME
      )
      try:
        results = self.session.execute(statement, (bytearray(key), key))
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Unable to fetch {} from datastore metadata'.format(key)
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
  
      try:
        return results[0].value
      except IndexError:
        return None
  
    def set_metadata(self, key, value):
      """ Set a datastore metadata value.
  
      Args:
        key: A string containing the key to set.
        value: A string containing the value to set.
      """
      if not isinstance(key, str):
        raise TypeError('key should be a string')
  
      if not isinstance(value, str):
        raise TypeError('value should be a string')
  
      statement = """
        INSERT INTO "{table}" ({key}, {column}, {value})
        VALUES (%(key)s, %(column)s, %(value)s)
      """.format(
        table=dbconstants.DATASTORE_METADATA_TABLE,
        key=ThriftColumn.KEY,
        column=ThriftColumn.COLUMN_NAME,
        value=ThriftColumn.VALUE
      )
      parameters = {'key': bytearray(key),
                    'column': key,
                    'value': bytearray(value)}
      try:
        self.session.execute(statement, parameters)
      except dbconstants.TRANSIENT_CASSANDRA_ERRORS:
        message = 'Unable to set datastore metadata for {}'.format(key)
        logging.exception(message)
        raise AppScaleDBConnectionError(message)
      except cassandra.InvalidRequest:
        self.create_table(dbconstants.DATASTORE_METADATA_TABLE,
                          dbconstants.DATASTORE_METADATA_SCHEMA)
        self.session.execute(statement, parameters)
  
    def get_indices(self, app_id):
      """ Gets the indices of the given application.
  
      Args:
        app_id: Name of the application.
      Returns:
        Returns a list of encoded entity_pb.CompositeIndex objects.
      """
      start_key = dbconstants.KEY_DELIMITER.join([app_id, 'index', ''])
      end_key = dbconstants.KEY_DELIMITER.join(
        [app_id, 'index', dbconstants.TERMINATING_STRING])
      result = self.range_query(
        dbconstants.METADATA_TABLE,
        dbconstants.METADATA_SCHEMA,
        start_key,
        end_key,
        dbconstants.MAX_NUMBER_OF_COMPOSITE_INDEXES,
        offset=0,
        start_inclusive=True,
        end_inclusive=True)
      list_result = []
      for list_item in result:
        for key, value in list_item.iteritems():
          list_result.append(value['data'])
      return list_result
  
    def valid_data_version(self):
      """ Checks whether or not the data layout can be used.
  
      Returns:
        A boolean.
      """
      try:
        version = self.get_metadata(VERSION_INFO_KEY)
      except cassandra.InvalidRequest:
        return False
  
      return version is not None and float(version) == EXPECTED_DATA_VERSION
  
files: AppDB/appscale/datastore/cassandra_env/cassandra_interface.py
files: """
   Datastore Constants
  """
  import cassandra.cluster
  
  SECRET_LOCATION = "/etc/appscale/secret.key"
  
  ERROR_DEFAULT = "DB_ERROR:"
  NONEXISTANT_TRANSACTION = "0"
  
  # The datastore's default HTTP port.
  DEFAULT_PORT = 4080
  
  # The datastore's default HTTPS port.
  DEFAULT_SSL_PORT = 8443
  
  # The lowest character to separate different fields in a row key.
  KEY_DELIMITER = '\x00'
  
  # The character used to separate kinds in an ancestry.
  KIND_SEPARATOR = '\x01'
  
  # HTTP code to indicate that the request is invalid.
  HTTP_BAD_REQUEST = 400
  
  # The length of an ID string. A constant length allows lexicographic ordering.
  ID_KEY_LENGTH = 10
  
  # The character between the kind and the ID/name of an entity.
  ID_SEPARATOR = ":"
  
  # The maximum number of composite indexes an application can have.
  MAX_NUMBER_OF_COMPOSITE_INDEXES = 1000
  
  # A string used to create end keys when doing range queries.
  TERMINATING_STRING = chr(255) * 500
  
  # Tombstone value for soft deletes.
  TOMBSTONE = "APPSCALE_SOFT_DELETE"
  
  TRANSIENT_CASSANDRA_ERRORS = (
    cassandra.Unavailable, cassandra.Timeout, cassandra.CoordinationFailure,
    cassandra.OperationTimedOut, cassandra.cluster.NoHostAvailable)
  
  # The database backends supported by the AppScale datastore.
  VALID_DATASTORES = ['cassandra']
  
  # Table names
  USERS_TABLE = "USERS__"
  APPS_TABLE = "APPS__"
  JOURNAL_TABLE = "JOURNAL__"
  
  ASC_PROPERTY_TABLE = "ASC_PROPERTY__"
  DSC_PROPERTY_TABLE = "DSC_PROPERTY__"
  COMPOSITE_TABLE = "COMPOSITE_INDEXES__"
  APP_ID_TABLE = "APP_IDS__"
  APP_ENTITY_TABLE = "ENTITIES__"
  APP_KIND_TABLE = "KINDS__"
  METADATA_TABLE = "METADATA__"
  DATASTORE_METADATA_TABLE = "DATASTORE_METADATA__"
  TRANSACTIONS_TABLE = 'TRANSACTIONS__'
  SCHEMA_TABLE = '__key__'
  
  INITIAL_TABLES = [ASC_PROPERTY_TABLE,
                    DSC_PROPERTY_TABLE,
                    APP_ID_TABLE,
                    APP_ENTITY_TABLE,
                    APP_KIND_TABLE,
                    COMPOSITE_TABLE,
                    METADATA_TABLE,
                    USERS_TABLE,
                    APPS_TABLE,
                    SCHEMA_TABLE,
                    DATASTORE_METADATA_TABLE,
                    TRANSACTIONS_TABLE]
  
  ###########################################
  # DB schemas for version 1 of the datastore
  ###########################################
  JOURNAL_SCHEMA = [
    "Encoded_Entity"]
  
  ENTITY_TABLE_SCHEMA = [
    "Encoded_Entity",
    "Txn_Num"]
  
  ###########################################
  # DB schema for version 2 of the datastore
  ###########################################
  
  # The schema of the table which holds the encoded entities
  APP_ENTITY_SCHEMA = [
    "entity",
    "txnID"]
  
  # Index tables store references are to entity table
  PROPERTY_SCHEMA = [
    "reference" ]
  APP_ID_SCHEMA = [
    "next_id" ]
  APP_KIND_SCHEMA = [
    "reference" ]
  COMPOSITE_SCHEMA = [
    "reference" ]
  METADATA_SCHEMA = [
    "data" ]
  
  USERS_SCHEMA = [
    "email",
    "pw",
    "date_creation", 
    "date_change",
    "date_last_login",
    "applications",
    "appdrop_rem_token",
    "appdrop_rem_token_exp",
    "visit_cnt",
    "cookie",
    "cookie_ip",
    "cookie_exp",
    "cksum",
    "enabled",
    "type",
    "is_cloud_admin",
    "capabilities" ]
  
  APPS_SCHEMA = [
    "name",
    "language",
    "version",
    "owner",
    "admins_list",
    "host",
    "port",
    "creation_date",
    "last_time_updated_date",
    "yaml_file",
    "cksum",
    "num_entries",
    "tar_ball",
    "enabled",
    "indexes" ]
  
  DATASTORE_METADATA_SCHEMA = [
    "version"]
  
  TRANSACTIONS_SCHEMA = [
    'operation',
    'operand',
    'exclude_indices'
  ]
  
  # All schema information for the keyspace is stored in the schema table.
  SCHEMA_TABLE_SCHEMA = ['schema']
  
  
  # Possible values in the 'action' column of the transaction table.
  class TxnActions(object):
    DELETE = '0'
    PUT = '1'
    GET = '2'
    ENQUEUE_TASK = '3'
  
  
  class Operations(object):
    """ Possible datastore operations on entities. """
    PUT = 'put'
    DELETE = 'delete'
  
  
  ###############################
  # Generic Datastore Exceptions
  ###############################
  class AppScaleDBError(Exception):
    """ Tossed for generic datastore errors
    """
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
  
    def __str__(self):
      return repr(self.value)
  
  class AppScaleDBConnectionError(Exception):
    """ Tossed when there is a bad connection
    """ 
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
    def __str__(self):
      return repr(self.value)
  
  class AppScaleMisconfiguredQuery(Exception):
    """ Tossed when a query is misconfigured
    """
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
    def __str__(self):
      return repr(self.value)
  
  class AppScaleBadArg(Exception):
    """ Bad Argument given for a function
    """
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
    def __str__(self):
      return repr(self.value)
  
  class ExcessiveTasks(Exception):
    """ Indicates that there are too many tasks for a transaction. """
    pass
  
files: AppDB/appscale/datastore/dbconstants.py
files: """
   Datastore Constants
  """
  import cassandra.cluster
  
  SECRET_LOCATION = "/etc/appscale/secret.key"
  
  ERROR_DEFAULT = "DB_ERROR:"
  NONEXISTANT_TRANSACTION = "0"
  
  # The datastore's default HTTP port.
  DEFAULT_PORT = 4080
  
  # The datastore's default HTTPS port.
  DEFAULT_SSL_PORT = 8443
  
  # The lowest character to separate different fields in a row key.
  KEY_DELIMITER = '\x00'
  
  # The character used to separate kinds in an ancestry.
  KIND_SEPARATOR = '\x01'
  
  # HTTP code to indicate that the request is invalid.
  HTTP_BAD_REQUEST = 400
  
  # The length of an ID string. A constant length allows lexicographic ordering.
  ID_KEY_LENGTH = 10
  
  # The character between the kind and the ID/name of an entity.
  ID_SEPARATOR = ":"
  
  # The maximum number of composite indexes an application can have.
  MAX_NUMBER_OF_COMPOSITE_INDEXES = 1000
  
  # A string used to create end keys when doing range queries.
  TERMINATING_STRING = chr(255) * 500
  
  # Tombstone value for soft deletes.
  TOMBSTONE = "APPSCALE_SOFT_DELETE"
  
  TRANSIENT_CASSANDRA_ERRORS = (
    cassandra.Unavailable, cassandra.Timeout, cassandra.CoordinationFailure,
    cassandra.OperationTimedOut, cassandra.cluster.NoHostAvailable)
  
  # The database backends supported by the AppScale datastore.
  VALID_DATASTORES = ['cassandra']
  
  # Table names
  USERS_TABLE = "USERS__"
  APPS_TABLE = "APPS__"
  JOURNAL_TABLE = "JOURNAL__"
  
  ASC_PROPERTY_TABLE = "ASC_PROPERTY__"
  DSC_PROPERTY_TABLE = "DSC_PROPERTY__"
  COMPOSITE_TABLE = "COMPOSITE_INDEXES__"
  APP_ID_TABLE = "APP_IDS__"
  APP_ENTITY_TABLE = "ENTITIES__"
  APP_KIND_TABLE = "KINDS__"
  METADATA_TABLE = "METADATA__"
  DATASTORE_METADATA_TABLE = "DATASTORE_METADATA__"
  TRANSACTIONS_TABLE = 'TRANSACTIONS__'
  SCHEMA_TABLE = '__key__'
  
  INITIAL_TABLES = [ASC_PROPERTY_TABLE,
                    DSC_PROPERTY_TABLE,
                    APP_ID_TABLE,
                    APP_ENTITY_TABLE,
                    APP_KIND_TABLE,
                    COMPOSITE_TABLE,
                    METADATA_TABLE,
                    USERS_TABLE,
                    APPS_TABLE,
                    SCHEMA_TABLE,
                    DATASTORE_METADATA_TABLE,
                    TRANSACTIONS_TABLE]
  
  ###########################################
  # DB schemas for version 1 of the datastore
  ###########################################
  JOURNAL_SCHEMA = [
    "Encoded_Entity"]
  
  ENTITY_TABLE_SCHEMA = [
    "Encoded_Entity",
    "Txn_Num"]
  
  ###########################################
  # DB schema for version 2 of the datastore
  ###########################################
  
  # The schema of the table which holds the encoded entities
  APP_ENTITY_SCHEMA = [
    "entity",
    "txnID"]
  
  # Index tables store references are to entity table
  PROPERTY_SCHEMA = [
    "reference" ]
  APP_ID_SCHEMA = [
    "next_id" ]
  APP_KIND_SCHEMA = [
    "reference" ]
  COMPOSITE_SCHEMA = [
    "reference" ]
  METADATA_SCHEMA = [
    "data" ]
  
  USERS_SCHEMA = [
    "email",
    "pw",
    "date_creation", 
    "date_change",
    "date_last_login",
    "applications",
    "appdrop_rem_token",
    "appdrop_rem_token_exp",
    "visit_cnt",
    "cookie",
    "cookie_ip",
    "cookie_exp",
    "cksum",
    "enabled",
    "type",
    "is_cloud_admin",
    "capabilities" ]
  
  APPS_SCHEMA = [
    "name",
    "language",
    "version",
    "owner",
    "admins_list",
    "host",
    "port",
    "creation_date",
    "last_time_updated_date",
    "yaml_file",
    "cksum",
    "num_entries",
    "tar_ball",
    "enabled",
    "indexes" ]
  
  DATASTORE_METADATA_SCHEMA = [
    "version"]
  
  TRANSACTIONS_SCHEMA = [
    'operation',
    'operand',
    'exclude_indices'
  ]
  
  # All schema information for the keyspace is stored in the schema table.
  SCHEMA_TABLE_SCHEMA = ['schema']
  
  
  # Possible values in the 'action' column of the transaction table.
  class TxnActions(object):
    DELETE = '0'
    PUT = '1'
    GET = '2'
    ENQUEUE_TASK = '3'
  
  
  ###############################
  # Generic Datastore Exceptions
  ###############################
  class AppScaleDBError(Exception):
    """ Tossed for generic datastore errors
    """
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
  
    def __str__(self):
      return repr(self.value)
  
  class AppScaleDBConnectionError(Exception):
    """ Tossed when there is a bad connection
    """ 
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
    def __str__(self):
      return repr(self.value)
  
  class AppScaleMisconfiguredQuery(Exception):
    """ Tossed when a query is misconfigured
    """
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
    def __str__(self):
      return repr(self.value)
  
  class AppScaleBadArg(Exception):
    """ Bad Argument given for a function
    """
    def __init__(self, value):
      Exception.__init__(self, value)
      self.value = value
    def __str__(self):
      return repr(self.value)
  
  class ExcessiveTasks(Exception):
    """ Indicates that there are too many tasks for a transaction. """
    pass
  
files: AppDB/appscale/datastore/dbconstants.py
files: #!/usr/bin/env python
  # Programmer: Navraj Chohan <nlake44@gmail.com>
  
  import sys
  import unittest
  
  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 ID_KEY_LENGTH
  from appscale.datastore.dbconstants import JOURNAL_SCHEMA
  from appscale.datastore.dbconstants import TOMBSTONE
  from appscale.datastore.cassandra_env.cassandra_interface import DatastoreProxy
  from appscale.datastore.cassandra_env.cassandra_interface import\
    deletions_for_entity
  from appscale.datastore.cassandra_env.cassandra_interface import\
    index_deletions
  from appscale.datastore.cassandra_env.cassandra_interface import\
    mutations_for_entity
  from appscale.datastore.unpackaged import APPSCALE_LIB_DIR
  from appscale.datastore.unpackaged import APPSCALE_PYTHON_APPSERVER
  from appscale.datastore.utils import encode_index_pb
  from appscale.datastore.utils import get_entity_key
  from appscale.datastore.utils import get_entity_kind
  from appscale.datastore.utils import get_index_key_from_params
  from appscale.datastore.utils import get_index_kv_from_tuple
  from appscale.datastore.utils import get_kind_key
  from appscale.datastore.zkappscale.zktransaction import TX_TIMEOUT
  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
  
  sys.path.append(APPSCALE_LIB_DIR)
  import appscale_info
  
  
  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):
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("increment_and_get_counter").and_return(0,1000)
      return zookeeper
  
    def test_get_entity_kind(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      item = Item(name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_entity_kind(key), "Item")
  
    def test_kind_key(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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)
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, 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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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)
      dd = DatastoreDistributed(db_batch, 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)
      dd = DatastoreDistributed(db_batch, 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)
      dd = DatastoreDistributed(db_batch, 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()
      dd = DatastoreDistributed(db_batch, self.get_zookeeper())
      dd.insert_composite_indexes([ent], [composite_index])
  
    def test_allocate_ids(self):
      PREFIX = "x"
      BATCH_SIZE = 1000
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, self.get_zookeeper())
      self.assertEquals(dd.allocate_ids(PREFIX, BATCH_SIZE), (1, 1000))
  
      dd = DatastoreDistributed(db_batch, self.get_zookeeper())
      self.assertEquals(dd.allocate_ids(PREFIX, None, max_id=1000), (1, 1000))
  
      try:
        # Unable to use self.assertRaises because of the optional argrument max_id
        ed = DatastoreDistributed(db_batch, self.get_zookeeper())
        dd.allocate_ids(PREFIX, BATCH_SIZE, max_id=10)
        raise "Allocate IDs should not let you set max_id and size"
      except ValueError:
        pass
  
    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()}})
  
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, 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)
      zookeeper = flexmock()
      zookeeper.should_receive("release_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, 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)
      zookeeper = flexmock()
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("notify_failed_transaction").and_return(True)
      dd = DatastoreDistributed(db_batch, 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()
      db_batch.should_receive('valid_data_version').and_return(True)
  
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("get_transaction_id").and_return(1)
  
      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')
      dd = DatastoreDistributed(db_batch, 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()
      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)
  
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").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')
      dd = DatastoreDistributed(db_batch, zookeeper)
  
      # Make sure it does not throw an exception
      txn_hash = {entity_key1: 1, entity_key2: 1}
      dd.put_entities(app_id, entity_list, txn_hash)
  
    def test_acquire_locks_for_trans(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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, None)
      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, None)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True)
  
      dd = DatastoreDistributed(db_batch, None)
      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()
      zookeeper.should_receive("acquire_lock").once()
      dd = DatastoreDistributed(db_batch, 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()
      zookeeper.should_receive("acquire_lock").once().and_raise(ZKTransactionException)
      zookeeper.should_receive("notify_failed_transaction").once()
      dd = DatastoreDistributed(db_batch, 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_acquire_locks_for_nontrans(self):
      app_id = 'test'
      PREFIX = 'x\x01'
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("get_transaction_id").and_return(1).and_return(2)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return({PREFIX:{}})
      db_batch.should_receive("batch_delete").and_return(None)
      dd = DatastoreDistributed(db_batch, zookeeper) 
      entity_proto1 = self.get_new_entity_proto(
        app_id, 'test_kind', "bob", "prop1name", "prop1val", ns="blah")
      entity_proto2 = self.get_new_entity_proto(
        app_id, "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entity_list = [entity_proto1, entity_proto2]
      self.assertEquals({'test\x00blah\x00test_kind:bob\x01': 2, 'test\x00blah\x00test_kind:nancy\x01': 1}, 
                        dd.acquire_locks_for_nontrans("test", entity_list))
  
    def test_delete_entities(self):
      app_id = 'test'
      entity_proto1 = self.get_new_entity_proto(
        app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      row_key = "test\x00blah\x00test_kind:bob\x01"
      row_values = {row_key: {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                              APP_ENTITY_SCHEMA[1]: '1'}}
  
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(row_values)
      db_batch.should_receive('batch_mutate')
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
  
      row_keys = [entity_proto1.key()]
  
      txn_hash = {row_key: 1}
      dd.delete_entities(app_id, row_keys, txn_hash)
       
    def test_release_put_locks_for_nontrans(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
      db_batch.should_receive("batch_delete").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "bob", "prop1name", 
                                                "prop1val", ns="blah")
      entity_proto2 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop2val", ns="blah")
      entities = [entity_proto1, entity_proto2]
      dd.release_locks_for_nontrans("test", entities, 
                    {'test\x00blah\x00test_kind:bob\x01': 1, 'test\x00blah\x00test_kind:nancy\x01': 2})
   
    def test_root_key_from_entity_key(self):
      zookeeper = flexmock()
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
      self.assertEquals("test\x00blah\x00test_kind:bob\x01", 
                        dd.get_root_key_from_entity_key("test\x00blah\x00test_kind:bob\x01nancy\x01"))
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop2val", ns="blah")
      self.assertEquals("test\x00blah\x00test_kind:nancy\x01", 
        dd.get_root_key_from_entity_key(entity_proto1.key()))
  
    def test_dynamic_get(self):
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop2val", ns="blah")
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
  
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key() 
      key.MergeFrom(entity_key)
      get_resp = datastore_pb.GetResponse()
      
      dd.dynamic_get("test", get_req, get_resp)     
      self.assertEquals(get_resp.entity_size(), 1)
  
      # Now test while in a transaction
      get_resp = datastore_pb.GetResponse()
      get_req.mutable_transaction().set_handle(1)
      dd.dynamic_get("test", get_req, get_resp)     
      self.assertEquals(get_resp.entity_size(), 1)
  
    def test_ancestor_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key() 
      key.MergeFrom(entity_key)
      ancestor.MergeFrom(entity_key)
      
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      dd = DatastoreDistributed(db_batch, zookeeper) 
      dd.ancestor_query(query, filter_info, None)
      # Now with a transaction
      transaction = query.mutable_transaction() 
      transaction.set_handle(2)
      dd.ancestor_query(query, filter_info, None)
  
    def test_ordered_ancestor_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name",
                                                "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
      ancestor.MergeFrom(entity_key)
  
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01":
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      dd = DatastoreDistributed(db_batch, zookeeper)
      dd.ordered_ancestor_query(query, filter_info, None)
  
      # Now with a transaction
      transaction = query.mutable_transaction()
      transaction.set_handle(2)
      dd.ordered_ancestor_query(query, filter_info, None) 
  
    
    def test_kindless_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key() 
      key.MergeFrom(entity_key)
      
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      zookeeper.should_receive("acquire_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, zookeeper) 
      filter_info = {
        '__key__' : [[0, 0]]
      }
      dd.kindless_query(query, filter_info)
  
    def test_dynamic_delete(self):
      del_request = flexmock()
      del_request.should_receive("key_list")
      del_request.should_receive("has_transaction").never()
      del_request.should_receive("transaction").never()
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      dd.dynamic_delete("appid", del_request)
  
      fake_element = flexmock()
      fake_element.should_receive("type").and_return("kind")
      fake_path = flexmock()
      fake_path.should_receive("element_list").and_return([fake_element])
      fake_key = flexmock()
      fake_key.should_receive("path").and_return(fake_path)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(True).twice()
      transaction = flexmock()
      transaction.should_receive("handle").and_return(1)
      del_request.should_receive("transaction").and_return(transaction).once()
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, None)
      flexmock(dd).should_receive("acquire_locks_for_trans").and_return({})
      flexmock(dd).should_receive("release_locks_for_nontrans").never()
      flexmock(utils).should_receive("get_entity_kind").and_return("kind")
      flexmock(dd).should_receive('delete_entities_txn')
      dd.dynamic_delete("appid", del_request)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(False).twice()
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, None)
      flexmock(dd).should_receive("acquire_locks_for_trans").never()
      flexmock(dd).should_receive("acquire_locks_for_nontrans").once().and_return({})
      flexmock(dd).should_receive("delete_entities").once()
      flexmock(dd).should_receive("release_locks_for_nontrans").once()
      dd.dynamic_delete("appid", del_request)
  
    def test_reverse_path(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
      key = "Project:Synapse\x01Module:Core\x01"
      self.assertEquals(dd.reverse_path(key), "Module:Core\x01Project:Synapse\x01")
  
    def test_remove_exists_filters(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper)
      self.assertEquals(dd.remove_exists_filters({}), {})
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      self.assertEquals(dd.remove_exists_filters(filter_info), filter_info)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EXISTS, "1")],
        "prop2": [(datastore_pb.Query_Filter.EXISTS, "2")]}
      self.assertEquals(dd.remove_exists_filters(filter_info), {})
  
    def test_is_zigzag_merge_join(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      query = datastore_pb.Query()
      dd = DatastoreDistributed(db_batch, zookeeper) 
      db_batch.should_receive("remove_exists_filters").and_return({})
      self.assertEquals(dd.is_zigzag_merge_join(query, {}, {}), False)
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      db_batch.should_receive("remove_exists_filters").and_return(filter_info)
           
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), True)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), False)
  
    def test_zigzag_merge_join(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      query = datastore_pb.Query()
      dd = DatastoreDistributed(db_batch, zookeeper) 
      flexmock(dd).should_receive("is_zigzag_merge_join").and_return(False)
      self.assertEquals(dd.zigzag_merge_join(None, None, None), None)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      flexmock(query).should_receive("kind").and_return("kind")
      flexmock(dd).should_receive("get_table_prefix").and_return("prefix")
      flexmock(dd).should_receive("__apply_filters").and_return([])
      flexmock(query).should_receive("limit").and_return(1)
      self.assertEquals(dd.zigzag_merge_join(query, filter_info, []), None)
  
    def test_index_deletions(self):
      old_entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # No deletions should occur when the entity doesn't change.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      self.assertListEqual([], index_deletions(old_entity, old_entity))
  
      # When a property changes, the previous index entries should be deleted.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(old_entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
  
      deletions = index_deletions(old_entity, new_entity)
      self.assertEqual(len(deletions), 2)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = old_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # When given an index, an entry should be removed from the composite table.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 3)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.COMPOSITE_TABLE)
  
      # No composite deletions should occur when the entity type differs.
      definition.set_entity_type('TestEntity')
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 2)
  
    def test_deletions_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # Deleting an entity with one property should remove four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      deletions = deletions_for_entity(entity)
      self.assertEqual(len(deletions), 4)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.APP_KIND_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # Deleting an entity with two properties and one composite index should
      # remove seven entries.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = deletions_for_entity(entity, (composite_index,))
      self.assertEqual(len(deletions), 7)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[4]['table'], dbconstants.COMPOSITE_TABLE)
      self.assertEqual(deletions[5]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[6]['table'], dbconstants.APP_KIND_TABLE)
  
    def test_mutations_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
      txn = 1
  
      # Adding an entity with one property should add four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      mutations = mutations_for_entity(entity, txn)
      self.assertEqual(len(mutations), 4)
      self.assertEqual(mutations[0]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[1]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[2]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      # Updating an entity with one property should delete two entries and add
      # four more.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
      mutations = mutations_for_entity(entity, txn, new_entity)
      self.assertEqual(len(mutations), 6)
      self.assertEqual(mutations[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[0]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[1]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[4]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[5]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # Updating one property of an entity with two properties and one composite
      # index should remove three entries and add seven more.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
  
      mutations = mutations_for_entity(entity, txn, new_entity,
                                       (composite_index,))
      self.assertEqual(len(mutations), 10)
      self.assertEqual(mutations[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[0]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[1]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[2]['table'], dbconstants.COMPOSITE_TABLE)
      self.assertEqual(mutations[2]['operation'], dbconstants.Operations.DELETE)
      self.assertEqual(mutations[3]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[4]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[5]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[6]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[7]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[8]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[9]['table'], dbconstants.COMPOSITE_TABLE)
  
    def test_apply_txn_changes(self):
      app = 'guestbook'
      txn = 1
      entity = self.get_new_entity_proto(app, *self.BASIC_ENTITY[1:])
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive('range_query').and_return([{
        'txn_entity_1': {'operation': dbconstants.TxnActions.PUT,
                         'operand': entity.Encode()}
      }])
  
      db_batch.should_receive('get_indices').and_return([])
  
      dd = DatastoreDistributed(db_batch, None)
      prefix = dd.get_table_prefix(entity)
      entity_key = get_entity_key(prefix, entity.key().path())
      db_batch.should_receive('batch_get_entity').and_return({entity_key: {}})
      db_batch.should_receive('batch_mutate')
  
      dd.apply_txn_changes(app, txn)
  
    def test_delete_entities_txn(self):
      app = 'guestbook'
      txn_hash = {'root_key': 1}
      txn_str = '1'.zfill(ID_KEY_LENGTH)
      entity = self.get_new_entity_proto(app, *self.BASIC_ENTITY[1:])
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
  
      keys = [entity.key()]
      prefix = dd.get_table_prefix(entity.key())
      entity_key = get_entity_key(prefix, entity.key().path())
      encoded_path = str(encode_index_pb(entity.key().path()))
      txn_keys = [dd._SEPARATOR.join([app, txn_str, '', encoded_path])]
      txn_values = {
        txn_keys[0]: {
          dbconstants.TRANSACTIONS_SCHEMA[0]: dbconstants.TxnActions.DELETE,
          dbconstants.TRANSACTIONS_SCHEMA[1]: entity_key,
          dbconstants.TRANSACTIONS_SCHEMA[2]: ''
        }
      }
  
      flexmock(dd).should_receive('get_root_key').and_return('root_key')
  
      db_batch.should_receive('batch_put_entity').with_args(
        dbconstants.TRANSACTIONS_TABLE,
        txn_keys,
        dbconstants.TRANSACTIONS_SCHEMA,
        txn_values,
        ttl=TX_TIMEOUT * 2
      )
      dd.delete_entities_txn(app, keys, txn_hash)
  
    def test_put_entities_txn(self):
      app = 'guestbook'
      txn_hash = {'root_key': 1}
      txn_str = '1'.zfill(ID_KEY_LENGTH)
      entity = self.get_new_entity_proto(app, *self.BASIC_ENTITY[1:])
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
  
      entities = [entity]
      encoded_path = str(encode_index_pb(entity.key().path()))
      txn_keys = [dd._SEPARATOR.join([app, txn_str, '', encoded_path])]
      txn_values = {
        txn_keys[0]: {
          dbconstants.TRANSACTIONS_SCHEMA[0]: dbconstants.TxnActions.PUT,
          dbconstants.TRANSACTIONS_SCHEMA[1]: entity.Encode(),
          dbconstants.TRANSACTIONS_SCHEMA[2]: ''
        }
      }
  
      flexmock(dd).should_receive('get_root_key').and_return('root_key')
  
      db_batch.should_receive('batch_put_entity').with_args(
        dbconstants.TRANSACTIONS_TABLE,
        txn_keys,
        dbconstants.TRANSACTIONS_SCHEMA,
        txn_values,
        ttl=TX_TIMEOUT * 2
      )
      dd.put_entities_txn(entities, txn_hash, app)
  
  if __name__ == "__main__":
    unittest.main()    
  
files: AppDB/test/unit/test_datastore_server.py
files: #!/usr/bin/env python
  # Programmer: Navraj Chohan <nlake44@gmail.com>
  
  import sys
  import unittest
  
  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 ID_KEY_LENGTH
  from appscale.datastore.dbconstants import JOURNAL_SCHEMA
  from appscale.datastore.dbconstants import TOMBSTONE
  from appscale.datastore.cassandra_env.cassandra_interface import DatastoreProxy
  from appscale.datastore.cassandra_env.cassandra_interface import\
    deletions_for_entity
  from appscale.datastore.cassandra_env.cassandra_interface import\
    index_deletions
  from appscale.datastore.cassandra_env.cassandra_interface import\
    mutations_for_entity
  from appscale.datastore.unpackaged import APPSCALE_LIB_DIR
  from appscale.datastore.unpackaged import APPSCALE_PYTHON_APPSERVER
  from appscale.datastore.utils import encode_index_pb
  from appscale.datastore.utils import get_entity_key
  from appscale.datastore.utils import get_entity_kind
  from appscale.datastore.utils import get_index_key_from_params
  from appscale.datastore.utils import get_index_kv_from_tuple
  from appscale.datastore.utils import get_kind_key
  from appscale.datastore.zkappscale.zktransaction import TX_TIMEOUT
  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
  
  sys.path.append(APPSCALE_LIB_DIR)
  import appscale_info
  
  
  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):
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("increment_and_get_counter").and_return(0,1000)
      return zookeeper
  
    def test_get_entity_kind(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      item = Item(name="Bob", _app="hello")
      key = db.model_to_protobuf(item)
      self.assertEquals(get_entity_kind(key), "Item")
  
    def test_kind_key(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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)
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, 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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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)
      dd = DatastoreDistributed(db_batch, 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)
      dd = DatastoreDistributed(db_batch, 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)
      dd = DatastoreDistributed(db_batch, 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()
      dd = DatastoreDistributed(db_batch, self.get_zookeeper())
      dd.insert_composite_indexes([ent], [composite_index])
  
    def test_allocate_ids(self):
      PREFIX = "x"
      BATCH_SIZE = 1000
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, self.get_zookeeper())
      self.assertEquals(dd.allocate_ids(PREFIX, BATCH_SIZE), (1, 1000))
  
      dd = DatastoreDistributed(db_batch, self.get_zookeeper())
      self.assertEquals(dd.allocate_ids(PREFIX, None, max_id=1000), (1, 1000))
  
      try:
        # Unable to use self.assertRaises because of the optional argrument max_id
        ed = DatastoreDistributed(db_batch, self.get_zookeeper())
        dd.allocate_ids(PREFIX, BATCH_SIZE, max_id=10)
        raise "Allocate IDs should not let you set max_id and size"
      except ValueError:
        pass
  
    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()}})
  
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, 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)
      zookeeper = flexmock()
      zookeeper.should_receive("release_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, 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)
      zookeeper = flexmock()
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("notify_failed_transaction").and_return(True)
      dd = DatastoreDistributed(db_batch, 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()
      db_batch.should_receive('valid_data_version').and_return(True)
  
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      zookeeper.should_receive("get_transaction_id").and_return(1)
  
      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')
      dd = DatastoreDistributed(db_batch, 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()
      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)
  
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").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')
      dd = DatastoreDistributed(db_batch, zookeeper)
  
      # Make sure it does not throw an exception
      txn_hash = {entity_key1: 1, entity_key2: 1}
      dd.put_entities(app_id, entity_list, txn_hash)
  
    def test_acquire_locks_for_trans(self):
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      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, None)
      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, None)
      flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \
        .and_return(True)
  
      dd = DatastoreDistributed(db_batch, None)
      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()
      zookeeper.should_receive("acquire_lock").once()
      dd = DatastoreDistributed(db_batch, 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()
      zookeeper.should_receive("acquire_lock").once().and_raise(ZKTransactionException)
      zookeeper.should_receive("notify_failed_transaction").once()
      dd = DatastoreDistributed(db_batch, 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_acquire_locks_for_nontrans(self):
      app_id = 'test'
      PREFIX = 'x\x01'
      zookeeper = flexmock()
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("get_transaction_id").and_return(1).and_return(2)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return({PREFIX:{}})
      db_batch.should_receive("batch_delete").and_return(None)
      dd = DatastoreDistributed(db_batch, zookeeper) 
      entity_proto1 = self.get_new_entity_proto(
        app_id, 'test_kind', "bob", "prop1name", "prop1val", ns="blah")
      entity_proto2 = self.get_new_entity_proto(
        app_id, "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
      entity_list = [entity_proto1, entity_proto2]
      self.assertEquals({'test\x00blah\x00test_kind:bob\x01': 2, 'test\x00blah\x00test_kind:nancy\x01': 1}, 
                        dd.acquire_locks_for_nontrans("test", entity_list))
  
    def test_delete_entities(self):
      app_id = 'test'
      entity_proto1 = self.get_new_entity_proto(
        app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
      row_key = "test\x00blah\x00test_kind:bob\x01"
      row_values = {row_key: {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                              APP_ENTITY_SCHEMA[1]: '1'}}
  
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(row_values)
      db_batch.should_receive('batch_mutate')
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
  
      row_keys = [entity_proto1.key()]
  
      txn_hash = {row_key: 1}
      dd.delete_entities(app_id, row_keys, txn_hash)
       
    def test_release_put_locks_for_nontrans(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
      db_batch.should_receive("batch_delete").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "bob", "prop1name", 
                                                "prop1val", ns="blah")
      entity_proto2 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop2val", ns="blah")
      entities = [entity_proto1, entity_proto2]
      dd.release_locks_for_nontrans("test", entities, 
                    {'test\x00blah\x00test_kind:bob\x01': 1, 'test\x00blah\x00test_kind:nancy\x01': 2})
   
    def test_root_key_from_entity_key(self):
      zookeeper = flexmock()
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
      self.assertEquals("test\x00blah\x00test_kind:bob\x01", 
                        dd.get_root_key_from_entity_key("test\x00blah\x00test_kind:bob\x01nancy\x01"))
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop2val", ns="blah")
      self.assertEquals("test\x00blah\x00test_kind:nancy\x01", 
        dd.get_root_key_from_entity_key(entity_proto1.key()))
  
    def test_dynamic_get(self):
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop2val", ns="blah")
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
  
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key() 
      key.MergeFrom(entity_key)
      get_resp = datastore_pb.GetResponse()
      
      dd.dynamic_get("test", get_req, get_resp)     
      self.assertEquals(get_resp.entity_size(), 1)
  
      # Now test while in a transaction
      get_resp = datastore_pb.GetResponse()
      get_req.mutable_transaction().set_handle(1)
      dd.dynamic_get("test", get_req, get_resp)     
      self.assertEquals(get_resp.entity_size(), 1)
  
    def test_ancestor_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key() 
      key.MergeFrom(entity_key)
      ancestor.MergeFrom(entity_key)
      
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      dd = DatastoreDistributed(db_batch, zookeeper) 
      dd.ancestor_query(query, filter_info, None)
      # Now with a transaction
      transaction = query.mutable_transaction() 
      transaction.set_handle(2)
      dd.ancestor_query(query, filter_info, None)
  
    def test_ordered_ancestor_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name",
                                                "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key()
      key.MergeFrom(entity_key)
      ancestor.MergeFrom(entity_key)
  
      filter_info = []
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01":
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      dd = DatastoreDistributed(db_batch, zookeeper)
      dd.ordered_ancestor_query(query, filter_info, None)
  
      # Now with a transaction
      transaction = query.mutable_transaction()
      transaction.set_handle(2)
      dd.ordered_ancestor_query(query, filter_info, None) 
  
    
    def test_kindless_query(self):
      query = datastore_pb.Query()
      ancestor = query.mutable_ancestor()
      entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                                "prop1val", ns="blah")
      entity_key = entity_proto1.key()
      get_req = datastore_pb.GetRequest()
      key = get_req.add_key() 
      key.MergeFrom(entity_key)
      
      tombstone1 = {'key': {APP_ENTITY_SCHEMA[0]:TOMBSTONE, APP_ENTITY_SCHEMA[1]: 1}}
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_get_entity").and_return(
                 {"test\x00blah\x00test_kind:nancy\x01": 
                   {
                     APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                     APP_ENTITY_SCHEMA[1]: 1
                   }
                 })
  
      db_batch.should_receive("batch_put_entity").and_return(None)
      entity_proto1 = {'test\x00blah\x00test_kind:nancy\x01':{APP_ENTITY_SCHEMA[0]:entity_proto1.Encode(),
                        APP_ENTITY_SCHEMA[1]: 1}}
      db_batch.should_receive("range_query").and_return([entity_proto1, tombstone1]).and_return([])
      zookeeper = flexmock()
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("is_in_transaction").and_return(False)
      zookeeper.should_receive("acquire_lock").and_return(True)
      dd = DatastoreDistributed(db_batch, zookeeper) 
      filter_info = {
        '__key__' : [[0, 0]]
      }
      dd.kindless_query(query, filter_info)
  
    def test_dynamic_delete(self):
      del_request = flexmock()
      del_request.should_receive("key_list")
      del_request.should_receive("has_transaction").never()
      del_request.should_receive("transaction").never()
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      dd.dynamic_delete("appid", del_request)
  
      fake_element = flexmock()
      fake_element.should_receive("type").and_return("kind")
      fake_path = flexmock()
      fake_path.should_receive("element_list").and_return([fake_element])
      fake_key = flexmock()
      fake_key.should_receive("path").and_return(fake_path)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(True).twice()
      transaction = flexmock()
      transaction.should_receive("handle").and_return(1)
      del_request.should_receive("transaction").and_return(transaction).once()
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, None)
      flexmock(dd).should_receive("acquire_locks_for_trans").and_return({})
      flexmock(dd).should_receive("release_locks_for_nontrans").never()
      flexmock(utils).should_receive("get_entity_kind").and_return("kind")
      flexmock(dd).should_receive('delete_entities_txn')
      dd.dynamic_delete("appid", del_request)
  
      del_request = flexmock()
      del_request.should_receive("key_list").and_return([fake_key])
      del_request.should_receive("has_transaction").and_return(False).twice()
      del_request.should_receive("has_mark_changes").and_return(False)
      dd = DatastoreDistributed(db_batch, None)
      flexmock(dd).should_receive("acquire_locks_for_trans").never()
      flexmock(dd).should_receive("acquire_locks_for_nontrans").once().and_return({})
      flexmock(dd).should_receive("delete_entities").once()
      flexmock(dd).should_receive("release_locks_for_nontrans").once()
      dd.dynamic_delete("appid", del_request)
  
    def test_reverse_path(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper) 
      key = "Project:Synapse\x01Module:Core\x01"
      self.assertEquals(dd.reverse_path(key), "Module:Core\x01Project:Synapse\x01")
  
    def test_remove_exists_filters(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      dd = DatastoreDistributed(db_batch, zookeeper)
      self.assertEquals(dd.remove_exists_filters({}), {})
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      self.assertEquals(dd.remove_exists_filters(filter_info), filter_info)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EXISTS, "1")],
        "prop2": [(datastore_pb.Query_Filter.EXISTS, "2")]}
      self.assertEquals(dd.remove_exists_filters(filter_info), {})
  
    def test_is_zigzag_merge_join(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      query = datastore_pb.Query()
      dd = DatastoreDistributed(db_batch, zookeeper) 
      db_batch.should_receive("remove_exists_filters").and_return({})
      self.assertEquals(dd.is_zigzag_merge_join(query, {}, {}), False)
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      db_batch.should_receive("remove_exists_filters").and_return(filter_info)
           
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), True)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop1": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), False)
  
    def test_zigzag_merge_join(self):
      zookeeper = flexmock()
      zookeeper.should_receive("get_transaction_id").and_return(1)
      zookeeper.should_receive("get_valid_transaction_id").and_return(1)
      zookeeper.should_receive("register_updated_key").and_return(1)
      zookeeper.should_receive("acquire_lock").and_return(True)
      zookeeper.should_receive("release_lock").and_return(True)
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive("batch_delete").and_return(None)
      db_batch.should_receive("batch_put_entity").and_return(None)
      db_batch.should_receive("batch_get_entity").and_return(None)
  
      query = datastore_pb.Query()
      dd = DatastoreDistributed(db_batch, zookeeper) 
      flexmock(dd).should_receive("is_zigzag_merge_join").and_return(False)
      self.assertEquals(dd.zigzag_merge_join(None, None, None), None)
  
      filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
        "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
      flexmock(query).should_receive("kind").and_return("kind")
      flexmock(dd).should_receive("get_table_prefix").and_return("prefix")
      flexmock(dd).should_receive("__apply_filters").and_return([])
      flexmock(query).should_receive("limit").and_return(1)
      self.assertEquals(dd.zigzag_merge_join(query, filter_info, []), None)
  
    def test_index_deletions(self):
      old_entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # No deletions should occur when the entity doesn't change.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      self.assertListEqual([], index_deletions(old_entity, old_entity))
  
      # When a property changes, the previous index entries should be deleted.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(old_entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
  
      deletions = index_deletions(old_entity, new_entity)
      self.assertEqual(len(deletions), 2)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = old_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # When given an index, an entry should be removed from the composite table.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 3)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.COMPOSITE_TABLE)
  
      # No composite deletions should occur when the entity type differs.
      definition.set_entity_type('TestEntity')
      deletions = index_deletions(old_entity, new_entity, (composite_index,))
      self.assertEqual(len(deletions), 2)
  
    def test_deletions_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
  
      # Deleting an entity with one property should remove four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      deletions = deletions_for_entity(entity)
      self.assertEqual(len(deletions), 4)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.APP_KIND_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # Deleting an entity with two properties and one composite index should
      # remove seven entries.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      deletions = deletions_for_entity(entity, (composite_index,))
      self.assertEqual(len(deletions), 7)
      self.assertEqual(deletions[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[1]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(deletions[2]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(deletions[4]['table'], dbconstants.COMPOSITE_TABLE)
      self.assertEqual(deletions[5]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(deletions[6]['table'], dbconstants.APP_KIND_TABLE)
  
    def test_mutations_for_entity(self):
      entity = self.get_new_entity_proto(*self.BASIC_ENTITY)
      txn = 1
  
      # Adding an entity with one property should add four entries.
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
      mutations = mutations_for_entity(entity, txn)
      self.assertEqual(len(mutations), 4)
      self.assertEqual(mutations[0]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[1]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[2]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      # Updating an entity with one property should delete two entries and add
      # four more.
      new_entity = entity_pb.EntityProto()
      new_entity.MergeFrom(entity)
      new_entity.property_list()[0].value().set_stringvalue('updated content')
      mutations = mutations_for_entity(entity, txn, new_entity)
      self.assertEqual(len(mutations), 6)
      self.assertEqual(mutations[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[0]['operation'], dbconstants.TxnActions.DELETE)
      self.assertEqual(mutations[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[1]['operation'], dbconstants.TxnActions.DELETE)
      self.assertEqual(mutations[2]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[3]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[4]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[5]['table'], dbconstants.DSC_PROPERTY_TABLE)
  
      prop = entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      prop = new_entity.add_property()
      prop.set_name('author')
      prop.set_multiple(0)
      value = prop.mutable_value()
      value.set_stringvalue('author1')
  
      # Updating one property of an entity with two properties and one composite
      # index should remove three entries and add seven more.
      composite_index = entity_pb.CompositeIndex()
      composite_index.set_id(123)
      composite_index.set_app_id('guestbook')
      definition = composite_index.mutable_definition()
      definition.set_entity_type('Greeting')
      prop1 = definition.add_property()
      prop1.set_name('content')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
      prop2 = definition.add_property()
      prop2.set_name('author')
      prop1.set_direction(datastore_pb.Query_Order.ASCENDING)
  
      mutations = mutations_for_entity(entity, txn, new_entity,
                                       (composite_index,))
      self.assertEqual(len(mutations), 10)
      self.assertEqual(mutations[0]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[0]['operation'], dbconstants.TxnActions.DELETE)
      self.assertEqual(mutations[1]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[1]['operation'], dbconstants.TxnActions.DELETE)
      self.assertEqual(mutations[2]['table'], dbconstants.COMPOSITE_TABLE)
      self.assertEqual(mutations[2]['operation'], dbconstants.TxnActions.DELETE)
      self.assertEqual(mutations[3]['table'], dbconstants.APP_ENTITY_TABLE)
      self.assertEqual(mutations[4]['table'], dbconstants.APP_KIND_TABLE)
      self.assertEqual(mutations[5]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[6]['table'], dbconstants.ASC_PROPERTY_TABLE)
      self.assertEqual(mutations[7]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[8]['table'], dbconstants.DSC_PROPERTY_TABLE)
      self.assertEqual(mutations[9]['table'], dbconstants.COMPOSITE_TABLE)
  
    def test_apply_txn_changes(self):
      app = 'guestbook'
      txn = 1
      entity = self.get_new_entity_proto(app, *self.BASIC_ENTITY[1:])
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      db_batch.should_receive('range_query').and_return([{
        'txn_entity_1': {'operation': dbconstants.TxnActions.PUT,
                         'operand': entity.Encode()}
      }])
  
      db_batch.should_receive('get_indices').and_return([])
  
      dd = DatastoreDistributed(db_batch, None)
      prefix = dd.get_table_prefix(entity)
      entity_key = get_entity_key(prefix, entity.key().path())
      db_batch.should_receive('batch_get_entity').and_return({entity_key: {}})
      db_batch.should_receive('batch_mutate')
  
      dd.apply_txn_changes(app, txn)
  
    def test_delete_entities_txn(self):
      app = 'guestbook'
      txn_hash = {'root_key': 1}
      txn_str = '1'.zfill(ID_KEY_LENGTH)
      entity = self.get_new_entity_proto(app, *self.BASIC_ENTITY[1:])
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
  
      keys = [entity.key()]
      prefix = dd.get_table_prefix(entity.key())
      entity_key = get_entity_key(prefix, entity.key().path())
      encoded_path = str(encode_index_pb(entity.key().path()))
      txn_keys = [dd._SEPARATOR.join([app, txn_str, '', encoded_path])]
      txn_values = {
        txn_keys[0]: {
          dbconstants.TRANSACTIONS_SCHEMA[0]: dbconstants.TxnActions.DELETE,
          dbconstants.TRANSACTIONS_SCHEMA[1]: entity_key,
          dbconstants.TRANSACTIONS_SCHEMA[2]: ''
        }
      }
  
      flexmock(dd).should_receive('get_root_key').and_return('root_key')
  
      db_batch.should_receive('batch_put_entity').with_args(
        dbconstants.TRANSACTIONS_TABLE,
        txn_keys,
        dbconstants.TRANSACTIONS_SCHEMA,
        txn_values,
        ttl=TX_TIMEOUT * 2
      )
      dd.delete_entities_txn(app, keys, txn_hash)
  
    def test_put_entities_txn(self):
      app = 'guestbook'
      txn_hash = {'root_key': 1}
      txn_str = '1'.zfill(ID_KEY_LENGTH)
      entity = self.get_new_entity_proto(app, *self.BASIC_ENTITY[1:])
  
      db_batch = flexmock()
      db_batch.should_receive('valid_data_version').and_return(True)
      dd = DatastoreDistributed(db_batch, None)
  
      entities = [entity]
      encoded_path = str(encode_index_pb(entity.key().path()))
      txn_keys = [dd._SEPARATOR.join([app, txn_str, '', encoded_path])]
      txn_values = {
        txn_keys[0]: {
          dbconstants.TRANSACTIONS_SCHEMA[0]: dbconstants.TxnActions.PUT,
          dbconstants.TRANSACTIONS_SCHEMA[1]: entity.Encode(),
          dbconstants.TRANSACTIONS_SCHEMA[2]: ''
        }
      }
  
      flexmock(dd).should_receive('get_root_key').and_return('root_key')
  
      db_batch.should_receive('batch_put_entity').with_args(
        dbconstants.TRANSACTIONS_TABLE,
        txn_keys,
        dbconstants.TRANSACTIONS_SCHEMA,
        txn_values,
        ttl=TX_TIMEOUT * 2
      )
      dd.put_entities_txn(entities, txn_hash, app)
  
  if __name__ == "__main__":
    unittest.main()    
  
files: AppDB/test/unit/test_datastore_server.py