Python datastore_distributed.DatastoreDistributed类代码示例

 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_is_zigzag_merge_join(self):
    zk_client = flexmock()
    zk_client.should_receive('add_listener')

    zookeeper = flexmock(handle=zk_client)
    zookeeper.should_receive("get_valid_transaction_id").and_return(1)
    zookeeper.should_receive("acquire_lock").and_return(True)
    zookeeper.should_receive("release_lock").and_return(True)
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version').and_return(True)
    db_batch.should_receive("batch_delete").and_return(None)
    db_batch.should_receive("batch_put_entity").and_return(None)
    db_batch.should_receive("batch_get_entity").and_return(None)

    query = datastore_pb.Query()
    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
    db_batch.should_receive("remove_exists_filters").and_return({})
    self.assertEquals(dd.is_zigzag_merge_join(query, {}, {}), False)
    filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
      "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
    db_batch.should_receive("remove_exists_filters").and_return(filter_info)
         
    self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), True)

    filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
      "prop1": [(datastore_pb.Query_Filter.EQUAL, "2")]}
    self.assertEquals(dd.is_zigzag_merge_join(query, filter_info, []), False)

  def testFetchKeys(self):
    entity_proto1 = self.get_new_entity_proto("test", "test_kind", "bob", "prop1name", 
                                              "prop1val", ns="blah")

    db_batch = flexmock()
    db_batch.should_receive('valid_data_version').and_return(True)
    db_batch.should_receive("batch_delete").and_return(None)
    db_batch.should_receive("batch_put_entity").and_return(None)
    db_batch.should_receive("batch_get_entity").and_return({'test\x00blah\x00test_kind:bob\x01':
                {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(), 
                 APP_ENTITY_SCHEMA[1]: 1}}).and_return({'test\x00blah\x00test_kind:bob\x01\x000000000002':
                {JOURNAL_SCHEMA[0]: entity_proto1.Encode()}})

    zk_client = flexmock()
    zk_client.should_receive('add_listener')

    zookeeper = flexmock(handle=zk_client)
    zookeeper.should_receive("acquire_lock").and_return(True)
    zookeeper.should_receive("release_lock").and_return(True)
    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)

    self.assertEquals(({'test\x00blah\x00test_kind:bob\x01':
                         {'txnID': 1, 'entity': entity_proto1.Encode()}
                       },
                       ['test\x00blah\x00test_kind:bob\x01']),
                       dd.fetch_keys([entity_proto1.key()]))

  def test_create_composite_index(self):
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version_sync').and_return(True)
    db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE)
    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager,
                              self.get_zookeeper())
    zk_handle = flexmock(ensure_path=lambda path: None,
                         get=lambda path: (None, flexmock(version=None)),
                         set=lambda path, value, version: None)
    dd.zookeeper.handle = zk_handle
    dd = flexmock(dd)
    index = entity_pb.CompositeIndex()
    index.set_app_id("appid")
    index.set_state(2)
    definition = index.mutable_definition()
    definition.set_entity_type("kind")
    definition.set_ancestor(0)
    prop1 = definition.add_property()
    prop1.set_name("prop1")
    prop1.set_direction(1) # ascending
    prop2 = definition.add_property()
    prop2.set_name("prop2")
    prop1.set_direction(1) # ascending

    index_id = yield dd.create_composite_index("appid", index)
    assert index_id > 0

  def test_put_entities(self):
    app_id = 'test'
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version').and_return(True)

    entity_proto1 = self.get_new_entity_proto(
      app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
    entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
    entity_proto2 = self.get_new_entity_proto(
      app_id, "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
    entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
    entity_list = [entity_proto1, entity_proto2]

    db_batch.should_receive('batch_get_entity').and_return(
      {entity_key1: {}, entity_key2: {}})
    db_batch.should_receive('batch_mutate')
    transaction_manager = flexmock(
      create_transaction_id=lambda project, xg: 1,
      delete_transaction_id=lambda project, txid: None)
    dd = DatastoreDistributed(db_batch, transaction_manager,
                              self.get_zookeeper())

    entity_lock = flexmock(EntityLock)
    entity_lock.should_receive('acquire')
    entity_lock.should_receive('release')

    dd.put_entities(app_id, entity_list)

  def test_get_composite_index_key(self):
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version_sync').and_return(True)
    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager,
                              self.get_zookeeper())
    dd = flexmock(dd)

    composite_index = entity_pb.CompositeIndex()
    composite_index.set_id(123)
    composite_index.set_app_id("appid")

    definition = composite_index.mutable_definition()
    definition.set_entity_type("kind")

    prop1 = definition.add_property()
    prop1.set_name("prop1")
    prop1.set_direction(1) # ascending
    prop2 = definition.add_property()
    prop2.set_name("prop2")
    prop1.set_direction(1) # ascending

    ent = self.get_new_entity_proto("appid", "kind", "entity_name", "prop1", "value", ns="")

    self.assertEquals(
      dd.get_composite_index_key(composite_index, ent),
      "appid\x00\x00123\x00\x9avalue\x01\x01\x00\x00kind:entity_name\x01")

  def test_zigzag_merge_join(self):
    zk_client = flexmock()
    zk_client.should_receive('add_listener')

    zookeeper = flexmock(handle=zk_client)
    zookeeper.should_receive("get_valid_transaction_id").and_return(1)
    zookeeper.should_receive("acquire_lock").and_return(True)
    zookeeper.should_receive("release_lock").and_return(True)
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version').and_return(True)
    db_batch.should_receive("batch_delete").and_return(None)
    db_batch.should_receive("batch_put_entity").and_return(None)
    db_batch.should_receive("batch_get_entity").and_return(None)

    query = datastore_pb.Query()
    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)
    flexmock(dd).should_receive("is_zigzag_merge_join").and_return(False)
    self.assertEquals(dd.zigzag_merge_join(None, None, None), None)

    filter_info = {"prop1":[(datastore_pb.Query_Filter.EQUAL, "1")],
      "prop2": [(datastore_pb.Query_Filter.EQUAL, "2")]}
    flexmock(query).should_receive("kind").and_return("kind")
    flexmock(dd).should_receive("get_table_prefix").and_return("prefix")
    flexmock(dd).should_receive("__apply_filters").and_return([])
    flexmock(query).should_receive("limit").and_return(1)
    self.assertEquals(dd.zigzag_merge_join(query, filter_info, []), None)

  def test_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('get_transaction_metadata').and_return({
      'puts': {entity.key().Encode(): entity.Encode()},
      'deletes': [],
      'tasks': [],
      'reads': set(),
      'start': datetime.datetime.utcnow(),
      'is_xg': False,
    })
    db_batch.should_receive('valid_data_version').and_return(True)
    db_batch.should_receive('group_updates').and_return([])

    db_batch.should_receive('get_indices').and_return([])

    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager,
                              self.get_zookeeper())
    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')

    entity_lock = flexmock(EntityLock)
    entity_lock.should_receive('acquire')
    entity_lock.should_receive('release')

    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_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_insert_composite_indexes(self):
    composite_index = entity_pb.CompositeIndex()
    composite_index.set_id(123)
    composite_index.set_app_id("appid")

    definition = composite_index.mutable_definition()
    definition.set_entity_type("kind")

    prop1 = definition.add_property()
    prop1.set_name("prop1")
    prop1.set_direction(1) # ascending
    prop2 = definition.add_property()
    prop2.set_name("prop2")
    prop1.set_direction(1) # ascending

    ent = self.get_new_entity_proto(
      "appid", "kind", "entity_name", "prop1", "value", ns="")

    db_batch = flexmock()
    db_batch.should_receive('valid_data_version_sync').and_return(True)
    db_batch.should_receive("batch_put_entity").and_return(ASYNC_NONE).once()
    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager,
                              self.get_zookeeper())
    yield dd.insert_composite_indexes([ent], [composite_index])

  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_put_entities(self):
    app_id = 'test'
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version_sync').and_return(True)

    entity_proto1 = self.get_new_entity_proto(
      app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
    entity_key1 = 'test\x00blah\x00test_kind:bob\x01'
    entity_proto2 = self.get_new_entity_proto(
      app_id, "test_kind", "nancy", "prop1name", "prop2val", ns="blah")
    entity_key2 = 'test\x00blah\x00test_kind:nancy\x01'
    entity_list = [entity_proto1, entity_proto2]
    async_result = gen.Future()
    async_result.set_result({entity_key1: {}, entity_key2: {}})

    db_batch.should_receive('batch_get_entity').and_return(async_result)
    db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)
    transaction_manager = flexmock(
      create_transaction_id=lambda project, xg: 1,
      delete_transaction_id=lambda project, txid: None,
      set_groups=lambda project, txid, groups: None)
    dd = DatastoreDistributed(db_batch, transaction_manager,
                              self.get_zookeeper())
    dd.index_manager = flexmock(
      projects={app_id: flexmock(indexes_pb=[])})

    async_true = gen.Future()
    async_true.set_result(True)
    entity_lock = flexmock(EntityLock)
    entity_lock.should_receive('acquire').and_return(async_true)
    entity_lock.should_receive('release')

    yield dd.put_entities(app_id, entity_list)

  def test_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_delete_entities(self):
    app_id = 'test'
    entity_proto1 = self.get_new_entity_proto(
      app_id, "test_kind", "bob", "prop1name", "prop1val", ns="blah")
    row_key = "test\x00blah\x00test_kind:bob\x01"
    row_values = {row_key: {APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                            APP_ENTITY_SCHEMA[1]: '1'}}
    async_result = gen.Future()
    async_result.set_result(row_values)

    zk_client = flexmock()
    zk_client.should_receive('add_listener')

    zookeeper = flexmock(handle=zk_client)
    zookeeper.should_receive("get_valid_transaction_id").and_return(1)
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version_sync').and_return(True)
    db_batch.should_receive("batch_get_entity").and_return(async_result)
    db_batch.should_receive('normal_batch').and_return(ASYNC_NONE)

    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)

    row_keys = [entity_proto1.key()]

    yield dd.delete_entities(entity_proto1.key(), 1, row_keys)