def testFetchPage_withCursor(self):
query = datastore_entities.LabInfo.query()
query = query.order(datastore_entities.LabInfo.key)
labs, _, cursor = datastore_util.FetchPage(query, 2)
self.assertEqual(2, len(labs))
labs, prev_cursor, next_cursor = datastore_util.FetchPage(
query, 2, page_cursor=cursor)
self.assertEqual(2, len(labs))
self.assertEqual('lab2', labs[0].lab_name)
self.assertEqual('lab3', labs[1].lab_name)
self.assertIsNotNone(next_cursor)
self.assertEqual(cursor, prev_cursor)
def testFetchPage_backwardsWithAncestorQuery(self):
self.ndb_host_0 = datastore_test_util.CreateHost(
cluster='free',
hostname='host_0',
timestamp=self.TIMESTAMP,
host_state=api_messages.HostState.RUNNING,
)
self.ndb_host_1 = datastore_test_util.CreateHost(
cluster='paid',
hostname='host_1',
timestamp=self.TIMESTAMP,
device_count_timestamp=self.TIMESTAMP,
)
self._CreateHostInfoHistory(self.ndb_host_1).put()
self.ndb_host_1.host_state = api_messages.HostState.UNKNOWN
self.ndb_host_1.timestamp += datetime.timedelta(hours=1)
self._CreateHostInfoHistory(self.ndb_host_1).put()
self.ndb_host_1.host_state = api_messages.HostState.GONE
self.ndb_host_1.timestamp += datetime.timedelta(hours=1)
self._CreateHostInfoHistory(self.ndb_host_1).put()
self._CreateHostInfoHistory(self.ndb_host_0).put()
self.ndb_host_0.host_state = api_messages.HostState.KILLING
self.ndb_host_0.timestamp += datetime.timedelta(hours=1)
self._CreateHostInfoHistory(self.ndb_host_0).put()
self.ndb_host_0.host_state = api_messages.HostState.GONE
self.ndb_host_0.timestamp += datetime.timedelta(hours=1)
self._CreateHostInfoHistory(self.ndb_host_0).put()
# First page
query = (datastore_entities.HostInfoHistory.query(ancestor=ndb.Key(
datastore_entities.HostInfo, self.ndb_host_0.hostname)).order(
-datastore_entities.HostInfoHistory.timestamp))
histories, prev_cursor, next_cursor = datastore_util.FetchPage(
query, 2)
self.assertEqual(2, len(histories))
self.assertIsNone(prev_cursor)
self.assertIsNotNone(next_cursor)
# Back to first page (ancestor query with backwards)
histories, prev_cursor, next_cursor = datastore_util.FetchPage(
query, 2, backwards=True)
self.assertEqual(2, len(histories))
self.assertEqual(self.ndb_host_0.hostname, histories[0].hostname)
self.assertEqual(api_messages.HostState.GONE, histories[0].host_state)
self.assertEqual(self.ndb_host_0.hostname, histories[1].hostname)
self.assertEqual(api_messages.HostState.KILLING,
histories[1].host_state)
def testFetchPage(self):
query = datastore_entities.LabInfo.query()
query = query.order(datastore_entities.LabInfo.key)
labs, _, _ = datastore_util.FetchPage(query, 2)
self.assertEqual(2, len(labs))
self.assertEqual('lab0', labs[0].lab_name)
self.assertEqual('lab1', labs[1].lab_name)
def ListNotes(self, request):
"""List notes of a host.
Args:
request: an API request.
Returns:
an api_messages.NoteCollection object.
"""
query = (
datastore_entities.Note.query()
.filter(datastore_entities.Note.hostname == request.hostname)
.order(-datastore_entities.Note.timestamp))
if not request.include_device_notes:
query = query.filter(
datastore_entities.Note.type == common.NoteType.HOST_NOTE)
note_entities, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor, backwards=request.backwards)
note_msgs = [
datastore_entities.ToMessage(entity) for entity in note_entities
]
return api_messages.NoteCollection(
notes=note_msgs,
more=bool(next_cursor),
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def ListSummaries(self, request):
"""Fetches a page of test run summaries.
Parameters:
max_results: Maximum number of test runs to return
page_token: Token for pagination
backwards: True to fetch previous page of results
labels: Labels to look for
test: Test name to look for
device_spec: A device spec to look for
state: List of test run states to include
device_build_id: Build ID to look for
test_package_info: Test package to look for
prev_test_run_id: Previous test run ID
filter_query: Additional filters to apply (applied to test name, run
target, labels, test package name and version, device build ID and
product)
"""
query = ndb_models.TestRunSummary.query().order(
-ndb_models.TestRunSummary.create_time,
ndb_models.TestRunSummary.key)
query = self._ApplyQueryFilters(query, request)
result_filter = self._BuildTestRunFilterFunction(request)
test_runs, prev_cursor, next_cursor = datastore_util.FetchPage(
query,
request.max_results,
page_cursor=request.page_token,
backwards=request.backwards,
result_filter=result_filter)
return mtt_messages.TestRunSummaryList(
test_runs=mtt_messages.ConvertList(test_runs,
mtt_messages.TestRunSummary),
prev_page_token=prev_cursor,
next_page_token=next_cursor)
def ListTestHarnessImages(self, request):
"""List test harness image metadata.
Args:
request: an API request.
Returns:
An api_messages.TestHarnessImageMetadataCollection object.
"""
query = datastore_entities.TestHarnessImageMetadata.query()
if request.tag_prefix:
query = (
query.filter(datastore_entities.TestHarnessImageMetadata.current_tags
>= request.tag_prefix).filter(
datastore_entities.TestHarnessImageMetadata
.current_tags < request.tag_prefix + u"\ufffd")
.order(-datastore_entities.TestHarnessImageMetadata.current_tags))
query = query.order(
-datastore_entities.TestHarnessImageMetadata.create_time)
entities, _, next_cursor = datastore_util.FetchPage(query, request.count,
request.cursor)
image_msgs = [datastore_entities.ToMessage(entity) for entity in entities]
return api_messages.TestHarnessImageMetadataCollection(
images=image_msgs, next_cursor=next_cursor)
def testFetchPage_backwards(self):
query = datastore_entities.LabInfo.query()
query = query.order(datastore_entities.LabInfo.key)
# Fetch 3 results
labs, _, cursor = datastore_util.FetchPage(query, 3)
self.assertEqual(3, len(labs))
self.assertEqual('lab0', labs[0].lab_name)
self.assertEqual('lab1', labs[1].lab_name)
self.assertEqual('lab2', labs[2].lab_name)
# Fetch 2 results backwards
labs, prev_cursor, next_cursor = datastore_util.FetchPage(
query, 2, page_cursor=cursor, backwards=True)
self.assertEqual(2, len(labs))
self.assertEqual('lab1', labs[0].lab_name)
self.assertEqual('lab2', labs[1].lab_name)
self.assertIsNotNone(prev_cursor)
self.assertEqual(cursor, next_cursor)
def testFetchPage_backwardsCountLargeThanRest(self):
# When a page is fetched backwards with a cursor, and the remaining results
# are smaller than one page, the first page should be fetched instead.
query = datastore_entities.LabInfo.query()
query = query.order(datastore_entities.LabInfo.key)
# Fetch 2 results
labs, _, cursor = datastore_util.FetchPage(query, 2)
self.assertEqual(2, len(labs))
self.assertEqual('lab0', labs[0].lab_name)
self.assertEqual('lab1', labs[1].lab_name)
# Fetch 3 results backwards
labs, prev_cursor, next_cursor = datastore_util.FetchPage(
query, 3, page_cursor=cursor, backwards=True)
self.assertEqual(3, len(labs))
self.assertEqual('lab0', labs[0].lab_name)
self.assertEqual('lab1', labs[1].lab_name)
self.assertEqual('lab2', labs[2].lab_name)
self.assertIsNone(prev_cursor)
self.assertIsNotNone(next_cursor)
def testFetchPage_withResultFilter(self):
"""Tests that a predicate function can be applied to the query results."""
even_lab_number_re = re.compile('^lab[02468]$')
def _Filter(lab):
return even_lab_number_re.search(lab.lab_name)
query = datastore_entities.LabInfo.query()
query = query.order(datastore_entities.LabInfo.key)
labs, _, _ = datastore_util.FetchPage(query, 4, result_filter=_Filter)
self.assertEqual(4, len(labs))
self.assertEqual('lab0', labs[0].lab_name)
self.assertEqual('lab2', labs[1].lab_name)
self.assertEqual('lab4', labs[2].lab_name)
self.assertEqual('lab6', labs[3].lab_name)
def _ListLabs(self, request):
"""ListLabs without owner filter. Some labs don't have config."""
query = datastore_entities.LabInfo.query()
query = query.order(datastore_entities.LabInfo.key)
labs, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, page_cursor=request.cursor)
lab_config_keys = [
ndb.Key(datastore_entities.LabConfig, lab.lab_name) for lab in labs]
lab_configs = ndb.get_multi(lab_config_keys)
lab_infos = [datastore_entities.ToMessage(lab, lab_config)
for lab, lab_config in zip(labs, lab_configs)]
return api_messages.LabInfoCollection(
lab_infos=lab_infos,
more=bool(next_cursor),
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def _ListLabsByOwner(self, request):
"""Lab owners are in LabConfig, so need to query by LabConfig."""
query = datastore_entities.LabConfig.query()
query = query.order(datastore_entities.LabConfig.key)
query = query.filter(
datastore_entities.LabConfig.owners == request.owner)
lab_configs, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, page_cursor=request.cursor)
lab_info_msgs = []
for lab_config in lab_configs:
lab_info_msgs.append(
api_messages.LabInfo(
lab_name=lab_config.lab_name,
owners=(lab_config.owners or [])))
return api_messages.LabInfoCollection(
lab_infos=lab_info_msgs,
more=bool(next_cursor),
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def ListDeviceBlocklist(self, request):
"""List device blocklists.
Args:
request: API request with key id.
Returns:
a list of DeviceBlocklistMessage objects.
"""
query = (datastore_entities.DeviceBlocklist.query().order(
-datastore_entities.DeviceBlocklist.create_timestamp))
device_blocklists, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor, backwards=request.backwards)
msgs = [
datastore_entities.ToMessage(blocklist)
for blocklist in device_blocklists
]
return api_messages.DeviceBlocklistCollection(device_blocklists=msgs,
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def ListHostConfigs(self, request):
"""List host configs.
Args:
request: an API request.
Returns:
an api_messages.HostConfigCollection object.
"""
query = datastore_entities.HostConfig.query()
if request.lab_name:
query = query.filter(
datastore_entities.HostConfig.lab_name == request.lab_name)
host_configs, _, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor)
host_config_msgs = [datastore_entities.ToMessage(host_config)
for host_config in host_configs]
return api_messages.HostConfigCollection(
host_configs=host_config_msgs, next_cursor=next_cursor)
def ListHistories(self, request):
"""List histories of a host.
Args:
request: an API request.
Returns:
an api_messages.HostInfoHistoryCollection object.
"""
query = (
datastore_entities.HostInfoHistory.query(
ancestor=ndb.Key(datastore_entities.HostInfo, request.hostname))
.order(-datastore_entities.HostInfoHistory.timestamp))
histories, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor, backwards=request.backwards)
history_msgs = [
datastore_entities.ToMessage(entity) for entity in histories
]
return api_messages.HostInfoHistoryCollection(
histories=history_msgs,
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def ListHistories(self, request):
"""List histories of a device.
Args:
request: an API request.
Returns:
an api_messages.DeviceInfoHistoryCollection object.
"""
device = device_manager.GetDevice(device_serial=request.device_serial)
query = (
datastore_entities.DeviceInfoHistory.query(ancestor=device.key).order(
-datastore_entities.DeviceInfoHistory.timestamp))
histories, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor, backwards=request.backwards)
history_msgs = [
datastore_entities.ToMessage(entity) for entity in histories
]
return api_messages.DeviceInfoHistoryCollection(
histories=history_msgs,
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def ListDevices(self, request):
"""Fetches a list of devices from NDB.
Args:
request: an API request.
Returns:
a DeviceInfoCollection object.
"""
logging.debug("ClusterDeviceApi.NDBListDevices request: %s", request)
query = datastore_entities.DeviceInfo.query().order(
datastore_entities.DeviceInfo.key)
if request.lab_name:
query = query.filter(
datastore_entities.DeviceInfo.lab_name == request.lab_name)
if request.cluster_id:
query = query.filter(
datastore_entities.DeviceInfo.clusters == request.cluster_id)
if request.hostname:
query = query.filter(
datastore_entities.DeviceInfo.hostname == request.hostname)
# We only consider device hidden here, since there is no simple way to do
# join like operation in datastore. We tried fetching host and use
# in(hostnames), but it was not scalable at all.
if not request.include_hidden:
query = query.filter(datastore_entities.DeviceInfo.hidden == False)
if request.product:
query = query.filter(
datastore_entities.DeviceInfo.product == request.product)
if request.device_serial:
query = query.filter(
datastore_entities.DeviceInfo.device_serial == request.device_serial)
if request.flated_extra_info:
query = query.filter(datastore_entities.DeviceInfo.flated_extra_info ==
request.flated_extra_info)
start_time = time.time()
if len(request.pools) == 1:
query = query.filter(
datastore_entities.DeviceInfo.pools == request.pools[0])
if len(request.device_states) == 1:
query = query.filter(
datastore_entities.DeviceInfo.state == request.device_states[0])
if len(request.host_groups) == 1:
query = query.filter(
datastore_entities.DeviceInfo.host_group == request.host_groups[0])
if len(request.device_types) == 1:
query = query.filter(
datastore_entities.DeviceInfo.device_type == request.device_types[0])
test_harnesses = request.test_harness + request.test_harnesses
if len(test_harnesses) == 1:
query = query.filter(
datastore_entities.DeviceInfo.test_harness == test_harnesses[0])
if len(request.hostnames) == 1:
query = query.filter(
datastore_entities.DeviceInfo.hostname == request.hostnames[0])
if len(request.run_targets) == 1:
query = query.filter(
datastore_entities.DeviceInfo.run_target == request.run_targets[0])
def _PostFilter(device):
if (request.pools and
not set(device.pools).intersection(set(request.pools))):
return
if request.device_states and device.state not in request.device_states:
return
if request.host_groups and device.host_group not in request.host_groups:
return
if not request.include_offline_devices and \
device.state not in common.DEVICE_ONLINE_STATES:
return
if request.device_types and \
device.device_type not in request.device_types:
return
if test_harnesses and device.test_harness not in test_harnesses:
return
if request.hostnames and device.hostname not in request.hostnames:
return
if request.run_targets and device.run_target not in request.run_targets:
return
return True
devices, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor, result_filter=_PostFilter)
logging.debug("Fetched %d devices in %r seconds.", len(devices),
time.time() - start_time)
start_time = time.time()
device_infos = [datastore_entities.ToMessage(d) for d in devices]
logging.debug("Tranformed devices to messages in %r seconds.",
time.time() - start_time)
return api_messages.DeviceInfoCollection(
device_infos=device_infos,
next_cursor=next_cursor,
prev_cursor=prev_cursor,
more=bool(next_cursor))
def ListHosts(self, request):
"""Fetches a list of hosts.
Args:
request: an API request.
Returns:
a HostInfoCollection object.
"""
if ((request.timestamp and not request.timestamp_operator) or
(not request.timestamp and request.timestamp_operator)):
raise endpoints.BadRequestException(
'"timestamp" and "timestamp_operator" must be set at the same time.')
query = datastore_entities.HostInfo.query()
if request.lab_name:
query = query.filter(
datastore_entities.HostInfo.lab_name == request.lab_name)
if request.assignee:
query = query.filter(
datastore_entities.HostInfo.assignee == request.assignee)
if request.is_bad is not None:
query = query.filter(datastore_entities.HostInfo.is_bad == request.is_bad)
if not request.include_hidden:
query = query.filter(datastore_entities.HostInfo.hidden == False)
if request.flated_extra_info:
query = query.filter(datastore_entities.HostInfo.flated_extra_info ==
request.flated_extra_info)
if len(request.host_groups) == 1:
query = query.filter(
datastore_entities.HostInfo.host_group == request.host_groups[0])
if len(request.hostnames) == 1:
query = query.filter(
datastore_entities.HostInfo.hostname == request.hostnames[0])
test_harnesses = request.test_harness + request.test_harnesses
if len(test_harnesses) == 1:
query = query.filter(
datastore_entities.HostInfo.test_harness == test_harnesses[0])
if len(request.test_harness_versions) == 1:
query = query.filter(
datastore_entities.HostInfo.test_harness_version ==
request.test_harness_versions[0])
if len(request.pools) == 1:
query = query.filter(
datastore_entities.HostInfo.pools == request.pools[0])
if len(request.host_states) == 1:
query = query.filter(
datastore_entities.HostInfo.host_state == request.host_states[0])
if len(request.recovery_states) == 1:
query = query.filter(
datastore_entities.HostInfo.recovery_state
== request.recovery_states[0])
hostnames_with_requested_update_states = set()
if request.host_update_states:
update_state_query = datastore_entities.HostUpdateState.query().filter(
datastore_entities.HostUpdateState.state.IN(
request.host_update_states))
hostnames_with_requested_update_states = set(
update_state.hostname for update_state in update_state_query.fetch(
projection=[datastore_entities.HostUpdateState.hostname]))
def _PostFilter(host):
if request.host_groups and host.host_group not in request.host_groups:
return
if request.hostnames and host.hostname not in request.hostnames:
return
if (test_harnesses and
host.test_harness not in test_harnesses):
return
if (request.test_harness_versions and
host.test_harness_version not in request.test_harness_versions):
return
if request.pools and not set(host.pools).intersection(set(request.pools)):
return
if request.host_states and host.host_state not in request.host_states:
return
if (request.recovery_states and
host.recovery_state not in request.recovery_states):
return
if request.timestamp:
if not host.timestamp:
return
return _CheckTimestamp(
host.timestamp, request.timestamp_operator, request.timestamp)
if request.host_update_states:
if host.hostname not in hostnames_with_requested_update_states:
return
return True
if request.timestamp:
query = query.order(
datastore_entities.HostInfo.timestamp,
datastore_entities.HostInfo.key)
else:
query = query.order(datastore_entities.HostInfo.key)
hosts, prev_cursor, next_cursor = datastore_util.FetchPage(
query, request.count, request.cursor, result_filter=_PostFilter)
host_update_state_keys = [
ndb.Key(datastore_entities.HostUpdateState, host.hostname)
for host in hosts]
host_update_states = ndb.get_multi(host_update_state_keys)
host_infos = []
for host, host_update_state in zip(hosts, host_update_states):
devices = []
if request.include_devices:
device_query = datastore_entities.DeviceInfo.query(ancestor=host.key)
if not request.include_hidden:
device_query = device_query.filter(
datastore_entities.DeviceInfo.hidden == False) devices = device_query.fetch()
host_infos.append(datastore_entities.ToMessage(
host, devices=devices,
host_update_state_entity=host_update_state))
return api_messages.HostInfoCollection(
host_infos=host_infos,
more=bool(next_cursor),
next_cursor=next_cursor,
prev_cursor=prev_cursor)
def testFetchPage_lastPage(self):
"""Tests that fetching the last page doesn't produce a next cursor."""
query = datastore_entities.LabInfo.query()
labs, _, next_cursor = datastore_util.FetchPage(query, 99)
self.assertEqual(10, len(labs))
self.assertIsNone(next_cursor)