def _save_coverage_information(context, result):
"""Saves coverage information in datastore using an atomic transaction."""
# Use ndb.transaction with retries below to mitigate risk of a race condition.
def _try_save_coverage_information():
"""Implements save_coverage_information function."""
coverage_info = data_handler.get_coverage_information(
context.fuzz_target.project_qualified_name(),
result.coverage_info.date,
create_if_needed=True)
# Intentionally skip edge and function coverage values as those would come
# from fuzzer coverage cron task (see src/go/server/cron/coverage.go).
coverage_info.corpus_size_units = result.coverage_info.corpus_size_units
coverage_info.corpus_size_bytes = result.coverage_info.corpus_size_bytes
coverage_info.corpus_location = result.coverage_info.corpus_location
coverage_info.corpus_backup_location = (
result.coverage_info.corpus_backup_location)
coverage_info.quarantine_size_units = (
result.coverage_info.quarantine_size_units)
coverage_info.quarantine_size_bytes = (
result.coverage_info.quarantine_size_bytes)
coverage_info.quarantine_location = result.coverage_info.quarantine_location
coverage_info.put()
try:
ndb.transaction(_try_save_coverage_information,
retries=data_handler.DEFAULT_FAIL_RETRIES)
except Exception as e:
raise CorpusPruningException(
'Failed to save corpus pruning result: %s.' % repr(e))
async def set_email(content: APIKeyEmail, response: Response):
def update_email(api_key: str, email: str, response: Response):
email_obj = ndb.Key("UserEmail", api_key).get()
if email_obj is None:
email_obj = UserEmail(id=api_key, email=email)
response.status_code = status.HTTP_201_CREATED
else:
email_obj.email = email
response.status_code = status.HTTP_200_OK
email_obj.put()
client = ndb.Client()
with client.context():
ndb.transaction(lambda: update_email(
content.api_key, content.email, response))
def test_delete_entity_in_transaction(ds_entity):
entity_id = test_utils.system.unique_resource_id()
ds_entity(KIND, entity_id, foo=42)
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
key = ndb.Key(KIND, entity_id)
assert key.get().foo == 42
def delete_entity():
assert key.delete() is None
assert key.get().foo == 42 # not deleted until commit
ndb.transaction(delete_entity)
assert key.get() is None
def test_delete_entity_in_transaction_then_rollback(ds_entity):
entity_id = test_utils.system.unique_resource_id()
ds_entity(KIND, entity_id, foo=42)
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
key = ndb.Key(KIND, entity_id)
assert key.get().foo == 42
def delete_entity():
assert key.delete() is None
raise Exception("Spurious error")
with pytest.raises(Exception):
ndb.transaction(delete_entity)
assert key.get().foo == 42
def test_delete_entity_in_transaction_with_global_cache(
client_context, ds_entity):
"""Regression test for #426
https://github.com/googleapis/python-ndb/issues/426
"""
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
entity_id = test_utils.system.unique_resource_id()
ds_entity(KIND, entity_id, foo=42)
global_cache = global_cache_module._InProcessGlobalCache()
with client_context.new(global_cache=global_cache).use():
key = ndb.Key(KIND, entity_id)
assert key.get().foo == 42
ndb.transaction(key.delete)
assert key.get() is None
def test_insert_entity_in_transaction_without_preallocating_id(dispose_of):
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
bar = ndb.StringProperty()
def save_entity():
# By not waiting on the Future, we don't force a call to AllocateIds
# before the transaction is committed.
SomeKind(foo=42, bar="none").put_async()
ndb.transaction(save_entity)
query = SomeKind.query()
eventually(query.fetch, length_equals(1))
retrieved = query.fetch()[0]
dispose_of(retrieved._key._key)
assert retrieved.foo == 42
assert retrieved.bar == "none"
def acquire_lock(key_name,
max_hold_seconds=DEFAULT_MAX_HOLD_SECONDS,
retries=None,
by_zone=True):
"""Acquire a lock for the given key name. Returns the expiration time if
succeeded, otherwise None. The lock holder is responsible for making sure it
doesn't assume the lock is still held after the expiration time."""
logs.log('Acquiring lock for %s.' % key_name)
failed_acquires = 0
total_wait = 0
wait_exponent = 1
if by_zone:
key_name_with_zone = _get_key_name_with_lock_zone(key_name)
if key_name_with_zone is None:
logs.log_error('Failed to get zone while trying to lock %s.' %
key_name)
return None
key_name = key_name_with_zone
bot_name = environment.get_value('BOT_NAME')
expiration_delta = datetime.timedelta(seconds=max_hold_seconds)
while total_wait < LOCK_CHECK_TIMEOUT:
try:
lock_entity = ndb.transaction(lambda: _try_acquire_lock(
key_name,
expiration_time=datetime.datetime.utcnow() + expiration_delta,
holder=bot_name),
retries=TRANSACTION_RETRIES)
if lock_entity.holder == bot_name:
logs.log('Got the lock.')
return lock_entity.expiration_time
except exceptions.Error:
pass
failed_acquires += 1
if retries and retries >= failed_acquires:
logs.log('Failed to acquire lock, exceeded max retries.')
return None
logs.log('Failed to acquire lock, waiting...')
# Exponential backoff.
max_sleep = (1 << wait_exponent) * LOCK_CHECK_SLEEP_MULTIPLIER
sleep_time = random.uniform(1.0, max_sleep)
time.sleep(sleep_time)
total_wait += sleep_time
wait_exponent = min(wait_exponent + 1, MAX_WAIT_EXPONENT)
logs.log('Timeout exceeded while trying to acquire lock, bailing.')
return None
def test_get_or_insert_get_in_transaction(ds_entity):
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
name = "Inigo Montoya"
assert SomeKind.get_by_id(name) is None
def do_the_thing():
ds_entity(KIND, name, foo=42)
return SomeKind.get_or_insert(name, foo=21)
entity = ndb.transaction(do_the_thing)
assert entity.foo == 42
def test_insert_entity_in_transaction(dispose_of):
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
bar = ndb.StringProperty()
def save_entity():
entity = SomeKind(foo=42, bar="none")
key = entity.put()
dispose_of(key._key)
return key
key = ndb.transaction(save_entity)
retrieved = key.get()
assert retrieved.foo == 42
assert retrieved.bar == "none"
def test_insert_entity_in_transaction(dispose_of):
commit_callback = mock.Mock()
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
bar = ndb.StringProperty()
def save_entity():
ndb.get_context().call_on_commit(commit_callback)
entity = SomeKind(foo=42, bar="none")
key = entity.put()
dispose_of(key._key)
return key
key = ndb.transaction(save_entity)
retrieved = key.get()
assert retrieved.foo == 42
assert retrieved.bar == "none"
commit_callback.assert_called_once_with()
def test_update_entity_in_transaction(ds_entity):
entity_id = test_utils.system.unique_resource_id()
ds_entity(KIND, entity_id, foo=42, bar="none")
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
bar = ndb.StringProperty()
def update_entity():
key = ndb.Key(KIND, entity_id)
entity = key.get()
entity.foo = 56
entity.bar = "high"
assert entity.put() == key
return key
key = ndb.transaction(update_entity)
retrieved = key.get()
assert retrieved.foo == 56
assert retrieved.bar == "high"
def claim_token_with_tx():
"""Claim a token and record the action, using a transaction.
Because this function uses a datastore transaction, it can handle many
concurrent requests claiming the same token, but ensure the token is
issued to no more than 1 request.
"""
client = ndb.Client()
path = '/token-with-tx'
request_id = new_request_id()
token_id = next_token_id()
with client.context():
# This uses an inner function so that we can capture the function call
# arguments from the calling scope.
def _inner():
return get_or_insert(token_id, request_id, path)
token_obj, new_claim = ndb.transaction(_inner)
token = token_obj.to_dict()
return {'token': token, 'new_claim': new_claim}