1#!/usr/bin/python
2
3"""A script that provides convertion between models.job and a protocol
4buffer object.
5
6This script contains only one class that takes an job instance and
7convert it into a protocol buffer object. The class will also be
8responsible for serializing the job instance via protocol buffers.
9
10"""
11
12# import python libraries
13import os
14import datetime
15import time
16import random
17import re
18
19# import autotest libraries
20from autotest_lib.tko import models
21from autotest_lib.tko import tko_pb2
22from autotest_lib.tko import utils
23
24__author__ = 'darrenkuo@google.com (Darren Kuo)'
25
26mktime = time.mktime
27datetime = datetime.datetime
28
29class JobSerializer(object):
30    """A class that takes a job object of the tko module and package
31    it with a protocol buffer.
32
33    This class will take a model.job object as input and create a
34    protocol buffer to include all the content of the job object. This
35    protocol buffer object will be serialized into a binary file.
36    """
37
38    def __init__(self):
39
40        self.job_type_dict = {'dir':str, 'tests':list, 'user':str,
41                              'label':str, 'machine':str,
42                              'queued_time':datetime,
43                              'started_time':datetime,
44                              'finished_time':datetime,
45                              'machine_owner':str,
46                              'machine_group':str, 'aborted_by':str,
47                              'aborted_on':datetime,
48                              'keyval_dict':dict}
49
50        self.test_type_dict = {'subdir':str, 'testname':str,
51                               'status':str, 'reason':str,
52                               'kernel':models.kernel, 'machine':str,
53                               'started_time':datetime,
54                               'finished_time':datetime,
55                               'iterations':list, 'attributes':dict,
56                               'labels':list}
57
58        self.kernel_type_dict = {'base':str, 'kernel_hash':str}
59
60        self.iteration_type_dict = {'index':int, 'attr_keyval':dict,
61                                    'perf_keyval':dict}
62
63
64    def deserialize_from_binary(self, infile):
65        """Takes in a binary file name and returns a tko job object.
66
67        The method first deserialize the binary into a protocol buffer
68        job object and then converts the job object into a tko job
69        object.
70
71        @param
72        infile: the name of the binary file that will be deserialized.
73
74        @return a tko job that is represented by the binary file will
75        be returned.
76        """
77
78        job_pb = tko_pb2.Job()
79
80        binary = open(infile, 'r')
81        try:
82            job_pb.ParseFromString(binary.read())
83        finally:
84            binary.close()
85
86        return self.get_tko_job(job_pb)
87
88
89    def serialize_to_binary(self, the_job, tag, binaryfilename):
90        """Serializes the tko job object into a binary by using a
91        protocol buffer.
92
93        The method takes a tko job object and constructs a protocol
94        buffer job object. Then invokes the native serializing
95        function on the object to get a binary string. The string is
96        then written to outfile.
97
98        Precondition: Assumes that all the information about the job
99        is already in the job object. Any fields that is None will be
100        provided a default value.
101
102        @param
103        the_job: the tko job object that will be serialized.
104        tag: contains the job name and the afe_job_id
105        binaryfilename: the name of the file that will be written to
106
107        @return the filename of the file that contains the
108        binary of the serialized object.
109        """
110
111        pb_job = tko_pb2.Job()
112        self.set_pb_job(the_job, pb_job, tag)
113
114        out = open(binaryfilename, 'wb')
115        try:
116            out.write(pb_job.SerializeToString())
117        finally:
118            out.close()
119
120
121    def set_afe_job_id_and_tag(self, pb_job, tag):
122        """Sets the pb job's afe_job_id and tag field.
123
124        @param
125        pb_job: the pb job that will have it's fields set.
126        tag: used to set pb_job.tag and pb_job.afe_job_id.
127        """
128        pb_job.tag = tag
129        pb_job.afe_job_id = utils.get_afe_job_id(tag)
130
131
132    # getter setter methods
133    def get_tko_job(self, job):
134        """Creates a a new tko job object from the pb job object.
135
136        Uses getter methods on the pb objects to extract all the
137        attributes and finally constructs a tko job object using the
138        models.job constructor.
139
140        @param
141        job: a pb job where data is being extracted from.
142
143        @return a tko job object.
144        """
145
146        fields_dict = self.get_trivial_attr(job, self.job_type_dict)
147
148        fields_dict['tests'] = [self.get_tko_test(test) for test in job.tests]
149
150        fields_dict['keyval_dict'] = dict((keyval.name, keyval.value)
151                                          for keyval in job.keyval_dict)
152
153        newjob = models.job(fields_dict['dir'], fields_dict['user'],
154                            fields_dict['label'],
155                            fields_dict['machine'],
156                            fields_dict['queued_time'],
157                            fields_dict['started_time'],
158                            fields_dict['finished_time'],
159                            fields_dict['machine_owner'],
160                            fields_dict['machine_group'],
161                            fields_dict['aborted_by'],
162                            fields_dict['aborted_on'],
163                            fields_dict['keyval_dict'])
164
165        newjob.tests.extend(fields_dict['tests'])
166
167        return newjob
168
169
170    def set_pb_job(self, tko_job, pb_job, tag):
171        """Set the fields for the new job object.
172
173        Method takes in a tko job and an empty protocol buffer job
174        object.  Then safely sets all the appropriate field by first
175        testing if the value in the original object is None.
176
177        @param
178        tko_job: a tko job instance that will have it's values
179        transfered to the new job
180        pb_job: a new instance of the job class provided in the
181        protocol buffer.
182        tag: used to set pb_job.tag and pb_job.afe_job_id.
183        """
184
185        self.set_trivial_attr(tko_job, pb_job, self.job_type_dict)
186        self.set_afe_job_id_and_tag(pb_job, tag)
187
188        for test in tko_job.tests:
189            newtest = pb_job.tests.add()
190            self.set_pb_test(test, newtest)
191
192        for key, val in tko_job.keyval_dict.iteritems():
193            newkeyval = pb_job.keyval_dict.add()
194            newkeyval.name = key
195            newkeyval.value = str(val)
196
197
198    def get_tko_test(self, test):
199        """Creates a tko test from pb_test.
200
201        Extracts data from pb_test by calling helper methods and
202        creates a tko test using the models.test constructor.
203
204        @param:
205        test: a pb_test where fields will be extracted from.
206
207        @return a new instance of models.test
208        """
209        fields_dict = self.get_trivial_attr(test, self.test_type_dict)
210
211        fields_dict['kernel'] = self.get_tko_kernel(test.kernel)
212
213        fields_dict['iterations'] = [self.get_tko_iteration(iteration)
214                                     for iteration in test.iterations]
215
216        fields_dict['attributes'] = dict((keyval.name, keyval.value)
217                                         for keyval in test.attributes)
218
219        fields_dict['labels'] = list(test.labels)
220
221        # The constructor for models.test accepts a "perf_values" list that
222        # represents performance values of the test.  The empty list argument
223        # in the constructor call below represents this value and makes this
224        # code adhere properly to the models.test constructor argument list.
225        # However, the effect of the empty list is that perf values are
226        # ignored in the job_serializer module.  This is ok for now because
227        # autotest does not use the current module.  If job_serializer is used
228        # in the future, we need to modify the "tko.proto" protobuf file to
229        # understand the notion of perf_values, then modify this file
230        # accordingly to use it.
231        return models.test(fields_dict['subdir'],
232                           fields_dict['testname'],
233                           fields_dict['status'],
234                           fields_dict['reason'],
235                           fields_dict['kernel'],
236                           fields_dict['machine'],
237                           fields_dict['started_time'],
238                           fields_dict['finished_time'],
239                           fields_dict['iterations'],
240                           fields_dict['attributes'],
241                           [],
242                           fields_dict['labels'])
243
244
245    def set_pb_test(self, tko_test, pb_test):
246        """Sets the various fields of test object of the tko protocol.
247
248        Method takes a tko test and a new test of the protocol buffer and
249        transfers the values in the tko test to the new test.
250
251        @param
252        tko_test: a tko test instance.
253        pb_test: an empty protocol buffer test instance.
254
255        """
256
257        self.set_trivial_attr(tko_test, pb_test, self.test_type_dict)
258
259        self.set_pb_kernel(tko_test.kernel, pb_test.kernel)
260
261        for current_iteration in tko_test.iterations:
262            pb_iteration = pb_test.iterations.add()
263            self.set_pb_iteration(current_iteration, pb_iteration)
264
265        for key, val in tko_test.attributes.iteritems():
266            newkeyval = pb_test.attributes.add()
267            newkeyval.name = key
268            newkeyval.value = str(val)
269
270        for current_label in tko_test.labels:
271            pb_test.labels.append(current_label)
272
273
274    def get_tko_kernel(self, kernel):
275        """Constructs a new tko kernel object from a pb kernel object.
276
277        Uses all the getter methods on the pb kernel object to extract
278        the attributes and constructs a new tko kernel object using
279        the model.kernel constructor.
280
281        @param
282        kernel: a pb kernel object where data will be extracted.
283
284        @return a new tko kernel object.
285        """
286
287        fields_dict = self.get_trivial_attr(kernel, self.kernel_type_dict)
288
289        return models.kernel(fields_dict['base'], [], fields_dict['kernel_hash'])
290
291
292    def set_pb_kernel(self, tko_kernel, pb_kernel):
293        """Set a specific kernel of a test.
294
295        Takes the same form of all the other setting methods.  It
296        seperates the string variables from the int variables and set
297        them safely.
298
299        @param
300        tko_kernel: a tko kernel.
301        pb_kernel: an empty protocol buffer kernel.
302
303        """
304
305        self.set_trivial_attr(tko_kernel, pb_kernel, self.kernel_type_dict)
306
307
308    def get_tko_iteration(self, iteration):
309        """Creates a new tko iteration with the data in the provided
310        pb iteration.
311
312        Uses the data in the pb iteration and the models.iteration
313        constructor to create a new tko iterations
314
315        @param
316        iteration: a pb iteration instance
317
318        @return a tko iteration instance with the same data.
319        """
320
321        fields_dict = self.get_trivial_attr(iteration,
322                                            self.iteration_type_dict)
323
324        fields_dict['attr_keyval'] = dict((keyval.name, keyval.value)
325                                          for keyval in iteration.attr_keyval)
326
327        fields_dict['perf_keyval'] = dict((keyval.name, keyval.value)
328                                          for keyval in iteration.perf_keyval)
329
330        return models.iteration(fields_dict['index'],
331                                fields_dict['attr_keyval'],
332                                fields_dict['perf_keyval'])
333
334
335    def set_pb_iteration(self, tko_iteration, pb_iteration):
336        """Sets all fields for a particular iteration.
337
338        Takes same form as all the other setting methods. Sets int,
339        str and datetime variables safely.
340
341        @param
342        tko_iteration: a tko test iteration.
343        pb_iteration: an empty pb test iteration.
344
345        """
346
347        self.set_trivial_attr(tko_iteration, pb_iteration,
348                              self.iteration_type_dict)
349
350        for key, val in tko_iteration.attr_keyval.iteritems():
351            newkeyval = pb_iteration.attr_keyval.add()
352            newkeyval.name = key
353            newkeyval.value = str(val)
354
355        for key, val in tko_iteration.perf_keyval.iteritems():
356            newkeyval = pb_iteration.perf_keyval.add()
357            newkeyval.name = key
358            newkeyval.value = str(val)
359
360
361    def get_trivial_attr(self, obj, objdict):
362        """Get all trivial attributes from the object.
363
364        This function is used to extract attributes from a pb job. The
365        dictionary specifies the types of each attribute in each tko
366        class.
367
368        @param
369        obj: the pb object that is being extracted.
370        objdict: the dict that specifies the type.
371
372        @return a dict of each attr name and it's corresponding value.
373        """
374
375        resultdict = {}
376        for field, field_type in objdict.items():
377            value = getattr(obj, field)
378            if field_type in (str, int, long):
379                resultdict[field] = field_type(value)
380            elif field_type == datetime:
381                resultdict[field] = (
382                            datetime.fromtimestamp(value/1000.0))
383
384        return resultdict
385
386
387    def set_trivial_attr(self, tko_obj, pb_obj, objdict):
388        """Sets all the easy attributes appropriately according to the
389        type.
390
391        This function is used to set all the trivial attributes
392        provided by objdict, the dictionary that specifies the types
393        of each attribute in each tko class.
394
395        @param
396        tko_obj: the original object that has the data being copied.
397        pb_obj: the new pb object that is being copied into.
398        objdict: specifies the type of each attribute in the class we
399        are working with.
400
401        """
402        for attr, attr_type in objdict.iteritems():
403            if attr_type == datetime:
404                t = getattr(tko_obj, attr)
405                if not t:
406                    self.set_attr_safely(pb_obj, attr, t, int)
407                else:
408                    t = mktime(t.timetuple()) + 1e-6 * t.microsecond
409                    setattr(pb_obj, attr, long(t*1000))
410            else:
411                value = getattr(tko_obj, attr)
412                self.set_attr_safely(pb_obj, attr, value, attr_type)
413
414
415    def set_attr_safely(self, var, attr, value, vartype):
416        """Sets a particular attribute of var if the provided value is
417        not None.
418
419        Checks if value is None. If not, set the attribute of the var
420        to be the default value. This is necessary for the special
421        required fields of the protocol buffer.
422
423        @param
424        var: the variable of which one of the attribute is being set.
425        attr: the attribute that is being set.
426        value: the value that is being checked
427        vartype: the expected type of the attr
428
429        """
430
431        supported_types = [int, long, str]
432        if vartype in supported_types:
433            if value is None:
434                value = vartype()
435            else:
436                assert isinstance(value, vartype), (
437                'Unexpected type %s for attr %s, should be %s' %
438                (type(value), attr, vartype))
439
440            setattr(var, attr, value)
441