1# Copyright 2014 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5"""This module provides utility functions to help managing servers in server 6database (defined in global config section AUTOTEST_SERVER_DB). 7 8""" 9 10import collections 11import json 12import socket 13import subprocess 14import sys 15 16import common 17 18import django.core.exceptions 19from autotest_lib.client.common_lib import utils 20from autotest_lib.client.common_lib.global_config import global_config 21from autotest_lib.frontend.server import models as server_models 22from autotest_lib.site_utils.lib import infra 23 24 25class ServerActionError(Exception): 26 """Exception raised when action on server failed. 27 """ 28 29 30def use_server_db(): 31 """Check if use_server_db is enabled in configuration. 32 33 @return: True if use_server_db is set to True in global config. 34 """ 35 return global_config.get_config_value( 36 'SERVER', 'use_server_db', default=False, type=bool) 37 38 39def warn_missing_role(role, exclude_server): 40 """Post a warning if Autotest instance has no other primary server with 41 given role. 42 43 @param role: Name of the role. 44 @param exclude_server: Server to be excluded from search for role. 45 """ 46 servers = server_models.Server.objects.filter( 47 roles__role=role, 48 status=server_models.Server.STATUS.PRIMARY).exclude( 49 hostname=exclude_server.hostname) 50 if not servers: 51 message = ('WARNING! There will be no server with role %s after it\'s ' 52 'removed from server %s. Autotest will not function ' 53 'normally without any server in role %s.' % 54 (role, exclude_server.hostname, role)) 55 print >> sys.stderr, message 56 57 58def get_servers(hostname=None, role=None, status=None): 59 """Find servers with given role and status. 60 61 @param hostname: hostname of the server. 62 @param role: Role of server, default to None. 63 @param status: Status of server, default to None. 64 65 @return: A list of server objects with given role and status. 66 """ 67 filters = {} 68 if hostname: 69 filters['hostname'] = hostname 70 if role: 71 filters['roles__role'] = role 72 if status: 73 filters['status'] = status 74 return list(server_models.Server.objects.filter(**filters)) 75 76 77def format_servers(servers): 78 """Format servers for printing. 79 80 Example output: 81 82 Hostname : server2 83 Status : primary 84 Roles : drone 85 Attributes : {'max_processes':300} 86 Date Created : 2014-11-25 12:00:00 87 Date Modified: None 88 Note : Drone in lab1 89 90 @param servers: Sequence of Server instances. 91 @returns: Formatted output as string. 92 """ 93 return '\n'.join(str(server) for server in servers) 94 95 96def format_servers_json(servers): 97 """Format servers for printing as JSON. 98 99 @param servers: Sequence of Server instances. 100 @returns: String. 101 """ 102 server_dicts = [] 103 for server in servers: 104 if server.date_modified is None: 105 date_modified = None 106 else: 107 date_modified = str(server.date_modified) 108 attributes = {k: v for k, v in server.attributes.values_list( 109 'attribute', 'value')} 110 server_dicts.append({'hostname': server.hostname, 111 'status': server.status, 112 'roles': server.get_role_names(), 113 'date_created': str(server.date_created), 114 'date_modified': date_modified, 115 'note': server.note, 116 'attributes': attributes}) 117 return json.dumps(server_dicts) 118 119 120def format_servers_nameonly(servers): 121 """format servers for printing names only 122 123 @param servers: Sequence of Server instances. 124 @returns: Formatted output as string. 125 """ 126 return '\n'.join(s.hostname for s in servers) 127 128 129def _get_servers_by_role(servers): 130 """Return a mapping from roles to servers. 131 132 @param servers: Iterable of servers. 133 @returns: Mapping of role strings to lists of servers. 134 """ 135 roles = [role for role, _ in server_models.ServerRole.ROLE.choices()] 136 servers_by_role = collections.defaultdict(list) 137 for server in servers: 138 for role in server.get_role_names(): 139 servers_by_role[role].append(server) 140 return servers_by_role 141 142 143def _format_role_servers_summary(role, servers): 144 """Format one line of servers for a role in a server list summary. 145 146 @param role: Role string. 147 @param servers: Iterable of Server instances. 148 @returns: String. 149 """ 150 servers_part = ', '.join( 151 '%s(%s)' % (server.hostname, server.status) 152 for server in servers) 153 return '%-15s: %s' % (role, servers_part) 154 155 156def check_server(hostname, role): 157 """Confirm server with given hostname is ready to be primary of given role. 158 159 If the server is a backup and failed to be verified for the role, remove 160 the role from its roles list. If it has no other role, set its status to 161 repair_required. 162 163 @param hostname: hostname of the server. 164 @param role: Role to be checked. 165 @return: True if server can be verified for the given role, otherwise 166 return False. 167 """ 168 # TODO(dshi): Add more logic to confirm server is ready for the role. 169 # For now, the function just checks if server is ssh-able. 170 try: 171 infra.execute_command(hostname, 'true') 172 return True 173 except subprocess.CalledProcessError as e: 174 print >> sys.stderr, ('Failed to check server %s, error: %s' % 175 (hostname, e)) 176 return False 177 178 179def verify_server(exist=True): 180 """Decorator to check if server with given hostname exists in the database. 181 182 @param exist: Set to True to confirm server exists in the database, raise 183 exception if not. If it's set to False, raise exception if 184 server exists in database. Default is True. 185 186 @raise ServerActionError: If `exist` is True and server does not exist in 187 the database, or `exist` is False and server exists 188 in the database. 189 """ 190 def deco_verify(func): 191 """Wrapper for the decorator. 192 193 @param func: Function to be called. 194 """ 195 def func_verify(*args, **kwargs): 196 """Decorator to check if server exists. 197 198 If exist is set to True, raise ServerActionError is server with 199 given hostname is not found in server database. 200 If exist is set to False, raise ServerActionError is server with 201 given hostname is found in server database. 202 203 @param func: function to be called. 204 @param args: arguments for function to be called. 205 @param kwargs: keyword arguments for function to be called. 206 """ 207 hostname = kwargs['hostname'] 208 try: 209 server = server_models.Server.objects.get(hostname=hostname) 210 except django.core.exceptions.ObjectDoesNotExist: 211 server = None 212 213 if not exist and server: 214 raise ServerActionError('Server %s already exists.' % 215 hostname) 216 if exist and not server: 217 raise ServerActionError('Server %s does not exist in the ' 218 'database.' % hostname) 219 if server: 220 kwargs['server'] = server 221 return func(*args, **kwargs) 222 return func_verify 223 return deco_verify 224 225 226def get_drones(): 227 """Get a list of drones in status primary. 228 229 @return: A list of drones in status primary. 230 """ 231 servers = get_servers(role=server_models.ServerRole.ROLE.DRONE, 232 status=server_models.Server.STATUS.PRIMARY) 233 return [s.hostname for s in servers] 234 235 236def delete_attribute(server, attribute): 237 """Delete the attribute from the host. 238 239 @param server: An object of server_models.Server. 240 @param attribute: Name of an attribute of the server. 241 """ 242 attributes = server.attributes.filter(attribute=attribute) 243 if not attributes: 244 raise ServerActionError('Server %s does not have attribute %s' % 245 (server.hostname, attribute)) 246 attributes[0].delete() 247 print 'Attribute %s is deleted from server %s.' % (attribute, 248 server.hostname) 249 250 251def change_attribute(server, attribute, value): 252 """Change the value of an attribute of the server. 253 254 @param server: An object of server_models.Server. 255 @param attribute: Name of an attribute of the server. 256 @param value: Value of the attribute of the server. 257 258 @raise ServerActionError: If the attribute already exists and has the 259 given value. 260 """ 261 attributes = server_models.ServerAttribute.objects.filter( 262 server=server, attribute=attribute) 263 if attributes and attributes[0].value == value: 264 raise ServerActionError('Attribute %s for Server %s already has ' 265 'value of %s.' % 266 (attribute, server.hostname, value)) 267 if attributes: 268 old_value = attributes[0].value 269 attributes[0].value = value 270 attributes[0].save() 271 print ('Attribute `%s` of server %s is changed from %s to %s.' % 272 (attribute, server.hostname, old_value, value)) 273 else: 274 server_models.ServerAttribute.objects.create( 275 server=server, attribute=attribute, value=value) 276 print ('Attribute `%s` of server %s is set to %s.' % 277 (attribute, server.hostname, value)) 278 279 280def get_shards(): 281 """Get a list of shards in status primary. 282 283 @return: A list of shards in status primary. 284 """ 285 servers = get_servers(role=server_models.ServerRole.ROLE.SHARD, 286 status=server_models.Server.STATUS.PRIMARY) 287 return [s.hostname for s in servers] 288 289 290def confirm_server_has_role(hostname, role): 291 """Confirm a given server has the given role, and its status is primary. 292 293 @param hostname: hostname of the server. 294 @param role: Name of the role to be checked. 295 @raise ServerActionError: If localhost does not have given role or it's 296 not in primary status. 297 """ 298 if hostname.lower() in ['localhost', '127.0.0.1']: 299 hostname = socket.gethostname() 300 hostname = utils.normalize_hostname(hostname) 301 302 servers = get_servers(role=role, status=server_models.Server.STATUS.PRIMARY) 303 for server in servers: 304 if hostname == utils.normalize_hostname(server.hostname): 305 return True 306 raise ServerActionError('Server %s does not have role of %s running in ' 307 'status primary.' % (hostname, role)) 308