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
120_SERVER_TABLE_FORMAT = ('%(hostname)-30s | %(status)-7s | %(roles)-20s |'
121                        ' %(date_created)-19s | %(date_modified)-19s |'
122                        ' %(note)s')
123
124
125def format_servers_table(servers):
126    """format servers for printing as a table.
127
128    Example output:
129
130        Hostname | Status  | Roles     | Date Created    | Date Modified | Note
131        server1  | backup  | scheduler | 2014-11-25 23:45:19 |           |
132        server2  | primary | drone     | 2014-11-25 12:00:00 |           | Drone
133
134    @param servers: Sequence of Server instances.
135    @returns: Formatted output as string.
136    """
137    result_lines = [(_SERVER_TABLE_FORMAT %
138                     {'hostname': 'Hostname',
139                      'status': 'Status',
140                      'roles': 'Roles',
141                      'date_created': 'Date Created',
142                      'date_modified': 'Date Modified',
143                      'note': 'Note'})]
144    for server in servers:
145        roles = ','.join(server.get_role_names())
146        result_lines.append(_SERVER_TABLE_FORMAT %
147                            {'hostname':server.hostname,
148                             'status': server.status or '',
149                             'roles': roles,
150                             'date_created': server.date_created,
151                             'date_modified': server.date_modified or '',
152                             'note': server.note or ''})
153    return '\n'.join(result_lines)
154
155
156def format_servers_summary(servers):
157    """format servers for printing a summary.
158
159    Example output:
160
161        scheduler      : server1(backup), server3(primary),
162        host_scheduler :
163        drone          : server2(primary),
164        devserver      :
165        database       :
166        crash_server   :
167        No Role        :
168
169    @param servers: Sequence of Server instances.
170    @returns: Formatted output as string.
171    """
172    servers_by_role = _get_servers_by_role(servers)
173    servers_with_roles = {server for role_servers in servers_by_role.itervalues()
174                          for server in role_servers}
175    servers_without_roles = [server for server in servers
176                             if server not in servers_with_roles]
177    result_lines = ['Roles and status of servers:', '']
178    for role, role_servers in servers_by_role.iteritems():
179        result_lines.append(_format_role_servers_summary(role, role_servers))
180    if servers_without_roles:
181        result_lines.append(
182                _format_role_servers_summary('No Role', servers_without_roles))
183    return '\n'.join(result_lines)
184
185
186def format_servers_nameonly(servers):
187    """format servers for printing names only
188
189    @param servers: Sequence of Server instances.
190    @returns: Formatted output as string.
191    """
192    return '\n'.join(s.hostname for s in servers)
193
194
195def _get_servers_by_role(servers):
196    """Return a mapping from roles to servers.
197
198    @param servers: Iterable of servers.
199    @returns: Mapping of role strings to lists of servers.
200    """
201    roles = [role for role, _ in server_models.ServerRole.ROLE.choices()]
202    servers_by_role = collections.defaultdict(list)
203    for server in servers:
204        for role in server.get_role_names():
205            servers_by_role[role].append(server)
206    return servers_by_role
207
208
209def _format_role_servers_summary(role, servers):
210    """Format one line of servers for a role in a server list summary.
211
212    @param role: Role string.
213    @param servers: Iterable of Server instances.
214    @returns: String.
215    """
216    servers_part = ', '.join(
217            '%s(%s)' % (server.hostname, server.status)
218            for server in servers)
219    return '%-15s: %s' % (role, servers_part)
220
221
222def check_server(hostname, role):
223    """Confirm server with given hostname is ready to be primary of given role.
224
225    If the server is a backup and failed to be verified for the role, remove
226    the role from its roles list. If it has no other role, set its status to
227    repair_required.
228
229    @param hostname: hostname of the server.
230    @param role: Role to be checked.
231    @return: True if server can be verified for the given role, otherwise
232             return False.
233    """
234    # TODO(dshi): Add more logic to confirm server is ready for the role.
235    # For now, the function just checks if server is ssh-able.
236    try:
237        infra.execute_command(hostname, 'true')
238        return True
239    except subprocess.CalledProcessError as e:
240        print >> sys.stderr, ('Failed to check server %s, error: %s' %
241                              (hostname, e))
242        return False
243
244
245def verify_server(exist=True):
246    """Decorator to check if server with given hostname exists in the database.
247
248    @param exist: Set to True to confirm server exists in the database, raise
249                  exception if not. If it's set to False, raise exception if
250                  server exists in database. Default is True.
251
252    @raise ServerActionError: If `exist` is True and server does not exist in
253                              the database, or `exist` is False and server exists
254                              in the database.
255    """
256    def deco_verify(func):
257        """Wrapper for the decorator.
258
259        @param func: Function to be called.
260        """
261        def func_verify(*args, **kwargs):
262            """Decorator to check if server exists.
263
264            If exist is set to True, raise ServerActionError is server with
265            given hostname is not found in server database.
266            If exist is set to False, raise ServerActionError is server with
267            given hostname is found in server database.
268
269            @param func: function to be called.
270            @param args: arguments for function to be called.
271            @param kwargs: keyword arguments for function to be called.
272            """
273            hostname = kwargs['hostname']
274            try:
275                server = server_models.Server.objects.get(hostname=hostname)
276            except django.core.exceptions.ObjectDoesNotExist:
277                server = None
278
279            if not exist and server:
280                raise ServerActionError('Server %s already exists.' %
281                                        hostname)
282            if exist and not server:
283                raise ServerActionError('Server %s does not exist in the '
284                                        'database.' % hostname)
285            if server:
286                kwargs['server'] = server
287            return func(*args, **kwargs)
288        return func_verify
289    return deco_verify
290
291
292def get_drones():
293    """Get a list of drones in status primary.
294
295    @return: A list of drones in status primary.
296    """
297    servers = get_servers(role=server_models.ServerRole.ROLE.DRONE,
298                          status=server_models.Server.STATUS.PRIMARY)
299    return [s.hostname for s in servers]
300
301
302def delete_attribute(server, attribute):
303    """Delete the attribute from the host.
304
305    @param server: An object of server_models.Server.
306    @param attribute: Name of an attribute of the server.
307    """
308    attributes = server.attributes.filter(attribute=attribute)
309    if not attributes:
310        raise ServerActionError('Server %s does not have attribute %s' %
311                                (server.hostname, attribute))
312    attributes[0].delete()
313    print 'Attribute %s is deleted from server %s.' % (attribute,
314                                                       server.hostname)
315
316
317def change_attribute(server, attribute, value):
318    """Change the value of an attribute of the server.
319
320    @param server: An object of server_models.Server.
321    @param attribute: Name of an attribute of the server.
322    @param value: Value of the attribute of the server.
323
324    @raise ServerActionError: If the attribute already exists and has the
325                              given value.
326    """
327    attributes = server_models.ServerAttribute.objects.filter(
328            server=server, attribute=attribute)
329    if attributes and attributes[0].value == value:
330        raise ServerActionError('Attribute %s for Server %s already has '
331                                'value of %s.' %
332                                (attribute, server.hostname, value))
333    if attributes:
334        old_value = attributes[0].value
335        attributes[0].value = value
336        attributes[0].save()
337        print ('Attribute `%s` of server %s is changed from %s to %s.' %
338                     (attribute, server.hostname, old_value, value))
339    else:
340        server_models.ServerAttribute.objects.create(
341                server=server, attribute=attribute, value=value)
342        print ('Attribute `%s` of server %s is set to %s.' %
343               (attribute, server.hostname, value))
344
345
346def get_shards():
347    """Get a list of shards in status primary.
348
349    @return: A list of shards in status primary.
350    """
351    servers = get_servers(role=server_models.ServerRole.ROLE.SHARD,
352                          status=server_models.Server.STATUS.PRIMARY)
353    return [s.hostname for s in servers]
354
355
356def confirm_server_has_role(hostname, role):
357    """Confirm a given server has the given role, and its status is primary.
358
359    @param hostname: hostname of the server.
360    @param role: Name of the role to be checked.
361    @raise ServerActionError: If localhost does not have given role or it's
362                              not in primary status.
363    """
364    if hostname.lower() in ['localhost', '127.0.0.1']:
365        hostname = socket.gethostname()
366    hostname = utils.normalize_hostname(hostname)
367
368    servers = get_servers(role=role, status=server_models.Server.STATUS.PRIMARY)
369    for server in servers:
370        if hostname == utils.normalize_hostname(server.hostname):
371            return True
372    raise ServerActionError('Server %s does not have role of %s running in '
373                            'status primary.' % (hostname, role))
374