def _UpdateHostConfigs(host_config_pbs, cluster_config_pb, lab_config_pb): """Update host configs in HostInfo entities to ndb. Args: host_config_pbs: a list of host config protos. cluster_config_pb: the cluster config proto. lab_config_pb: the lab config proto. """ logging.debug('Updating host configs for <lab: %s, cluster: %s.>', lab_config_pb.lab_name, cluster_config_pb.cluster_name) host_config_keys = [] for host_config_pb in host_config_pbs: host_config_keys.append( ndb.Key(datastore_entities.HostConfig, host_config_pb.hostname)) entities = ndb.get_multi(host_config_keys) entities_to_update = [] # Update the exist host config entity. for entity, host_config_pb in zip(entities, host_config_pbs): host_config_msg = lab_config_util.HostConfig( host_config_pb, cluster_config_pb, lab_config_pb) new_host_config_entity = datastore_entities.HostConfig.FromMessage( host_config_msg) if not _CheckConfigEntitiesEqual(entity, new_host_config_entity): logging.debug('Updating host config entity: %s.', new_host_config_entity) entities_to_update.append(new_host_config_entity) ndb.put_multi(entities_to_update) logging.debug('Host configs updated.')
def testGetPredefinedMessage_OK(self): lab_name = "alab" predefined_message_entities = [ datastore_entities.PredefinedMessage( lab_name=lab_name, type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, content="offline_reason1", used_count=2), datastore_entities.PredefinedMessage( lab_name=lab_name, type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content="recovery_action1", used_count=5), ] ndb.put_multi(predefined_message_entities) message_1 = note_manager.GetPredefinedMessage( api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, lab_name, "offline_reason1") self.assertEqual("offline_reason1", message_1.content) self.assertEqual(api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, message_1.type) message_2 = note_manager.GetPredefinedMessage( api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, lab_name, "recovery_action1") self.assertEqual("recovery_action1", message_2.content) self.assertEqual(api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, message_2.type)
def _SetHostRecoveryState(hostname, recovery_state, assignee=None): """Set host's recovery state.""" host = GetHost(hostname) if not host: logging.error("Host %s doesn't exist.", hostname) return host.assignee = assignee entities_to_update = [] if recovery_state == common.RecoveryState.VERIFIED: host.recovery_state = common.RecoveryState.VERIFIED host.assignee = host.assignee host.last_recovery_time = common.Now() host.timestamp = common.Now() entities_to_update.append(_CreateHostInfoHistory(host)) devices = GetDevicesOnHost(hostname) for device in devices: if (device.recovery_state and device.recovery_state != common.RecoveryState.UNKNOWN): entities_to_update.extend( _BuildDeviceRecoveryState(device, common.RecoveryState.VERIFIED)) # After it's verified, it goes into a state as no-one is recovering it. host.recovery_state = common.RecoveryState.UNKNOWN host.assignee = None else: host.recovery_state = recovery_state host.last_recovery_time = common.Now() host.timestamp = common.Now() entities_to_update.append(_CreateHostInfoHistory(host)) entities_to_update.append(host) ndb.put_multi(entities_to_update)
def UpdateGoneHost(hostname): """Set a host and its devices to GONE.""" logging.info("Set host %s and its devices to GONE.", hostname) host = GetHost(hostname) if host.host_state == api_messages.HostState.GONE: logging.info("Host %s is already GONE.", hostname) return entities_to_update = [] now = common.Now() host_state_history, host_history = _UpdateHostState( host, api_messages.HostState.GONE, now) entities_to_update.append(host) if host_state_history: entities_to_update.append(host_state_history) if host_history: entities_to_update.append(host_history) devices = GetDevicesOnHost(hostname) for device in devices or []: if device.state == common.DeviceState.GONE: continue logging.debug("Set device %s to GONE.", device.device_serial) device_state_history, device_history = _UpdateDeviceState( device, common.DeviceState.GONE, now) entities_to_update.append(device) if device_state_history: entities_to_update.append(device_state_history) if device_history: entities_to_update.append(device_history) _DoCountDeviceForHost(host, devices) ndb.put_multi(entities_to_update)
def testSyncInventoryGroupsToNdbHostGroupConfig(self): ndb.put_multi([ datastore_entities.HostGroupConfig( id='foo_jump', lab_name='foo', parent_groups=['foo_all', 'foo_bar']), datastore_entities.HostGroupConfig(id='foo_bar', lab_name='foo', parent_groups=[ 'foo_all', ]) ]) config_syncer_gcs_to_ndb.SyncInventoryGroupsToNDB() ndb.get_context().clear_cache() res = datastore_entities.HostGroupConfig.query( datastore_entities.HostGroupConfig.lab_name == 'foo').fetch() for g in res: if g.name is None: self.assertIsNone(g) group_map = {g.name: g for g in res} self.assertLen(group_map, 10) self.assertSameElements(group_map['all'].parent_groups, []) self.assertSameElements(group_map['jump'].parent_groups, ['server']) self.assertSameElements(group_map['dhcp'].parent_groups, ['server']) self.assertSameElements(group_map['pxe'].parent_groups, ['server']) self.assertSameElements(group_map['server'].parent_groups, []) self.assertSameElements(group_map['dtf'].parent_groups, ['tf']) self.assertSameElements(group_map['storage_tf'].parent_groups, ['tf'])
def _UpdateHostConfigByInventoryData(lab_name, data): """Updates HostConfig inventory_groups. Args: lab_name: the lab name. data: the InventoryData object. """ keys = [] entities_from_file = { host.name: _CreateHostConfigEntityFromHostInventory(lab_name, host) for host in data.hosts.values() } for hostname in entities_from_file: keys.append(ndb.Key(datastore_entities.HostConfig, hostname)) entities_from_ndb = {} for entity in ndb.get_multi(keys): if entity: entities_from_ndb[entity.hostname] = entity need_update = {} for hostname in entities_from_file: old = entities_from_ndb.get(hostname) new = entities_from_file[hostname] if not old: logging.debug('Creating host config %s', new) need_update[hostname] = new elif old.inventory_groups != new.inventory_groups: old.inventory_groups = new.inventory_groups logging.debug('Updating host config %s', old) need_update[hostname] = old ndb.put_multi(need_update.values())
def _UpdateHostGroupConfigByInventoryData(lab_name, data): """Updates HostGroupConfig. The methods load new HostGroupConfig from inventory file, creates new HostGroupConfig, updates exist HostGroupConfig and remove obsolete HostGroupConfig. Args: lab_name: the lab name. data: inventory data which was parsed by the ansible inventory module. """ entity_from_file = { group.key: group for group in [ _CreateHostGroupConfigEntityFromGroupInventory(lab_name, group) for group in data.groups.values() ] } logging.debug('Loaded %d HostGroupConfigs of lab %s', len(entity_from_file), lab_name) entity_from_ndb = { g.key: g for g in datastore_entities.HostGroupConfig.query( datastore_entities.HostGroupConfig.lab_name == lab_name).fetch() } need_update = [] for key, group in entity_from_file.items(): if not _CheckConfigEntitiesEqual(entity_from_ndb.get(key), group): need_update.append(group) ndb.put_multi(need_update) logging.debug('Updated %d HostGroupConfigs.', len(need_update)) need_delete = set(entity_from_ndb.keys()).difference(entity_from_file.keys()) ndb.delete_multi(need_delete) logging.debug('Removed %d HostGroupConfigs', len(need_delete))
def _UpdateHostWithHostChangedEvent(event): """update the host with a host state changed event. Args: event: HostEvent object. """ host = GetHost(event.hostname) if not host: host = datastore_entities.HostInfo(id=event.hostname) # For HostStateChangedEvent other than hostname, state and timestamp, # everything else are optional. host.hostname = event.hostname host.lab_name = event.lab_name or host.lab_name or common.UNKNOWN_LAB_NAME host.timestamp = event.timestamp host.test_harness = event.test_harness or host.test_harness host.test_harness_version = ( event.test_harness_version or host.test_harness_version) host.extra_info = event.data or host.extra_info host.hidden = False host_state_history, host_history = _UpdateHostState( host, event.host_state, event.timestamp) entities_to_update = [host] if host_state_history: entities_to_update.append(host_state_history) if host_history: entities_to_update.append(host_history) ndb.put_multi(entities_to_update) return
def _UpdateClusterConfigs(cluster_configs): """Update cluster configs in ClusterInfo entity to ndb. Args: cluster_configs: a list of cluster configs proto. """ logging.debug('Updating cluster configs.') cluster_config_keys = set() for cluster_config in cluster_configs: cluster_config_keys.add( ndb.Key(datastore_entities.ClusterConfig, cluster_config.cluster_name)) entities = ndb.get_multi(cluster_config_keys) name_to_cluster = {} for entity in entities: if entity: name_to_cluster[entity.cluster_name] = entity name_to_entity = {} for cluster_config in cluster_configs: new_config_entity = datastore_entities.ClusterConfig.FromMessage( cluster_config) if not _CheckConfigEntitiesEqual( name_to_cluster.get(cluster_config.cluster_name), new_config_entity): logging.debug('Updating cluster config entity: %s.', new_config_entity) if cluster_config.cluster_name in name_to_entity: logging.warning( '%s has duplicated configs.', cluster_config.cluster_name) name_to_entity[cluster_config.cluster_name] = new_config_entity ndb.put_multi(name_to_entity.values()) logging.debug('Cluster configs updated.')
def _SetDeviceRecoveryState( hostname, device_serial, recovery_state): """Set device's recovery state.""" device = GetDevice(hostname, device_serial) if not device: logging.error("Device (%s, %s) doesn't exist.", hostname, device_serial) return entities_to_update = _BuildDeviceRecoveryState(device, recovery_state) ndb.put_multi(entities_to_update)
def SyncHarnessImageMetadata(): """The job to read image manifest from GCR and write to NDB.""" if not env_config.CONFIG.should_sync_harness_image: return common.HTTP_OK logging.debug('"should_sync_harness_image" is enabled.') creds = auth.GetComputeEngineCredentials() try: response = requests.get( url=_LIST_TAGS_URL, headers={'Authorization': _AUTHORIZATION_TMPL.format(creds.token)}) except requests.exceptions.HTTPError as http_error: logging.exception( 'Error in http request to get image manifest from GCR.') raise http_error response_json = response.json() manifest = response_json.get('manifest') if not manifest: raise KeyError('No valid image manifest is in the response: {}'.format( response_json)) time_now = datetime.datetime.utcnow() entities_to_update = [] for digest, data in manifest.items(): current_tags = data.get('tag', []) analysis_result = _AnalyseTradefedDockerImageTags(current_tags) historical_tags = [] if analysis_result['is_historical_golden']: historical_tags.append(_GOLDEN_TAG) key = ndb.Key( datastore_entities.TestHarnessImageMetadata, _IMAGE_METADATA_KEY_TMPL.format(_DEFAULT_TRADEFED_REPO, digest)) time_created = datetime.datetime.utcfromtimestamp( int(data['timeCreatedMs']) / 1000) entity = datastore_entities.TestHarnessImageMetadata( key=key, repo_name=_DEFAULT_TRADEFED_REPO, digest=digest, test_harness=_TRADEFED_HARNESS_NAME, test_harness_version=analysis_result['tradefed_version_number'], current_tags=current_tags, historical_tags=historical_tags, create_time=time_created, sync_time=time_now) entities_to_update.append(entity) ndb.put_multi(entities_to_update) # Delete stale image metadata datastore_util.DeleteEntitiesUpdatedEarlyThanSomeTimeAgo( datastore_entities.TestHarnessImageMetadata, datastore_entities.TestHarnessImageMetadata.sync_time, _STALE_METADATA_MAX_AGE) return common.HTTP_OK
def _UpdateLabs(clusters): """Update lab NDB entities based on hosts. Args: clusters: a list of ClusterInfo. 1. Add lab if the lab doesn't exist yet. 2. Refresh the host update state summary in all labs based on the underlying host groups. """ logging.info('Updating labs') labs_query = datastore_entities.LabInfo.query() labs_by_lab_names = {lab.lab_name: lab for lab in labs_query} clusters_by_lab_names = collections.defaultdict(list) for cluster_info in clusters: lab_name = cluster_info.lab_name or common.UNKNOWN_LAB_NAME clusters_by_lab_names[lab_name].append(cluster_info) labs = [] for lab_name, cluster_infos in clusters_by_lab_names.items(): lab_host_update_state_summary = datastore_entities.HostUpdateStateSummary() lab_host_update_state_summaries_by_version = collections.defaultdict( datastore_entities.HostUpdateStateSummary) host_count_by_harness_version = collections.Counter() for cluster_info in cluster_infos: if cluster_info and cluster_info.host_update_state_summary: lab_host_update_state_summary += cluster_info.host_update_state_summary if cluster_info and cluster_info.host_update_state_summaries_by_version: for summary in cluster_info.host_update_state_summaries_by_version: lab_host_update_state_summaries_by_version[ summary.target_version] += summary if cluster_info and cluster_info.host_count_by_harness_version: host_count_by_harness_version += collections.Counter( cluster_info.host_count_by_harness_version) if lab_name in labs_by_lab_names: lab = labs_by_lab_names[lab_name] else: lab = datastore_entities.LabInfo( id=lab_name, lab_name=lab_name) lab.populate( host_update_state_summary=lab_host_update_state_summary, host_count_by_harness_version=host_count_by_harness_version, host_update_state_summaries_by_version=list( lab_host_update_state_summaries_by_version.values()), update_timestamp=_Now()) labs.append(lab) ndb.put_multi(labs) logging.info('Updated labs.')
def testListPredefinedMessages_filtersAndOrdering(self): """Test list PredefinedMessages.""" pred_msg_entities = [ datastore_entities.PredefinedMessage( lab_name='lab-name-1', type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content='content-1', used_count=2), datastore_entities.PredefinedMessage( lab_name='lab-name-2', type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content='content-2', used_count=1), datastore_entities.PredefinedMessage( lab_name='lab-name-2', type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content='content-3', used_count=3), datastore_entities.PredefinedMessage( lab_name='lab-name-4', type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, content='content-4', used_count=3), ] ndb.put_multi(pred_msg_entities) api_request = { 'type': 'DEVICE_RECOVERY_ACTION', 'lab_name': 'lab-name-2' } api_response = self.testapp.post_json( '/_ah/api/PredefinedMessageApi.ListPredefinedMessages', api_request) pred_msgs = protojson.decode_message( api_messages.PredefinedMessageCollection, api_response.body).predefined_messages # The results are filtered by type and lab name. self.assertEqual(2, len(pred_msgs)) # The results are sorted by count in descending order. self.assertEqual(pred_msg_entities[1].lab_name, pred_msgs[1].lab_name) self.assertEqual(pred_msg_entities[1].content, pred_msgs[1].content) self.assertEqual(pred_msg_entities[1].type, pred_msgs[1].type) self.assertEqual(pred_msg_entities[1].used_count, pred_msgs[1].used_count) self.assertEqual(pred_msg_entities[2].lab_name, pred_msgs[0].lab_name) self.assertEqual(pred_msg_entities[2].content, pred_msgs[0].content) self.assertEqual(pred_msg_entities[2].type, pred_msgs[0].type) self.assertEqual(pred_msg_entities[2].used_count, pred_msgs[0].used_count)
def _UpdateHostWithDeviceSnapshotEvent(event): """update the host if the event is host info. Update host state to RUNNING if the olds state is GONE. Args: event: HostEvent dictionary. Returns: a HostEntity. """ host = GetHost(event.hostname) if not host: host = datastore_entities.HostInfo(id=event.hostname) host.hostname = event.hostname host.lab_name = event.lab_name # TODO: deprecate physical_cluster, use host_group. host.physical_cluster = event.cluster_id host.host_group = event.host_group host.timestamp = event.timestamp host.test_harness = event.test_harness host.test_harness_version = event.test_harness_version host.hidden = False # TODO: deprecate clusters, use pools. if event.cluster_id: host.clusters = [event.cluster_id] + event.next_cluster_ids else: host.clusters = event.next_cluster_ids[:] host.pools = event.pools entities_to_update = [host] if _IsNewTestHarnessInstance(host, event): # If it's a new instance, we change the host state to RUNNING. host_state_history, host_history = _UpdateHostState( host, api_messages.HostState.RUNNING, event.timestamp) if host_state_history: entities_to_update.append(host_state_history) if host_history: entities_to_update.append(host_history) # Extra info need to be update after checking _IsNewTestHarnessInstance, # since we use insit_harness_start_time_ms in extra info. host.extra_info = event.data ndb.put_multi(entities_to_update) return host
def _UpdateHostUpdateStateWithEvent( event, target_version=common.UNKNOWN_TEST_HARNESS_VERSION): """Update the host with a host update state change event. Args: event: HostEvent object. target_version: The test harness version which the host updates to. """ entities_to_update = [] host_update_state_enum = api_messages.HostUpdateState(event.host_update_state) host_update_state = datastore_entities.HostUpdateState.get_by_id( event.hostname) if not host_update_state: host_update_state = datastore_entities.HostUpdateState( id=event.hostname, hostname=event.hostname) if (host_update_state.update_timestamp and event.timestamp and host_update_state.update_timestamp > event.timestamp): logging.info("Ignore outdated event.") else: host_update_state.populate( state=host_update_state_enum, update_timestamp=event.timestamp, update_task_id=event.host_update_task_id, display_message=event.host_update_state_display_message, target_version=target_version) entities_to_update.append(host_update_state) host_update_state_history = datastore_entities.HostUpdateStateHistory( parent=ndb.Key(datastore_entities.HostInfo, event.hostname), hostname=event.hostname, state=host_update_state_enum, update_timestamp=event.timestamp, update_task_id=event.host_update_task_id, display_message=event.host_update_state_display_message, target_version=target_version) entities_to_update.append(host_update_state_history) ndb.put_multi(entities_to_update)
def _DoUpdateDevicesInNDB(reported_devices, event): """Update device entities to ndb. Args: reported_devices: device serial to device data mapping. event: the event have hostname, cluster info and timestamp. """ entities_to_update = [] device_keys = [] for device_serial in reported_devices.keys(): device_key = ndb.Key( datastore_entities.HostInfo, event.hostname, datastore_entities.DeviceInfo, device_serial) device_keys.append(device_key) # If the device doesn't exist, the corresponding entry will be None. devices = ndb.get_multi(device_keys) for device, device_key in zip(devices, device_keys): entities_to_update.extend( _UpdateDeviceInNDB( device, device_key, reported_devices.get(device_key.id()), event)) ndb.put_multi(entities_to_update)
def HideHost(hostname): """Hide a host and its devices.""" logging.info("Hide host %s.", hostname) host = GetHost(hostname) if not host: return None if host.hidden: logging.info("Host %s is already hidden.", hostname) return host now = common.Now() entities_to_update = [] host.hidden = True host.timestamp = now entities_to_update.append(host) devices = GetDevicesOnHost(hostname) for device in devices or []: if device.hidden: continue logging.debug("Hide device %s.", device.device_serial) device.hidden = True device.timestamp = now entities_to_update.append(device) ndb.put_multi(entities_to_update) return host
def _DoUpdateGoneDevicesInNDB(missing_device_keys, timestamp): """Do update gone devices in NDB within transactional.""" entities_to_update = [] devices = ndb.get_multi(missing_device_keys) for device in devices: if device.timestamp and device.timestamp > timestamp: logging.debug("Ignore outdated event.") continue if (device.state == common.DeviceState.GONE and device.timestamp and device.timestamp >= timestamp - ONE_MONTH): logging.debug("Ignore gone device.") continue if device.timestamp and device.timestamp < timestamp - ONE_MONTH: device.hidden = True device.timestamp = timestamp device_state_history, device_history = _UpdateDeviceState( device, common.DeviceState.GONE, timestamp) entities_to_update.append(device) if device_state_history: entities_to_update.append(device_state_history) if device_history: entities_to_update.append(device_history) ndb.put_multi(entities_to_update)
def testUpdatePredefinedMessage_failAlreadyExist(self): predefined_messages = [ datastore_entities.PredefinedMessage( lab_name='lab1', type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, content='content-1'), datastore_entities.PredefinedMessage( lab_name='lab1', type=api_messages.PredefinedMessageType.DEVICE_OFFLINE_REASON, content='content-2'), ] keys = ndb.put_multi(predefined_messages) api_request = { 'id': keys[0].id(), # the id of the 1st message 'content': 'content-2', # the content of the 2nd message } api_response = self.testapp.post_json( '/_ah/api/PredefinedMessageApi.UpdatePredefinedMessage', api_request, expect_errors=True) self.assertEqual('409 Conflict', api_response.status)
def AddOrUpdateNote(self, request): """Add or update a host note. Args: request: an API request. Returns: an api_messages.Note. """ time_now = datetime.datetime.utcnow() host_note_entity = datastore_util.GetOrCreateEntity( datastore_entities.Note, entity_id=request.id, hostname=request.hostname, type=common.NoteType.HOST_NOTE) host_note_entity.populate( user=request.user, message=request.message, timestamp=time_now, event_time=request.event_time) entities_to_update = [host_note_entity] try: offline_reason_entity = note_manager.PreparePredefinedMessageForNote( common.PredefinedMessageType.HOST_OFFLINE_REASON, message_id=request.offline_reason_id, lab_name=request.lab_name, content=request.offline_reason) except note_manager.InvalidParameterError as err: raise endpoints.BadRequestException("Invalid offline reason: [%s]" % err) if offline_reason_entity: host_note_entity.offline_reason = offline_reason_entity.content entities_to_update.append(offline_reason_entity) try: recovery_action_entity = note_manager.PreparePredefinedMessageForNote( common.PredefinedMessageType.HOST_RECOVERY_ACTION, message_id=request.recovery_action_id, lab_name=request.lab_name, content=request.recovery_action) except note_manager.InvalidParameterError as err: raise endpoints.BadRequestException("Invalid recovery action: [%s]" % err) if recovery_action_entity: host_note_entity.recovery_action = recovery_action_entity.content entities_to_update.append(recovery_action_entity) keys = ndb.put_multi(entities_to_update) host_note_msg = datastore_entities.ToMessage(host_note_entity) host_note_event_msg = api_messages.NoteEvent( note=host_note_msg, lab_name=request.lab_name) note_manager.PublishMessage(host_note_event_msg, common.PublishEventType.HOST_NOTE_EVENT) note_key = keys[0] if request.id != note_key.id(): # If ids are different, then a new note is created, we should create # a history snapshot. device_manager.CreateAndSaveHostInfoHistoryFromHostNote( request.hostname, note_key.id()) return host_note_msg
def testListPredefinedMessages_countAndCursor(self): """Test list PredefinedMessages.""" pred_msg_entities = [ datastore_entities.PredefinedMessage( lab_name='lab-name-2', type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content='content-1', used_count=4), datastore_entities.PredefinedMessage( lab_name='lab-name-2', type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content='content-2', used_count=3), datastore_entities.PredefinedMessage( lab_name='lab-name-2', type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content='content-3', used_count=2), datastore_entities.PredefinedMessage( lab_name='lab-name-2', type=api_messages.PredefinedMessageType.DEVICE_RECOVERY_ACTION, content='content-4', used_count=1), ] ndb.put_multi(pred_msg_entities) # look up the first page api_request = { 'type': 'DEVICE_RECOVERY_ACTION', 'lab_name': 'lab-name-2', 'count': 2, } api_response = self.testapp.post_json( '/_ah/api/PredefinedMessageApi.ListPredefinedMessages', api_request) pred_msg_collection = protojson.decode_message( api_messages.PredefinedMessageCollection, api_response.body) pred_msgs = pred_msg_collection.predefined_messages self.assertEqual(2, len(pred_msgs)) self.assertEqual(pred_msg_entities[0].content, pred_msgs[0].content) self.assertEqual(pred_msg_entities[1].content, pred_msgs[1].content) # look up the second page with next_cursor api_request = { 'type': 'DEVICE_RECOVERY_ACTION', 'lab_name': 'lab-name-2', 'count': 2, 'cursor': pred_msg_collection.next_cursor, } api_response = self.testapp.post_json( '/_ah/api/PredefinedMessageApi.ListPredefinedMessages', api_request) pred_msg_collection = protojson.decode_message( api_messages.PredefinedMessageCollection, api_response.body) pred_msgs = pred_msg_collection.predefined_messages self.assertEqual(2, len(pred_msgs)) self.assertEqual(pred_msg_entities[2].content, pred_msgs[0].content) self.assertEqual(pred_msg_entities[3].content, pred_msgs[1].content) # look up the first page again with prev_cursor of the second page api_request = { 'type': 'DEVICE_RECOVERY_ACTION', 'lab_name': 'lab-name-2', 'count': 2, 'cursor': pred_msg_collection.prev_cursor, 'backwards': True, } api_response = self.testapp.post_json( '/_ah/api/PredefinedMessageApi.ListPredefinedMessages', api_request) pred_msg_collection = protojson.decode_message( api_messages.PredefinedMessageCollection, api_response.body) pred_msgs = pred_msg_collection.predefined_messages self.assertEqual(2, len(pred_msgs)) self.assertEqual(pred_msg_entities[0].content, pred_msgs[0].content) self.assertEqual(pred_msg_entities[1].content, pred_msgs[1].content)
def _UpdateClusters(hosts): """Update cluster NDB entities based on hosts. Args: hosts: list of HostInfo entity with required field, from the entire system. Returns: list of ClusterInfo, clusters to upsert. """ logging.info('Updating clusters') cluster_to_hosts = collections.defaultdict(list) for host in hosts: cluster_to_hosts[host.physical_cluster].append(host) clusters_to_delete = [] clusters_to_upsert = [] query = datastore_entities.ClusterInfo.query() for cluster in query: if cluster.cluster not in cluster_to_hosts: clusters_to_delete.append(cluster.key) ndb.delete_multi(clusters_to_delete) logging.debug('Deleted clusters due to no hosts: %s', clusters_to_delete) query = datastore_entities.HostUpdateState.query() update_states_by_hostname = { update_state.hostname: update_state for update_state in query.fetch()} for cluster, hosts in six.iteritems(cluster_to_hosts): cluster_id = cluster or common.UNKNOWN_CLUSTER_NAME cluster_entity = datastore_entities.ClusterInfo(id=cluster_id) for host in hosts: if host.lab_name: cluster_entity.lab_name = host.lab_name break else: cluster_entity.lab_name = common.UNKNOWN_LAB_NAME cluster_entity.cluster = cluster_id cluster_entity.total_devices = 0 cluster_entity.offline_devices = 0 cluster_entity.available_devices = 0 cluster_entity.allocated_devices = 0 cluster_entity.device_count_timestamp = _Now() host_update_states = [] host_update_states_by_target_version = collections.defaultdict(list) host_count_by_harness_version = collections.Counter() for host in hosts: cluster_entity.total_devices += host.total_devices or 0 cluster_entity.offline_devices += host.offline_devices or 0 cluster_entity.available_devices += host.available_devices or 0 cluster_entity.allocated_devices += host.allocated_devices or 0 host_update_state = update_states_by_hostname.get(host.hostname) if host_update_state: host_update_states.append(host_update_state) host_update_states_by_target_version[ host_update_state.target_version].append(host_update_state) if host.test_harness_version: host_count_by_harness_version[host.test_harness_version] += 1 else: host_count_by_harness_version[common.UNKNOWN_TEST_HARNESS_VERSION] += 1 cluster_entity.host_count_by_harness_version = host_count_by_harness_version cluster_entity.host_update_state_summary = _CreateHostUpdateStateSummary( host_update_states) for version, states in host_update_states_by_target_version.items(): cluster_entity.host_update_state_summaries_by_version.append( _CreateHostUpdateStateSummary(states, target_version=version)) clusters_to_upsert.append(cluster_entity) ndb.put_multi(clusters_to_upsert) logging.debug('Updated clusters.') return clusters_to_upsert
def _MarkHostUpdateStateIfTimedOut(hostname): """Mark HostUpdateState as TIMED_OUT if it times out. Args: hostname: text, the host to check the update timeouts. Returns: An instance of HostUpdateState entity, None if it does not exist previously. """ host_update_state = datastore_entities.HostUpdateState.get_by_id(hostname) if not host_update_state: logging.info('No update state is found for host: %s.', hostname) return timeout_timedelta = _GetCustomizedOrDefaultHostUpdateTimeout(hostname) now = _Now() entities_to_update = [] if (host_update_state.state and host_update_state.state in common.NON_FINAL_HOST_UPDATE_STATES): if host_update_state.update_timestamp: update_state_age = now - host_update_state.update_timestamp if timeout_timedelta < update_state_age: display_message = ( _TIMEDOUT_DISPLAY_MESSAGE_TMPL % ( hostname, host_update_state.state, host_update_state.update_timestamp, update_state_age.total_seconds(), timeout_timedelta.total_seconds())) logging.info('Marking update state as TIMED_OUT: %s', display_message) host_update_state.state = api_messages.HostUpdateState.TIMED_OUT host_update_state.update_timestamp = now host_update_state.populate( state=api_messages.HostUpdateState.TIMED_OUT, update_timestamp=now, display_message=display_message) entities_to_update.append(host_update_state) host_update_state_history = datastore_entities.HostUpdateStateHistory( parent=ndb.Key(datastore_entities.HostInfo, hostname), hostname=host_update_state.hostname, state=host_update_state.state, update_timestamp=now, update_task_id=host_update_state.update_task_id, display_message=display_message) entities_to_update.append(host_update_state_history) else: logging.debug('Host<%s> is in HostUpdateState<%s> since %s.', hostname, host_update_state.state, host_update_state.update_timestamp) else: logging.debug('Host<%s> has no timestamp in the HostUpdateState. ' 'Auto adding a timestamp on it.', hostname) host_update_state.update_timestamp = now entities_to_update.append(host_update_state) host_update_state_history = datastore_entities.HostUpdateStateHistory( parent=ndb.Key(datastore_entities.HostInfo, hostname), hostname=host_update_state.hostname, state=host_update_state.state, update_timestamp=now, update_task_id=host_update_state.update_task_id) entities_to_update.append(host_update_state_history) ndb.put_multi(entities_to_update) return host_update_state
def BatchUpdateNotesWithPredefinedMessage(self, request): """Batch update notes with the same predefined message. Args: request: an API request. Returns: an api_messages.NoteCollection object. """ time_now = datetime.datetime.utcnow() host_note_entities = [] for note in request.notes: note_id = int(note.id) if note.id is not None else None host_note_entity = datastore_util.GetOrCreateEntity( datastore_entities.Note, entity_id=note_id, hostname=note.hostname, type=common.NoteType.HOST_NOTE) host_note_entity.populate( user=request.user, message=request.message, timestamp=time_now, event_time=request.event_time) host_note_entities.append(host_note_entity) try: offline_reason_entity = note_manager.PreparePredefinedMessageForNote( common.PredefinedMessageType.HOST_OFFLINE_REASON, message_id=request.offline_reason_id, lab_name=request.lab_name, content=request.offline_reason, delta_count=len(host_note_entities)) except note_manager.InvalidParameterError as err: raise endpoints.BadRequestException("Invalid offline reason: [%s]" % err) if offline_reason_entity: for host_note_entity in host_note_entities: host_note_entity.offline_reason = offline_reason_entity.content offline_reason_entity.put() try: recovery_action_entity = note_manager.PreparePredefinedMessageForNote( common.PredefinedMessageType.HOST_RECOVERY_ACTION, message_id=request.recovery_action_id, lab_name=request.lab_name, content=request.recovery_action, delta_count=len(host_note_entities)) except note_manager.InvalidParameterError as err: raise endpoints.BadRequestException("Invalid recovery action: [%s]" % err) if recovery_action_entity: for host_note_entity in host_note_entities: host_note_entity.recovery_action = recovery_action_entity.content recovery_action_entity.put() note_keys = ndb.put_multi(host_note_entities) host_note_entities = ndb.get_multi(note_keys) note_msgs = [] for host_note_entity in host_note_entities: host_note_msg = datastore_entities.ToMessage(host_note_entity) note_msgs.append(host_note_msg) host_note_event_msg = api_messages.NoteEvent( note=host_note_msg, lab_name=request.lab_name) note_manager.PublishMessage( host_note_event_msg, common.PublishEventType.HOST_NOTE_EVENT) for request_note, updated_note_key in zip(request.notes, note_keys): if not request_note.id: # If ids are not provided, then a new note is created, we should create # a history snapshot. device_manager.CreateAndSaveHostInfoHistoryFromHostNote( request_note.hostname, updated_note_key.id()) return api_messages.NoteCollection( notes=note_msgs, more=False, next_cursor=None, prev_cursor=None)
def BatchUpdateHostMetadata(self, request): """Update HostMetadata on multiple hosts. Args: request: an API request. Request Params: hostname: list of strings, the name of hosts. test_harness_image: string, the url to test harness image. user: string, the user sending the request. Returns: a message_types.VoidMessage object. Raises: endpoints.BadRequestException, when request does not match existing hosts. """ host_configs = ndb.get_multi( ndb.Key(datastore_entities.HostConfig, hostname) for hostname in request.hostnames) host_metadatas = ndb.get_multi( ndb.Key(datastore_entities.HostMetadata, hostname) for hostname in request.hostnames) hosts_no_permission = [] hosts_not_enabled = [] metadatas_to_update = [] for hostname, config, metadata in zip( request.hostnames, host_configs, host_metadatas): if not config or not config.enable_ui_update: hosts_not_enabled.append(hostname) continue if request.user not in config.owners: hosts_no_permission.append(hostname) continue if not metadata: metadata = datastore_entities.HostMetadata( id=hostname, hostname=hostname) if not harness_image_metadata_syncer.AreHarnessImagesEqual( metadata.test_harness_image, request.test_harness_image): event = host_event.HostEvent( time=datetime.datetime.utcnow(), type=_HOST_UPDATE_STATE_CHANGED_EVENT_NAME, hostname=hostname, host_update_state=_HOST_UPDATE_STATE_PENDING, data={"host_update_target_image": request.test_harness_image}) device_manager.HandleDeviceSnapshotWithNDB(event) metadata.populate(test_harness_image=request.test_harness_image) metadatas_to_update.append(metadata) ndb.put_multi(metadatas_to_update) if not hosts_no_permission and not hosts_not_enabled: return message_types.VoidMessage() error_message = "" if hosts_no_permission: error_message += ( "Request user %s is not in the owner list of hosts [%s]. " % (request.user, ", ".join(hosts_no_permission))) if hosts_not_enabled: error_message += ("Hosts [%s] are not enabled to be updated from UI. " % ", ".join(hosts_not_enabled)) raise endpoints.BadRequestException(error_message)
def testSyncHarnessImageMetadata_OverwriteExistingEntities( self, mock_requests, mock_util_datetime, mock_syncer_datetime, mock_auth): """Test sync harness image metadata.""" time_now = datetime.datetime(2020, 12, 24) time_created = datetime.datetime(2020, 12, 10) time_created_ms = str( int((time_created - datetime.datetime(1970, 1, 1)).total_seconds() * 1000)) mock_util_datetime.datetime.utcnow.return_value = time_now mock_syncer_datetime.datetime.utcnow.return_value = time_now mock_syncer_datetime.datetime.utcfromtimestamp = ( datetime.datetime.utcfromtimestamp) mock_requests.get().json.return_value = { 'manifest': { 'sha1': { 'tag': [ '111111', 'golden', 'canary', 'golden_tradefed_image_20201210_1200_RC00', ], 'timeCreatedMs': time_created_ms, }, 'sha2': { 'tag': [ '2222222', 'golden_tradefed_image_20201210_0600_RC00', ], 'timeCreatedMs': time_created_ms, }, 'sha3': { 'tag': [ '3333333', 'staging', ], 'timeCreatedMs': time_created_ms, }, } } existing_entities = [ datastore_entities.TestHarnessImageMetadata( key=ndb.Key(datastore_entities.TestHarnessImageMetadata, 'gcr.io/dockerized-tradefed/tradefed:sha1'), repo_name='gcr.io/dockerized-tradefed/tradefed', digest='sha1', test_harness=harness_image_metadata_syncer. _TRADEFED_HARNESS_NAME, test_harness_version='111111', current_tags=['111111', 'canary', 'staging'], create_time=time_created, sync_time=time_now), datastore_entities.TestHarnessImageMetadata( key=ndb.Key(datastore_entities.TestHarnessImageMetadata, 'gcr.io/dockerized-tradefed/tradefed:sha2'), repo_name='gcr.io/dockerized-tradefed/tradefed', digest='sha2', test_harness=harness_image_metadata_syncer. _TRADEFED_HARNESS_NAME, test_harness_version='2222222', current_tags=[ '2222222', 'golden', 'golden_tradefed_image_20201210_0600_RC00' ], historical_tags=['golden'], create_time=time_created, sync_time=time_now), ] ndb.put_multi(existing_entities) harness_image_metadata_syncer.SyncHarnessImageMetadata() keys = [ ndb.Key(datastore_entities.TestHarnessImageMetadata, 'gcr.io/dockerized-tradefed/tradefed:sha1'), ndb.Key(datastore_entities.TestHarnessImageMetadata, 'gcr.io/dockerized-tradefed/tradefed:sha2'), ndb.Key(datastore_entities.TestHarnessImageMetadata, 'gcr.io/dockerized-tradefed/tradefed:sha3'), ] entity_1, entity_2, entity_3 = ndb.get_multi(keys) self.assertEqual('sha1', entity_1.digest) self.assertEqual('111111', entity_1.test_harness_version) self.assertEqual(time_created, entity_1.create_time) self.assertEqual(time_now, entity_1.sync_time) self.assertCountEqual([ '111111', 'golden', 'canary', 'golden_tradefed_image_20201210_1200_RC00' ], entity_1.current_tags) self.assertCountEqual(['golden'], entity_1.historical_tags) self.assertEqual('sha2', entity_2.digest) self.assertEqual('2222222', entity_2.test_harness_version) self.assertEqual(time_created, entity_2.create_time) self.assertEqual(time_now, entity_2.sync_time) self.assertCountEqual([ '2222222', 'golden_tradefed_image_20201210_0600_RC00', ], entity_2.current_tags) self.assertCountEqual(['golden'], entity_2.historical_tags) self.assertEqual('sha3', entity_3.digest) self.assertEqual('3333333', entity_3.test_harness_version) self.assertEqual(time_created, entity_3.create_time) self.assertEqual(time_now, entity_3.sync_time) self.assertCountEqual(['3333333', 'staging'], entity_3.current_tags) self.assertEmpty(entity_3.historical_tags)