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