tko/job_serializer.py

   1 #!/usr/bin/python
   2
   3 """A script that provides convertion between models.job and a protocol
   4 buffer object.
   5
   6 This script contains only one class that takes an job instance and
   7 convert it into a protocol buffer object. The class will also be
   8 responsible for serializing the job instance via protocol buffers.
   9
  10 """
  11
  12 # import python libraries
  13 import os
  14 import datetime
  15 import time
  16 import random
  17 import re
  18
  19 # import autotest libraries
  20 from autotest_lib.tko import models
  21 from autotest_lib.tko import tko_pb2
  22 from autotest_lib.tko import utils
  23
  24 __author__ = 'darrenkuo@google.com (Darren Kuo)'
  25
  26 mktime = time.mktime
  27 datetime = datetime.datetime
  28
  29 class 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         return models.test(fields_dict['subdir'],
 222                            fields_dict['testname'],
 223                            fields_dict['status'],
 224                            fields_dict['reason'],
 225                            fields_dict['kernel'],
 226                            fields_dict['machine'],
 227                            fields_dict['started_time'],
 228                            fields_dict['finished_time'],
 229                            fields_dict['iterations'],
 230                            fields_dict['attributes'],
 231                            fields_dict['labels'])
 232
 233
 234     def set_pb_test(self, tko_test, pb_test):
 235         """Sets the various fields of test object of the tko protocol.
 236
 237         Method takes a tko test and a new test of the protocol buffer and
 238         transfers the values in the tko test to the new test.
 239
 240         @param
 241         tko_test: a tko test instance.
 242         pb_test: an empty protocol buffer test instance.
 243
 244         """
 245
 246         self.set_trivial_attr(tko_test, pb_test, self.test_type_dict)
 247
 248         self.set_pb_kernel(tko_test.kernel, pb_test.kernel)
 249
 250         for current_iteration in tko_test.iterations:
 251             pb_iteration = pb_test.iterations.add()
 252             self.set_pb_iteration(current_iteration, pb_iteration)
 253
 254         for key, val in tko_test.attributes.iteritems():
 255             newkeyval = pb_test.attributes.add()
 256             newkeyval.name = key
 257             newkeyval.value = str(val)
 258
 259         for current_label in tko_test.labels:
 260             pb_test.labels.append(current_label)
 261
 262
 263     def get_tko_kernel(self, kernel):
 264         """Constructs a new tko kernel object from a pb kernel object.
 265
 266         Uses all the getter methods on the pb kernel object to extract
 267         the attributes and constructs a new tko kernel object using
 268         the model.kernel constructor.
 269
 270         @param
 271         kernel: a pb kernel object where data will be extracted.
 272
 273         @return a new tko kernel object.
 274         """
 275
 276         fields_dict = self.get_trivial_attr(kernel, self.kernel_type_dict)
 277
 278         return models.kernel(fields_dict['base'], [], fields_dict['kernel_hash'])
 279
 280
 281     def set_pb_kernel(self, tko_kernel, pb_kernel):
 282         """Set a specific kernel of a test.
 283
 284         Takes the same form of all the other setting methods.  It
 285         seperates the string variables from the int variables and set
 286         them safely.
 287
 288         @param
 289         tko_kernel: a tko kernel.
 290         pb_kernel: an empty protocol buffer kernel.
 291
 292         """
 293
 294         self.set_trivial_attr(tko_kernel, pb_kernel, self.kernel_type_dict)
 295
 296
 297     def get_tko_iteration(self, iteration):
 298         """Creates a new tko iteration with the data in the provided
 299         pb iteration.
 300
 301         Uses the data in the pb iteration and the models.iteration
 302         constructor to create a new tko iterations
 303
 304         @param
 305         iteration: a pb iteration instance
 306
 307         @return a tko iteration instance with the same data.
 308         """
 309
 310         fields_dict = self.get_trivial_attr(iteration,
 311                                             self.iteration_type_dict)
 312
 313         fields_dict['attr_keyval'] = dict((keyval.name, keyval.value)
 314                                           for keyval in iteration.attr_keyval)
 315
 316         fields_dict['perf_keyval'] = dict((keyval.name, keyval.value)
 317                                           for keyval in iteration.perf_keyval)
 318
 319         return models.iteration(fields_dict['index'],
 320                                 fields_dict['attr_keyval'],
 321                                 fields_dict['perf_keyval'])
 322
 323
 324     def set_pb_iteration(self, tko_iteration, pb_iteration):
 325         """Sets all fields for a particular iteration.
 326
 327         Takes same form as all the other setting methods. Sets int,
 328         str and datetime variables safely.
 329
 330         @param
 331         tko_iteration: a tko test iteration.
 332         pb_iteration: an empty pb test iteration.
 333
 334         """
 335
 336         self.set_trivial_attr(tko_iteration, pb_iteration,
 337                               self.iteration_type_dict)
 338
 339         for key, val in tko_iteration.attr_keyval.iteritems():
 340             newkeyval = pb_iteration.attr_keyval.add()
 341             newkeyval.name = key
 342             newkeyval.value = str(val)
 343
 344         for key, val in tko_iteration.perf_keyval.iteritems():
 345             newkeyval = pb_iteration.perf_keyval.add()
 346             newkeyval.name = key
 347             newkeyval.value = str(val)
 348
 349
 350     def get_trivial_attr(self, obj, objdict):
 351         """Get all trivial attributes from the object.
 352
 353         This function is used to extract attributes from a pb job. The
 354         dictionary specifies the types of each attribute in each tko
 355         class.
 356
 357         @param
 358         obj: the pb object that is being extracted.
 359         objdict: the dict that specifies the type.
 360
 361         @return a dict of each attr name and it's corresponding value.
 362         """
 363
 364         resultdict = {}
 365         for field, field_type in objdict.items():
 366             value = getattr(obj, field)
 367             if field_type in (str, int, long):
 368                 resultdict[field] = field_type(value)
 369             elif field_type == datetime:
 370                 resultdict[field] = (
 371                             datetime.fromtimestamp(value/1000.0))
 372
 373         return resultdict
 374
 375
 376     def set_trivial_attr(self, tko_obj, pb_obj, objdict):
 377         """Sets all the easy attributes appropriately according to the
 378         type.
 379
 380         This function is used to set all the trivial attributes
 381         provided by objdict, the dictionary that specifies the types
 382         of each attribute in each tko class.
 383
 384         @param
 385         tko_obj: the original object that has the data being copied.
 386         pb_obj: the new pb object that is being copied into.
 387         objdict: specifies the type of each attribute in the class we
 388         are working with.
 389
 390         """
 391         for attr, attr_type in objdict.iteritems():
 392             if attr_type == datetime:
 393                 t = getattr(tko_obj, attr)
 394                 if not t:
 395                     self.set_attr_safely(pb_obj, attr, t, int)
 396                 else:
 397                     t = mktime(t.timetuple()) + 1e-6 * t.microsecond
 398                     setattr(pb_obj, attr, long(t*1000))
 399             else:
 400                 value = getattr(tko_obj, attr)
 401                 self.set_attr_safely(pb_obj, attr, value, attr_type)
 402
 403
 404     def set_attr_safely(self, var, attr, value, vartype):
 405         """Sets a particular attribute of var if the provided value is
 406         not None.
 407
 408         Checks if value is None. If not, set the attribute of the var
 409         to be the default value. This is necessary for the special
 410         required fields of the protocol buffer.
 411
 412         @param
 413         var: the variable of which one of the attribute is being set.
 414         attr: the attribute that is being set.
 415         value: the value that is being checked
 416         vartype: the expected type of the attr
 417
 418         """
 419
 420         supported_types = [int, long, str]
 421         if vartype in supported_types:
 422             if value is None:
 423                 value = vartype()
 424             else:
 425                 assert isinstance(value, vartype), (
 426                 'Unexpected type %s for attr %s, should be %s' %
 427                 (type(value), attr, vartype))
 428
 429             setattr(var, attr, value)