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"""Schedule Info APIs implemented using Google Cloud Endpoints."""
15
16import datetime
17import endpoints
18
19from google.appengine.ext import ndb
20
21from webapp.src import vtslab_status as Status
22from webapp.src.endpoint import endpoint_base
23from webapp.src.proto import model
24from webapp.src.utils import email_util
25
26SCHEDULE_INFO_RESOURCE = endpoints.ResourceContainer(model.ScheduleInfoMessage)
27SCHEDULE_SUSPEND_RESOURCE = endpoints.ResourceContainer(
28    model.ScheduleSuspendMessage)
29
30
31@endpoints.api(name="schedule", version="v1")
32class ScheduleInfoApi(endpoint_base.EndpointBase):
33    """Endpoint API for schedule_info."""
34
35    @endpoints.method(
36        SCHEDULE_INFO_RESOURCE,
37        model.DefaultResponse,
38        path="clear",
39        http_method="POST",
40        name="clear")
41    def clear(self, request):
42        """Clears test schedule info in DB."""
43        schedule_query = model.ScheduleModel.query(
44            model.ScheduleModel.schedule_type != "green")
45        existing_schedules = schedule_query.fetch(keys_only=True)
46        if existing_schedules and len(existing_schedules) > 0:
47            ndb.delete_multi(existing_schedules)
48        return model.DefaultResponse(
49            return_code=model.ReturnCodeMessage.SUCCESS)
50
51    @endpoints.method(
52        SCHEDULE_INFO_RESOURCE,
53        model.DefaultResponse,
54        path="set",
55        http_method="POST",
56        name="set")
57    def set(self, request):
58        """Sets the schedule info based on `request`."""
59        exist_on_both = self.GetCommonAttributes(request, model.ScheduleModel)
60        # check duplicates
61        exclusions = [
62            "name", "schedule_type", "schedule", "param", "timestamp",
63            "children_jobs", "error_count", "suspended"
64        ]
65        # list of protorpc message fields.
66        duplicate_checklist = [x for x in exist_on_both if x not in exclusions]
67        empty_list_field = []
68        query = model.ScheduleModel.query()
69        for attr_name in duplicate_checklist:
70            if model.ScheduleModel._properties[attr_name]._repeated:
71                value = request.get_assigned_value(attr_name)
72                if value:
73                    query = query.filter(
74                        getattr(model.ScheduleModel, attr_name).IN(
75                            request.get_assigned_value(attr_name)))
76                else:
77                    # empty list cannot be queried.
78                    empty_list_field.append(attr_name)
79            else:
80                query = query.filter(
81                    getattr(model.ScheduleModel, attr_name) ==
82                    request.get_assigned_value(attr_name))
83        duplicated_schedules = query.fetch()
84
85        if empty_list_field:
86            duplicated_schedules = [
87                schedule for schedule in duplicated_schedules
88                if all(
89                    [not getattr(schedule, attr) for attr in empty_list_field])
90            ]
91
92        if duplicated_schedules:
93            schedule = duplicated_schedules[0]
94        else:
95            schedule = model.ScheduleModel()
96            for attr_name in exist_on_both:
97                setattr(schedule, attr_name,
98                        request.get_assigned_value(attr_name))
99            schedule.schedule_type = "test"
100            schedule.error_count = 0
101            schedule.suspended = False
102            schedule.priority_value = Status.GetPriorityValue(schedule.priority)
103
104        schedule.timestamp = datetime.datetime.now()
105        schedule.put()
106
107        return model.DefaultResponse(
108            return_code=model.ReturnCodeMessage.SUCCESS)
109
110    @endpoints.method(
111        endpoint_base.GET_REQUEST_RESOURCE,
112        model.ScheduleResponseMessage,
113        path="get",
114        http_method="POST",
115        name="get")
116    def get(self, request):
117        """Gets the schedules from datastore."""
118        return_list, more = self.Get(request=request,
119                                     metaclass=model.ScheduleModel,
120                                     message=model.ScheduleInfoMessage)
121
122        return model.ScheduleResponseMessage(
123            schedules=return_list, has_next=more)
124
125    @endpoints.method(
126        endpoint_base.COUNT_REQUEST_RESOURCE,
127        model.CountResponseMessage,
128        path="count",
129        http_method="POST",
130        name="count")
131    def count(self, request):
132        """Gets total number of ScheduleModel entities stored in datastore."""
133        filters = self.CreateFilterList(
134            filter_string=request.filter, metaclass=model.ScheduleModel)
135
136        count = self.Count(metaclass=model.ScheduleModel, filters=filters)
137
138        return model.CountResponseMessage(count=count)
139
140    @endpoints.method(
141        SCHEDULE_SUSPEND_RESOURCE,
142        model.ScheduleSuspendMessage,
143        path="suspend",
144        http_method="POST",
145        name="suspend")
146    def suspend(self, request):
147        """Toggles a schedule from suspend to resume, or vice versa."""
148        schedules_to_put = []
149        schedules_to_return = []
150        for schedule in request.schedules:
151            schedule_key = ndb.key.Key(urlsafe=schedule.urlsafe_key)
152            schedule_entity = schedule_key.get()
153            if schedule.suspend:  # to suspend
154                schedule_entity.suspended = True
155            else:  # to resume
156                schedule_entity.error_count = 0
157                schedule_entity.suspended = False
158            schedules_to_put.append(schedule_entity)
159            schedules_to_return.append({"urlsafe_key": schedule.urlsafe_key,
160                                        "suspend": schedule_entity.suspended})
161            # TODO(jongmok): Minimize a number of emails by merging schedules.
162            email_util.send_schedule_suspension_notification(schedule_entity)
163
164        ndb.put_multi(schedules_to_put)
165        return model.ScheduleSuspendMessage(schedules=schedules_to_return)
166
167
168@endpoints.api(name="green_schedule_info", version="v1")
169class GreenScheduleInfoApi(endpoint_base.EndpointBase):
170    """Endpoint API for green_schedule_info."""
171
172    @endpoints.method(
173        SCHEDULE_INFO_RESOURCE,
174        model.DefaultResponse,
175        path="clear",
176        http_method="POST",
177        name="clear")
178    def clear(self, request):
179        """Clears green build schedule info in DB."""
180        schedule_query = model.ScheduleModel.query(
181            model.ScheduleModel.schedule_type == "green")
182        existing_schedules = schedule_query.fetch(keys_only=True)
183        if existing_schedules and len(existing_schedules) > 0:
184            ndb.delete_multi(existing_schedules)
185        return model.DefaultResponse(
186            return_code=model.ReturnCodeMessage.SUCCESS)
187
188    @endpoints.method(
189        SCHEDULE_INFO_RESOURCE,
190        model.DefaultResponse,
191        path="set",
192        http_method="POST",
193        name="set")
194    def set(self, request):
195        """Sets the green build schedule info based on `request`."""
196        schedule = model.ScheduleModel()
197        schedule.name = request.name
198        schedule.manifest_branch = request.manifest_branch
199        schedule.build_target = request.build_target
200        schedule.device_pab_account_id = request.device_pab_account_id
201        schedule.test_name = request.test_name
202        schedule.schedule = request.schedule
203        schedule.priority = request.priority
204        schedule.device = request.device
205        schedule.shards = request.shards
206        schedule.gsi_branch = request.gsi_branch
207        schedule.gsi_build_target = request.gsi_build_target
208        schedule.gsi_pab_account_id = request.gsi_pab_account_id
209        schedule.gsi_vendor_version = request.gsi_vendor_version
210        schedule.test_branch = request.test_branch
211        schedule.test_build_target = request.test_build_target
212        schedule.test_pab_account_id = request.test_pab_account_id
213        schedule.timestamp = datetime.datetime.now()
214        schedule.schedule_type = "green"
215        schedule.put()
216
217        return model.DefaultResponse(
218            return_code=model.ReturnCodeMessage.SUCCESS)
219