def post(self):
entities = ndb.get_multi(
ndb.Key(urlsafe=u) for u in self.request.params.getall('key'))
for e in entities:
e.status = 'complete'
ndb.put_multi(entities)
self.redirect('/admin/responses')
def update_mappings_for_job(job, mappings):
"""Clear existing mappings for a job, and replace them."""
existing_fuzzers_query = data_types.Fuzzer.query(
data_types.Fuzzer.jobs == job.name)
existing_fuzzers = {
fuzzer.name: fuzzer
for fuzzer in existing_fuzzers_query
}
modified_fuzzers = []
for fuzzer_name in mappings:
fuzzer = existing_fuzzers.pop(fuzzer_name, None)
if fuzzer:
continue
fuzzer = data_types.Fuzzer.query(
data_types.Fuzzer.name == fuzzer_name).get()
if not fuzzer:
logs.log_error('An unknown fuzzer %s was selected for job %s.' %
(fuzzer_name, job.name))
continue
fuzzer.jobs.append(job.name)
modified_fuzzers.append(fuzzer)
update_mappings_for_fuzzer(fuzzer)
# Removing the remaining values in exisiting_fuzzers as
# they are no longer mapped.
for fuzzer in existing_fuzzers.values():
fuzzer.jobs.remove(job.name)
modified_fuzzers.append(fuzzer)
update_mappings_for_fuzzer(fuzzer)
ndb.put_multi(modified_fuzzers)
def update_affected_commits(bug_id, result, project, ecosystem, public):
"""Update affected commits."""
to_put = []
to_delete = []
for commit in result.commits:
affected_commit = osv.AffectedCommit(
id=bug_id + '-' + commit,
bug_id=bug_id,
commit=commit,
confidence=result.confidence,
project=project,
ecosystem=ecosystem,
public=public)
to_put.append(affected_commit)
# Delete any affected commits that no longer apply. This can happen in cases
# where a FixResult comes in later and we had previously marked a commit prior
# to the fix commit as being affected by a vulnerability.
for existing in osv.AffectedCommit.query(osv.AffectedCommit.bug_id == bug_id):
if existing.commit not in result.commits:
to_delete.append(existing.key)
ndb.put_multi(to_put)
ndb.delete_multi(to_delete)
def setUp(self):
"""Put data in the local ndb table the tests to query from and set
bandit selection environment variable."""
test_helpers.patch_environ(self)
data = []
strategy1 = data_types.FuzzStrategyProbability()
strategy1.strategy_name = 'fork,corpus_subset,recommended_dict,'
strategy1.probability = 0.33
strategy1.engine = 'libFuzzer'
data.append(strategy1)
strategy2 = data_types.FuzzStrategyProbability()
strategy2.strategy_name = ('random_max_len,corpus_mutations_ml_rnn,'
'value_profile,recommended_dict,')
strategy2.probability = 0.34
strategy2.engine = 'libFuzzer'
data.append(strategy2)
strategy3 = data_types.FuzzStrategyProbability()
strategy3.strategy_name = ('corpus_mutations_radamsa,'
'random_max_len,corpus_subset,')
strategy3.probability = 0.33
strategy3.engine = 'libFuzzer'
data.append(strategy3)
ndb.put_multi(data)
distribution = fuzz_task.get_strategy_distribution_from_ndb()
environment.set_value('USE_BANDIT_STRATEGY_SELECTION', True)
environment.set_value('STRATEGY_SELECTION_DISTRIBUTION', distribution)
def process_package_info_task(message):
"""Process project info."""
package_name = message.attributes['package_name']
ecosystem = message.attributes['ecosystem']
repo_url = message.attributes['repo_url']
tags_info = osv.get_tags(repo_url)
if tags_info.latest_tag:
info = osv.PackageInfo(id=f'{ecosystem}/{package_name}')
info.latest_tag = tags_info.latest_tag
info.put()
infos = []
for tag in tags_info.tags:
tag_info = osv.PackageTagInfo(id=f'{ecosystem}/{package_name}-{tag}')
tag_info.package = package_name
tag_info.ecosystem = ecosystem
tag_info.tag = tag
tag_info.bugs = find_bugs_for_tag(package_name, tag, public=True)
tag_info.bugs_private = find_bugs_for_tag(package_name,
tag,
public=False)
infos.append(tag_info)
ndb.put_multi(infos)
def make_bugs_public():
"""Mark bugs public."""
if not request.headers.get('X-Appengine-Cron'):
abort(403)
monorail_account = get_monorail_service_account()
monorail_client = monorail.Client('oss-fuzz', monorail_account)
query = osv.Bug.query(osv.Bug.public == False) # pylint: disable=singleton-comparison
to_mark_public = []
for bug in query:
issue_id = bug.issue_id
if not issue_id:
logging.info('Missing issue_id for %s.', bug.key.id())
continue
try:
issue = monorail_client.get_issue(issue_id)
except requests.exceptions.HTTPError:
logging.error('Failed to get issue %s.', issue_id)
continue
labels = [label['label'].lower() for label in issue['labels']]
if 'restrict-view-commit' not in labels:
bug.public = True
logging.info('Marking %s as public.', bug.key.id())
to_mark_public.append(bug)
make_affected_commits_public(bug)
if to_mark_public:
ndb.put_multi(to_mark_public)
return 'done'
def test_map(dispose_of):
class SomeKind(ndb.Model):
foo = ndb.StringProperty()
ref = ndb.KeyProperty()
class OtherKind(ndb.Model):
foo = ndb.StringProperty()
foos = ("aa", "bb", "cc", "dd", "ee")
others = [OtherKind(foo=foo) for foo in foos]
other_keys = ndb.put_multi(others)
for key in other_keys:
dispose_of(key._key)
things = [SomeKind(foo=foo, ref=key) for foo, key in zip(foos, other_keys)]
keys = ndb.put_multi(things)
for key in keys:
dispose_of(key._key)
eventually(SomeKind.query().fetch, _length_equals(5))
eventually(OtherKind.query().fetch, _length_equals(5))
@ndb.tasklet
def get_other_foo(thing):
other = yield thing.ref.get_async()
raise ndb.Return(other.foo)
query = SomeKind.query().order(SomeKind.foo)
assert query.map(get_other_foo) == foos
def mark_complete():
entities = ndb.get_multi(ndb.Key(urlsafe=u)
for u in request.values.getlist('key'))
for e in entities:
e.status = 'complete'
ndb.put_multi(entities)
return util.redirect('/admin/responses')
def make_affected_commits_public(bug):
"""Make related AffectedCommit entities public."""
to_update = []
query = osv.AffectedCommit.query(osv.AffectedCommit.bug_id == bug.key.id())
for affected_commit in query:
affected_commit.public = True
to_update.append(affected_commit)
if to_update:
ndb.put_multi(to_update)
def _update_employees(employee_dicts):
"""Given a JSON string in the format "[{employee info 1}, {employee info 2}, ...]",
create new employee records and update existing records as necessary.
Then determine whether any employees have been terminated since the last update,
and mark these employees as such.
"""
logging.info('Updating employees...')
db_employee_dict = {
employee.username: employee
for employee in Employee.query()
}
all_employees, new_employees = [], []
current_usernames = set()
for d in employee_dicts:
existing_employee = db_employee_dict.get(d['username'])
if existing_employee is None:
new_employee = Employee.create_from_dict(d, persist=False)
all_employees.append(new_employee)
new_employees.append(new_employee)
elif existing_employee.is_employee_data_changed(d, is_terminated=False):
existing_employee.update_from_dict(d)
all_employees.append(existing_employee)
current_usernames.add(d['username'])
logging.info('Processed {} employees'.format(len(all_employees)))
# Only updates the changed and new employees data
ndb.put_multi(all_employees)
# Figure out if there are any employees in the DB that aren't in the S3
# dump. These are terminated employees, and we need to mark them as such.
db_usernames = set(db_employee_dict.keys())
terminated_usernames = db_usernames - current_usernames
terminated_employees = []
for username in terminated_usernames:
employee = db_employee_dict[username]
# Only updating non-terminated employees
if not employee.terminated:
employee.terminated = True
terminated_employees.append(employee)
ndb.put_multi(terminated_employees)
logging.info('Done updating employees.')
def setUp(self):
self.total_fuzzer_jobs = 7000
self.platforms = ['LINUX', 'WINDOWS']
fuzzer_jobs = []
for platform in self.platforms:
for i in range(self.total_fuzzer_jobs):
fuzzer_job = data_types.FuzzerJob(
fuzzer='libFuzzer',
job='libfuzzer_asan_{}_{:06d}'.format(platform, i),
platform=platform)
fuzzer_jobs.append(fuzzer_job)
ndb.put_multi(fuzzer_jobs)
# Should be removed.
data_types.FuzzerJobs(id='LINUX-2', platform='LINUX').put()
def test_delete_multi_with_transactional(dispose_of):
"""Regression test for issue #271
https://github.com/googleapis/python-ndb/issues/271
"""
N = 10
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
@ndb.transactional()
def delete_them(entities):
ndb.delete_multi([entity.key for entity in entities])
foos = list(range(N))
entities = [SomeKind(foo=foo) for foo in foos]
keys = ndb.put_multi(entities)
dispose_of(*(key._key for key in keys))
entities = ndb.get_multi(keys)
assert [entity.foo for entity in entities] == foos
assert delete_them(entities) is None
entities = ndb.get_multi(keys)
assert entities == [None] * N
def put_multi(entities):
"""Put multiple entities, working around a limitation in the NDB library with
the maximum number of keys allowed."""
result = []
for chunk in _gen_chunks(entities, _MODIFY_BATCH_SIZE):
result.extend(ndb.put_multi(chunk))
return result
def test_query_keys_only() -> None:
stored_keys = ndb.put_multi(
[SimpleModel(id=f"test{i}", int_prop=i) for i in range(10)])
assert len(stored_keys) == 10
resp = SimpleModel.query().fetch(keys_only=True)
assert len(resp) == 10
assert all(isinstance(i, ndb.Key) for i in resp)
def sync_projects(projects):
"""Sync projects with cloud datastore."""
project_query = Project.query()
projects_to_remove = [
project.key for project in project_query
if project.name not in projects
]
ndb.delete_multi(projects_to_remove)
existing_projects = {project.name for project in project_query}
new_projects = [
Project(name=project) for project in projects
if project not in existing_projects
]
ndb.put_multi(new_projects)
def test_query_all() -> None:
stored_keys = ndb.put_multi(
[SimpleModel(id=f"test{i}", int_prop=i) for i in range(10)])
assert len(stored_keys) == 10
query_all = SimpleModel.query().order(SimpleModel.int_prop).fetch()
assert len(query_all) == 10
for i, model in enumerate(query_all):
assert model.int_prop == i
def put_multi(entities: Sequence[Model]) -> List[str]:
"""Stores a sequence of Model instances.
Args:
entities: list(datastore_services.Model). A list of Model instances.
Returns:
list(str). A list with the stored keys.
"""
return ndb.put_multi(list(entities))
def put_multi(entities: List[TYPE_MODEL_SUBCLASS]) -> List[str]:
"""Stores a sequence of Model instances.
Args:
entities: list(datastore_services.Model). A list of Model instances.
Returns:
list(str). A list with the stored keys.
"""
return ndb.put_multi(entities)
def setUp(self):
"""Put data in the local ndb table the tests to query from."""
test_helpers.patch_environ(self)
test_helpers.patch(
self, ['bot.fuzzers.engine_common.decide_with_probability'])
self.mock.decide_with_probability.return_value = True
data = []
strategy1 = data_types.FuzzStrategyProbability()
strategy1.strategy_name = 'corpus_mutations_ml_rnn,corpus_subset,'
strategy1.probability = 1
strategy1.engine = 'afl'
data.append(strategy1)
ndb.put_multi(data)
distribution = fuzz_task.get_strategy_distribution_from_ndb()
environment.set_value('USE_BANDIT_STRATEGY_SELECTION', True)
environment.set_value('STRATEGY_SELECTION_DISTRIBUTION', distribution)
def test_query_in_range() -> None:
stored_keys = ndb.put_multi(
[SimpleModel(id=f"test{i}", int_prop=i) for i in range(10)])
assert len(stored_keys) == 10
resp = (SimpleModel.query(SimpleModel.int_prop >= 1,
SimpleModel.int_prop < 5).order(
SimpleModel.int_prop).fetch())
assert len(resp) == 4
for i, model in enumerate(resp):
assert model.int_prop == i + 1
def import_new_oss_fuzz_entries(self, repo, oss_fuzz_source):
"""Import new entries."""
exported = []
vulnerabilities_path = os.path.join(
osv.repo_path(repo), oss_fuzz_source.directory_path or '')
for bug in osv.Bug.query(
osv.Bug.source_of_truth == osv.SourceOfTruth.INTERNAL):
if bug.status != osv.BugStatus.PROCESSED:
continue
if not bug.public:
continue
source_name, _ = osv.parse_source_id(bug.source_id)
if source_name != oss_fuzz_source.name:
continue
vulnerability_path = os.path.join(vulnerabilities_path,
osv.source_path(bug))
os.makedirs(os.path.dirname(vulnerability_path), exist_ok=True)
if os.path.exists(vulnerability_path):
continue
logging.info('Writing %s', bug.key.id())
osv.vulnerability_to_yaml(bug.to_vulnerability(),
vulnerability_path)
# The source of truth is now this yaml file.
bug.source_of_truth = osv.SourceOfTruth.SOURCE_REPO
exported.append(bug)
# Commit Vulnerability changes back to the oss-fuzz source repository.
repo.index.add_all()
diff = repo.index.diff_to_tree(repo.head.peel().tree)
if not diff:
logging.info('No new entries, skipping committing.')
return
logging.info('Commiting and pushing new entries')
if osv.push_source_changes(repo, 'Import from OSS-Fuzz',
self._git_callbacks(oss_fuzz_source)):
ndb.put_multi(exported)
def test_query_projection() -> None:
stored_keys = ndb.put_multi([
SimpleModel(id=f"test{i}", int_prop=i, str_prop=f"{i}")
for i in range(10)
])
assert len(stored_keys) == 10
resp = SimpleModel.query().fetch(projection=[SimpleModel.int_prop])
assert len(resp) == 10
assert all(i.int_prop is not None for i in resp)
with pytest.raises(ndb.model.UnprojectedPropertyError):
assert all(i.str_prop is None for i in resp)
def get_or_insert(token_id, request_id, path):
"""If a token doesn't exist, create it and record who it was issued to.
This function has a potential bug where 2 or more concurrent requests
cause a token to be issued to more than one request.
The request_id and path are here so that it is easier to see what went
wrong later.
"""
token_obj = Token.get_by_id(token_id)
new_claim = False
if token_obj is None:
token_obj = Token(id=token_id, request_id=request_id, path=path)
# There should be exactly one Claim entity for each Token.
claim_obj = Claim(token_id=token_id, request_id=request_id, path=path)
# Using put_multi here is more efficient than separate datastore puts
# but is not required to be transaction safe.
ndb.put_multi([token_obj, claim_obj])
new_claim = True
return token_obj, new_claim
def process_export(old_export_path, new_export_path):
for table, cls in get_tables():
logging.info('Processing ' + table)
table_suffix = '/WCA_export_' + table + '.tsv'
old_rows = read_table(old_export_path + table_suffix, cls, False)
new_rows = read_table(new_export_path + table_suffix, cls, True)
logging.info('Old: %d' % len(old_rows))
logging.info('New: %d' % len(new_rows))
write_table(new_export_path + table_suffix, new_rows, cls)
modifier = get_modifier(table)
objects_to_put = []
keys_to_delete = []
for key in new_rows:
row = new_rows[key]
if key in old_rows and old_rows[key] == row:
continue
else:
obj = cls(id=key)
obj.ParseFromDict(row)
if modifier:
modifier(obj)
objects_to_put += [obj]
for key, row in old_rows.items():
if key in new_rows:
continue
else:
keys_to_delete += [ndb.Key(cls, key)]
logging.info('Putting %d objects' % len(objects_to_put))
while objects_to_put:
batch_size = 500
subslice = objects_to_put[:batch_size]
objects_to_put = objects_to_put[batch_size:]
ndb.put_multi(subslice)
logging.info('Deleting %d objects' % len(keys_to_delete))
ndb.delete_multi(keys_to_delete)
def test_multi_with_lots_of_keys(dispose_of):
"""Regression test for issue #318.
https://github.com/googleapis/python-ndb/issues/318
"""
N = 1001
class SomeKind(ndb.Model):
foo = ndb.IntegerProperty()
foos = list(range(N))
entities = [SomeKind(foo=foo) for foo in foos]
keys = ndb.put_multi(entities)
dispose_of(*(key._key for key in keys))
assert len(keys) == N
entities = ndb.get_multi(keys)
assert [entity.foo for entity in entities] == foos
ndb.delete_multi(keys)
entities = ndb.get_multi(keys)
assert entities == [None] * N
def test_transactional_composable(dispose_of):
"""Regression test for Issue #366.
https://github.com/googleapis/python-ndb/issues/366
"""
class OtherKind(ndb.Model):
bar = ndb.IntegerProperty()
class SomeKind(ndb.Model):
foos = ndb.KeyProperty(repeated=True)
bar = ndb.IntegerProperty(default=42)
others = [OtherKind(bar=bar) for bar in range(5)]
other_keys = ndb.put_multi(others)
for key in other_keys:
dispose_of(key._key)
entity = SomeKind(foos=other_keys[1:])
entity_key = entity.put()
dispose_of(entity_key._key)
@ndb.transactional()
def get_entities(*keys):
entities = []
for entity in ndb.get_multi(keys):
entities.append(entity)
if isinstance(entity, SomeKind):
entities.extend(get_foos(entity))
return entities
@ndb.transactional()
def get_foos(entity):
return ndb.get_multi(entity.foos)
results = get_entities(entity_key, other_keys[0])
assert [result.bar for result in results] == [42, 1, 2, 3, 4, 0]
def test_multi_get_weirdness_with_redis(dispose_of):
"""Regression test for issue #294.
https://github.com/googleapis/python-ndb/issues/294
"""
class SomeKind(ndb.Model):
foo = ndb.StringProperty()
objects = [SomeKind(foo=str(i)) for i in range(10)]
keys = ndb.put_multi(objects)
for key in keys:
dispose_of(key._key)
ndb.get_multi(keys)
one_object = random.choice(keys).get()
one_object.foo = "CHANGED"
one_object.put()
objects_upd = ndb.get_multi(keys)
keys_upd = [obj.key for obj in objects_upd]
assert len(keys_upd) == len(keys)
assert len(set(keys_upd)) == len(set(keys))
assert set(keys_upd) == set(keys)
def _throttled_put(to_put):
"""Throttled ndb put."""
for batch, is_last in _batcher(to_put, _DATASTORE_BATCH_SIZE):
ndb.put_multi(batch)
if not is_last:
time.sleep(_DATASTORE_BATCH_SLEEP)
def operate_on_multiple_keys_at_once(list_of_entities):
list_of_keys = ndb.put_multi(list_of_entities)
list_of_entities = ndb.get_multi(list_of_keys)
ndb.delete_multi(list_of_keys)
def execute(args):
"""Build keywords for jobs."""
jobs = list(data_types.Job.query())
if args.non_dry_run:
ndb.put_multi(jobs)
print("Done building keywords for jobs.")
