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