1# Copyright 2017 Google Inc. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Host Info APIs implemented using Google Cloud Endpoints."""
15
16import datetime
17import endpoints
18import logging
19
20from google.appengine.api import users
21from google.appengine.ext import ndb
22
23from webapp.src import vtslab_status as Status
24from webapp.src.endpoint import endpoint_base
25from webapp.src.proto import model
26
27HOST_INFO_RESOURCE = endpoints.ResourceContainer(model.HostInfoMessage)
28
29# Product type name for null device.
30_NULL_DEVICE_PRODUCT_TYPE = "null"
31
32
33def AddNullDevices(hostname, null_device_count):
34    """Adds null devices to DeviceModel data store.
35
36    Args:
37        hostname: string, the host name.
38        null_device_count: integer, the number of null devices.
39    """
40    device_query = model.DeviceModel.query(
41        model.DeviceModel.hostname == hostname,
42        model.DeviceModel.product == _NULL_DEVICE_PRODUCT_TYPE
43    )
44    null_devices = device_query.fetch()
45    existing_null_device_count = len(null_devices)
46
47    if existing_null_device_count < null_device_count:
48        devices_to_put = []
49        for _ in range(null_device_count - existing_null_device_count):
50            device = model.DeviceModel()
51            device.hostname = hostname
52            device.serial = "n/a"
53            device.product = _NULL_DEVICE_PRODUCT_TYPE
54            device.status = Status.DEVICE_STATUS_DICT["ready"]
55            device.scheduling_status = Status.DEVICE_SCHEDULING_STATUS_DICT[
56                "free"]
57            device.timestamp = datetime.datetime.now()
58            devices_to_put.append(device)
59        if devices_to_put:
60            ndb.put_multi(devices_to_put)
61
62
63@endpoints.api(name='host', version='v1')
64class HostInfoApi(endpoint_base.EndpointBase):
65    """Endpoint API for host_info."""
66
67    @endpoints.method(
68        HOST_INFO_RESOURCE,
69        model.DefaultResponse,
70        path='set',
71        http_method='POST',
72        name='set')
73    def set(self, request):
74        """Sets the host info based on the `request`."""
75        if users.get_current_user():
76            username = users.get_current_user().email()
77        else:
78            username = "anonymous"
79
80        devices_to_put = []
81        for request_device in request.devices:
82            device_query = model.DeviceModel.query(
83                model.DeviceModel.serial == request_device.serial
84            )
85            existing_device = device_query.fetch()
86            if existing_device:
87                device = existing_device[0]
88            else:
89                device = model.DeviceModel()
90                device.serial = request_device.serial
91                device.scheduling_status = Status.DEVICE_SCHEDULING_STATUS_DICT[
92                    "free"]
93            if not device.product or request_device.product != "error":
94                device.product = request_device.product
95
96            device.username = username
97            device.hostname = request.hostname
98            device.status = request_device.status
99            device.timestamp = datetime.datetime.now()
100            devices_to_put.append(device)
101        if devices_to_put:
102            ndb.put_multi(devices_to_put)
103
104        return model.DefaultResponse(
105            return_code=model.ReturnCodeMessage.SUCCESS)
106
107    @endpoints.method(
108        endpoint_base.GET_REQUEST_RESOURCE,
109        model.DeviceResponseMessage,
110        path="get",
111        http_method="POST",
112        name="get")
113    def get(self, request):
114        """Gets the devices from datastore."""
115        return_list, more = self.Get(request=request,
116                                     metaclass=model.DeviceModel,
117                                     message=model.DeviceInfoMessage)
118
119        return model.DeviceResponseMessage(devices=return_list, has_next=more)
120
121    @endpoints.method(
122        endpoint_base.COUNT_REQUEST_RESOURCE,
123        model.CountResponseMessage,
124        path="count",
125        http_method="POST",
126        name="count")
127    def count(self, request):
128        """Gets total number of DeviceModel entities stored in datastore."""
129        filters = self.CreateFilterList(
130            filter_string=request.filter, metaclass=model.DeviceModel)
131
132        count = self.Count(metaclass=model.DeviceModel, filters=filters)
133
134        return model.CountResponseMessage(count=count)
135