def Keys(self, pattern): """Retrieve the set of keys in the database matching a glob-style pattern. NOTE: This command runs in O(n), where n is the number of keys in the database. It should only be used for maintenance or debugging purposes. Args: pattern - Glob-style pattern to retrieve keys for. Returns: An Operation object with the result of the query. The response_value field contains a list of keys matching the pattern. """ keys = [] try: # Modify the pattern to account for the vbucket prefix. shard_pattern = '*|%s' % pattern # Get the keys from each shard. for shard in self.shards: keys.extend(shard.keys(shard_pattern)) # Split out the vbucket info from the returned keys. op = Operation( success=True, response_value=['|'.join(key.split('|')[1:]) for key in keys]) except Exception: op = Operation(success=False, traceback=traceback.format_exc()) return op
def SetAdd(self, key, *values): """Add values to the Set under key. Args: key - The key of the Set. values - List of values to add to the Set. Returns: An Operation object with the result of the query. The response_value field contains a list of integers, the sum of which represents the number of new items added to the set. """ if key in self.val_dict: if type(self.val_dict[key]) != set: return Operation(success=False) before = len(self.val_dict[key]) self.val_dict[key].update(values) added = len(self.val_dict[key]) - before else: self.val_dict[key] = set(values) added = len(self.val_dict[key]) return Operation(success=True, response_value=added)
def StateGetBatchGenerator(self, id_tuples): """Retrieve a batch of State objects. Args: id_tuples - List of tuples of the form (client_id, name) of the State objects to retrieve. Returns: Operation object with the query results. If successful, the response_value field contains a List of tuples of the form ((client_id, name), state) """ # Convert the ID tuples to keys and retrieve from the DB. keys = [ self._MakeKey(KEY_STATE, client, name) for client, name in id_tuples ] op = self._GetBatch(keys) # Unpack the StateStruct objects and join the retrieved values to their # original ID tuples. if op.response_value: try: states_unpacked = itertools.imap(taba_state.UnpackState, op.response_value) op.response_value = itertools.izip(id_tuples, states_unpacked) except Exception: LOG.error("Error unpacking State objects") LOG.error(traceback.format_exc()) op = Operation(success=False, traceback=traceback.format_exc()) return op
def testStatePutFail(self): """Test State Put when the operation fails""" TUPLES = [ (('c1', 'n1'), StateStruct('s1', 0)), (('c2', 'n2'), StateStruct('s2', 0)), ] enc_tuples = [ ('state:c1:n1', S1_PACKED), ('state:c2:n2', S2_PACKED), ] self.mox.StubOutWithMock(taba_server_storage_manager, 'LOG') taba_server_storage_manager.LOG.error(mox.IgnoreArg()).MultipleTimes() self.mox.StubOutWithMock(self.engine, 'BatchPut') self.engine.BatchPut(enc_tuples).AndReturn(Operation(False)) self.mox.ReplayAll() op = self.vssm.StatePutBatch(TUPLES) self.assertFalse(op.success) self.assertTrue(op.response_value is None) self.assertEqual(self.engine.val_dict.get('state:c1:n1'), None) self.assertEqual(self.engine.val_dict.get('state:c2:n2'), None) self.assertEqual(self.engine.val_dict.get('clients'), None) self.assertEqual(self.engine.val_dict.get('names:all'), None) self.assertEqual(self.engine.val_dict.get('names:c1'), None) self.assertEqual(self.engine.val_dict.get('names:c2'), None) self.mox.VerifyAll()
def _GetBatch(self, keys): """Retrieve several keys in a single operation. Values will be decoded before returning. Args: keys - List of keys to retrieve. Returns: Operation object with the query results. """ try: # Retrieve the encoded valued from the storage engine. op = self.engine.BatchGet(keys) if not op.success: LOG.error("Batch Get operation failed") LOG.error(op) return op op.response_value = [v if v else None for v in op.response_value] except Exception: LOG.error("Exception in Batch Get operation") LOG.error(traceback.format_exc()) op = Operation(success=False, traceback=traceback.format_exc()) return op
def StateDeleteBatch(self, id_tuples): """Delete a batch of State objects. Args: id_tuples - List of tuples of the form (client_id, name) of the State objects to delete. Returns: Operation object with the query results. """ keys = [ self._MakeKey(KEY_STATE, client, name) for client, name in id_tuples ] try: op = self.engine.BatchDelete(keys) if not op.success: LOG.error("Batch Delete operation failed") LOG.error(op) except Exception: LOG.error("Exception in Batch Delete operation") LOG.error(traceback.format_exc()) op = Operation(success=False, traceback=traceback.format_exc()) return op
def _PutBatch(self, key_value_tuples): """Store several keys in a single operation. Values will be encoded before storing. Args: key_value_tuples - List of tuples of (key, value) to store. Returns: Operation object with the result of the query. """ try: # Put the values to the storege engine. op = self.engine.BatchPut(key_value_tuples) if not op.success: LOG.error("Batch Put operation failed") LOG.error(op) return op except Exception: LOG.error("Exception in Batch Get operation") LOG.error(traceback.format_exc()) op = Operation(success=False, traceback=traceback.format_exc()) return op
def _CheckedOp(self, description, engine_fn, *engine_fn_args): """Execute an operation which returns an Operation object, and check it for success or exceptions, with logging in the case of an error. Args: description - String describing the operation which will be appended to error messages. engine_fn - Callback to execute, which return an Operation object. enging_fn_args - List of arguments to bass to engine_fn Returns: Operation object with the query results. """ try: op = engine_fn(*engine_fn_args) if not op.success: LOG.error("Error %s" % description) LOG.error(op) return op except Exception: LOG.error("Exception %s" % description) LOG.error(traceback.format_exc()) op = Operation(success=False, traceback=traceback.format_exc()) return op
def BatchCheckAndMultiSet(self, keys, callback_fn): """Set a batch of keys transactionally using optimistic locking. Each key is locked and retrieved. The retrieved values are passed, one at a time, to the callback function. The return value of the callback is then put in the key. If the value changed between the time it was locked and when the Put happens, the Put will fail, and the whole operation will be retried, until it succeeds. Args: keys - List of keys to update transactionally. callback_fn - Callable which, given a key and a value, returns the updated value for that key, which will be put back in the DB. If a put fails due to a lock error, this function will be called again with the updated value. Returns: A CompountOperation with the result of the queries. """ responses = [] for key in keys: val = self.val_dict.get(key) new_kvs = callback_fn(key, val) for key, new_val in new_kvs: self.val_dict[key] = new_val responses.append(True) return Operation(success=True, response_value=responses)
def BatchGet(self, keys): """Return the values for a set of keys. Args: keys - List of keys to lookup. Returns: CompoundOperation with the results of the lookups. The response_value field has the form [(key, value)., ...] """ responses = [] for key in keys: if type(self.val_dict.get(key)) == set: return Operation(success=False) responses.append(self.val_dict.get(key)) return Operation(success=True, response_value=responses)
def BatchPut(self, key_value_tuples): """Put a batch of values into keys. Args: key_value_tuples - A list of tuples of the form (key, value) Returns: A CompoundOperation with the results of the queries. """ responses = [] for key, value in key_value_tuples: if type(self.val_dict.get(key)) == set: return Operation(success=False) self.val_dict[key] = value responses.append(True) op = Operation(success=True, response_value=responses) return op
def HashGetAll(self, key): """Retrieve an entire Hashtable. Args: key - The key of the Hashtable. Returns: An Operation object with the result of the query. """ if key in self.val_dict: if type(self.val_dict[key]) != dict: return Operation(success=False) result = self.val_dict[key] else: result = None return Operation(success=True, response_value=result)
def SetMembers(self, key): """Retrieve all the members of the Set at key. Args: key: - The key of the Set. Returns: An Operation object with the result of the query. The response_value field contains a set() object with the members of the Set. """ if key in self.val_dict: if type(self.val_dict[key]) != set: return Operation(success=False) val = self.val_dict[key] else: val = set([]) return Operation(success=True, response_value=val)
def HashGet(self, key, field): """Retrieve the value stored at a field in a Hashtable. Args: key - The key of the Hashtable. field - Field in the Hashtable to retrieve. Returns: An Operation object with the result of the query. """ if key in self.val_dict: if type(self.val_dict[key]) != dict: return Operation(success=False) result = self.val_dict[key].get(field) else: result = None return Operation(success=True, response_value=result)
def HashMultiPut(self, key, mapping): """Put a set of fields and values to a Hashtable. Args: key - The key of the Hashtable. mapping - Dictionary of keys and values to set. Returns: An Operation object with the result of the query. """ if key in self.val_dict: if type(self.val_dict[key]) != dict: return Operation(success=False) self.val_dict[key].extend(mapping) else: self.val_dict[key] = mapping return Operation(success=True, response_value=True)
def _ShardCheckAndSetBatch(self, shard, keys, vkeys, values): retries = 0 while True: # Open a transactional pipeline. pipe = shard.pipeline(True) try: # Lock the keys to start the operation. pipe.watch(*vkeys) # Batch Get the keys to be updated. sub_pipe = pipe.pipeline(False) map(sub_pipe.get, vkeys) values = sub_pipe.execute() # Get the new values from the client. put_kv_tuples = [] for i, (key, value) in enumerate(itertools.izip(keys, values)): put_kv_tuples.extend(self.get_updates_fn(key, value)) # Yield periodically. if i % 1000 == 0: gevent.sleep(0) # Put the new values into the DB. pipe.multi() for key, value in put_kv_tuples: vbucket = self.engine._GetVbucket(key) vkey = self.engine._MakeVkey(key, vbucket) pipe.set(vkey, value) response = pipe.execute() except redis.WatchError: # Lock error occurred. Try the operation again. gevent.sleep(0.1 * retries + 0.1 * random.random()) retries += 1 if retries > MAX_TRANSACTION_RETRIES: raise continue finally: # Make sure we always reset the pipe. pipe.reset() # If we make it here without a WatchError, the operation succeeded. op = Operation(success=True, response_value=response, retries=retries) break return op
def SetIsMember(self, key, value): """Test whether a value is a member of the Set at key. Args: key - The key of the Set. value - The String value to test membership in the Set. Returns: A boolean indicating whether the value is in the set. """ if key in self.val_dict: if type(self.val_dict[key]) != set: return Operation(success=False) val = value in self.val_dict[key] else: val = False return Operation(success=True, response_value=val)
def HashBatchGet(self, keys, field): """Retrieve the value of a field from several Hashtables. Args: keys - The keys of the Hashtables to lookup. field - The field to get from each Hashtable. Returns: """ result = [] for key in keys: if key in self.val_dict: if type(self.val_dict[key]) != dict: return Operation(success=False) result.append(self.val_dict[key].get(field)) else: result.append(None) return Operation(success=True, response_value=result)
def HashPut(self, key, field, value): """Put a value to a field of a Hashtable. Args: key - The key of the Hashtable field - Field in the Hashtable to set. value - Value to set. Returns: An Operation object with the result of the query. """ if key in self.val_dict: if type(self.val_dict[key]) != dict: return Operation(success=False) self.val_dict[key][field] = value result = 0 else: self.val_dict[key] = {field: value} result = 1 return Operation(success=True, response_value=result)
def _ShardGreenletWraper(shard_num, vkey_tuples): indices, keys, vkeys, values = [list(i) for i in zip(*vkey_tuples)] try: sub_op = shard_execute_fn(self.shards[shard_num], keys, vkeys, values) except Exception: sub_op = Operation(success=False, traceback=traceback.format_exc()) # Aggregate the results. op.AddOp(sub_op) if sub_op.response_value: responses.extend(zip(indices, sub_op.response_value))
def SetRemove(self, key, *values): """Remove values from the Set at key. Args: key - The key of the Set. values - List of values to remove from the Set. Returns: An Operation object with the result of the query. The response_value field contains a list of integers, the sum of which is the number of items removed from the Set. """ if key in self.val_dict: if type(self.val_dict[key]) != set: return Operation(success=False) before = len(self.val_dict[key]) self.val_dict[key] = self.val_dict[key] - set(values) removed = before - len(self.val_dict[key]) else: removed = 0 return Operation(success=True, response_value=removed)
def HashDelete(self, key, field): """Delete a field from a Hashtable. Args: key - The key of the Hashtable. field - Field to delete from the Hashtable. Returns: An Operation object with the result of the query. """ if key in self.val_dict: if type(self.val_dict[key]) != dict: return Operation(success=False) if field in self.val_dict[key]: del self.val_dict[key][field] removed = 1 else: removed = 0 else: removed = 0 return Operation(success=True, response_value=removed)
def BatchDelete(self, keys): """Delete a ket of keys. Args: keys - A list of keys to delete. Returns: A CompountOperation with the results of the queries. """ responses = [] for key in keys: if key in self.val_dict: del self.val_dict[key] responses.append(1) else: responses.append(0) return Operation(success=True, response_value=responses)
def _UpdateIdSets(self, id_tuples): """Update the sets of Client IDs and Taba Names. Args: id_tuples - List of tuples of the form (client_id, name). Returns: Operation object with the query results. """ op = CompoundOperation() # Build a map of Client ID to Taba Names client_id_names_map = defaultdict(list) for client_id, name in id_tuples: client_id_names_map[client_id].append(name) try: # Add the Client IDs to the Set of all Client IDs. op_clients = self.ClientIdsAdd(client_id_names_map.keys()) op.AddOp(op_clients) # Add all the Taba Names to the Set of all names across all Client IDs. all_names = set() all_names.update(*client_id_names_map.values()) op_names_all = self.TabaNamesForAllAdd(all_names) op.AddOp(op_names_all) # Add the Taba Names for each Client ID. for client_id, names in client_id_names_map.iteritems(): op_names = self.TabaNamesForClientAdd(client_id, names) op.AddOp(op_names) except Exception: LOG.error("Error updating Client/Name sets") LOG.error(traceback.format_exc()) op.AddOp(Operation(success=False, traceback=traceback.format_exc())) return op
def _ShardSetRemove(shard, keys, vkeys, values): pipe = shard.pipeline() map(pipe.srem, vkeys, values) removed = pipe.execute() return Operation(success=True, response_value=removed)
def _ShardHashPut(shard, keys, vkeys, vals): added = shard.hset(vkeys[0], field, value) return Operation(success=True, response_value=[added])
def _ShardHashMultiPut(shard, keys, vkeys, vals): response = shard.hmset(vkeys[0], mapping) return Operation(success=True, response_value=[response])
def _ShardHashGetAll(shard, keys, vkeys, vals): result = shard.hgetall(vkeys[0]) return Operation(success=True, response_value=[result])
def _ShardHashBatchGet(shard, keys, vkeys, vals): pipe = shard.pipeline() for vkey in vkeys: pipe.hget(vkey, field) result = pipe.execute() return Operation(success=True, response_value=result)
def _ShardHashDelete(shard, keys, vkeys, vals): added = shard.hdel(vkeys[0], field) return Operation(success=True, response_value=[added])