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