blob: cec7dad6bdf2fb6d703340fe57a91136b118d68d [file] [log] [blame]
showard7c785282008-05-29 19:45:12 +00001"""
2Extensions to Django's model logic.
3"""
4
showarda5288b42009-07-28 20:06:08 +00005import django.core.exceptions
Prashanth Balasubramanian75be1d32014-11-25 18:03:09 -08006from django.db import connection
7from django.db import connections
8from django.db import models as dbmodels
9from django.db import transaction
showarda5288b42009-07-28 20:06:08 +000010from django.db.models.sql import query
showard7e67b432010-01-20 01:13:04 +000011import django.db.models.sql.where
MK Ryu5cfd96a2015-01-30 15:31:23 -080012from autotest_lib.client.common_lib.cros.graphite import autotest_stats
Prashanth B489b91d2014-03-15 12:17:16 -070013from autotest_lib.frontend.afe import rdb_model_extensions
showard7c785282008-05-29 19:45:12 +000014
Prashanth B489b91d2014-03-15 12:17:16 -070015
16class ValidationError(django.core.exceptions.ValidationError):
jadmanski0afbb632008-06-06 21:10:57 +000017 """\
showarda5288b42009-07-28 20:06:08 +000018 Data validation error in adding or updating an object. The associated
jadmanski0afbb632008-06-06 21:10:57 +000019 value is a dictionary mapping field names to error strings.
20 """
showard7c785282008-05-29 19:45:12 +000021
showarda5288b42009-07-28 20:06:08 +000022def _quote_name(name):
23 """Shorthand for connection.ops.quote_name()."""
24 return connection.ops.quote_name(name)
25
26
beepscc9fc702013-12-02 12:45:38 -080027class LeasedHostManager(dbmodels.Manager):
28 """Query manager for unleased, unlocked hosts.
29 """
30 def get_query_set(self):
31 return (super(LeasedHostManager, self).get_query_set().filter(
32 leased=0, locked=0))
33
34
showard7c785282008-05-29 19:45:12 +000035class ExtendedManager(dbmodels.Manager):
jadmanski0afbb632008-06-06 21:10:57 +000036 """\
37 Extended manager supporting subquery filtering.
38 """
showard7c785282008-05-29 19:45:12 +000039
showardf828c772010-01-25 21:49:42 +000040 class CustomQuery(query.Query):
showard7e67b432010-01-20 01:13:04 +000041 def __init__(self, *args, **kwargs):
showardf828c772010-01-25 21:49:42 +000042 super(ExtendedManager.CustomQuery, self).__init__(*args, **kwargs)
showard7e67b432010-01-20 01:13:04 +000043 self._custom_joins = []
44
45
showarda5288b42009-07-28 20:06:08 +000046 def clone(self, klass=None, **kwargs):
showardf828c772010-01-25 21:49:42 +000047 obj = super(ExtendedManager.CustomQuery, self).clone(klass)
showard7e67b432010-01-20 01:13:04 +000048 obj._custom_joins = list(self._custom_joins)
showarda5288b42009-07-28 20:06:08 +000049 return obj
showard08f981b2008-06-24 21:59:03 +000050
showard7e67b432010-01-20 01:13:04 +000051
52 def combine(self, rhs, connector):
showardf828c772010-01-25 21:49:42 +000053 super(ExtendedManager.CustomQuery, self).combine(rhs, connector)
showard7e67b432010-01-20 01:13:04 +000054 if hasattr(rhs, '_custom_joins'):
55 self._custom_joins.extend(rhs._custom_joins)
56
57
58 def add_custom_join(self, table, condition, join_type,
59 condition_values=(), alias=None):
60 if alias is None:
61 alias = table
62 join_dict = dict(table=table,
63 condition=condition,
64 condition_values=condition_values,
65 join_type=join_type,
66 alias=alias)
67 self._custom_joins.append(join_dict)
68
69
showard7e67b432010-01-20 01:13:04 +000070 @classmethod
71 def convert_query(self, query_set):
72 """
showardf828c772010-01-25 21:49:42 +000073 Convert the query set's "query" attribute to a CustomQuery.
showard7e67b432010-01-20 01:13:04 +000074 """
75 # Make a copy of the query set
76 query_set = query_set.all()
77 query_set.query = query_set.query.clone(
showardf828c772010-01-25 21:49:42 +000078 klass=ExtendedManager.CustomQuery,
showard7e67b432010-01-20 01:13:04 +000079 _custom_joins=[])
80 return query_set
showard43a3d262008-11-12 18:17:05 +000081
82
showard7e67b432010-01-20 01:13:04 +000083 class _WhereClause(object):
84 """Object allowing us to inject arbitrary SQL into Django queries.
showard43a3d262008-11-12 18:17:05 +000085
showard7e67b432010-01-20 01:13:04 +000086 By using this instead of extra(where=...), we can still freely combine
87 queries with & and |.
showarda5288b42009-07-28 20:06:08 +000088 """
showard7e67b432010-01-20 01:13:04 +000089 def __init__(self, clause, values=()):
90 self._clause = clause
91 self._values = values
showarda5288b42009-07-28 20:06:08 +000092
showard7e67b432010-01-20 01:13:04 +000093
Dale Curtis74a314b2011-06-23 14:55:46 -070094 def as_sql(self, qn=None, connection=None):
showard7e67b432010-01-20 01:13:04 +000095 return self._clause, self._values
96
97
98 def relabel_aliases(self, change_map):
99 return
showard43a3d262008-11-12 18:17:05 +0000100
101
showard8b0ea222009-12-23 19:23:03 +0000102 def add_join(self, query_set, join_table, join_key, join_condition='',
showard7e67b432010-01-20 01:13:04 +0000103 join_condition_values=(), join_from_key=None, alias=None,
104 suffix='', exclude=False, force_left_join=False):
105 """Add a join to query_set.
106
107 Join looks like this:
108 (INNER|LEFT) JOIN <join_table> AS <alias>
109 ON (<this table>.<join_from_key> = <join_table>.<join_key>
110 and <join_condition>)
111
showard0957a842009-05-11 19:25:08 +0000112 @param join_table table to join to
113 @param join_key field referencing back to this model to use for the join
114 @param join_condition extra condition for the ON clause of the join
showard7e67b432010-01-20 01:13:04 +0000115 @param join_condition_values values to substitute into join_condition
116 @param join_from_key column on this model to join from.
showard8b0ea222009-12-23 19:23:03 +0000117 @param alias alias to use for for join
118 @param suffix suffix to add to join_table for the join alias, if no
119 alias is provided
showard0957a842009-05-11 19:25:08 +0000120 @param exclude if true, exclude rows that match this join (will use a
showarda5288b42009-07-28 20:06:08 +0000121 LEFT OUTER JOIN and an appropriate WHERE condition)
showardc4780402009-08-31 18:31:34 +0000122 @param force_left_join - if true, a LEFT OUTER JOIN will be used
123 instead of an INNER JOIN regardless of other options
showard0957a842009-05-11 19:25:08 +0000124 """
showard7e67b432010-01-20 01:13:04 +0000125 join_from_table = query_set.model._meta.db_table
126 if join_from_key is None:
127 join_from_key = self.model._meta.pk.name
128 if alias is None:
129 alias = join_table + suffix
130 full_join_key = _quote_name(alias) + '.' + _quote_name(join_key)
131 full_join_condition = '%s = %s.%s' % (full_join_key,
132 _quote_name(join_from_table),
133 _quote_name(join_from_key))
showard43a3d262008-11-12 18:17:05 +0000134 if join_condition:
135 full_join_condition += ' AND (' + join_condition + ')'
136 if exclude or force_left_join:
showarda5288b42009-07-28 20:06:08 +0000137 join_type = query_set.query.LOUTER
showard43a3d262008-11-12 18:17:05 +0000138 else:
showarda5288b42009-07-28 20:06:08 +0000139 join_type = query_set.query.INNER
showard43a3d262008-11-12 18:17:05 +0000140
showardf828c772010-01-25 21:49:42 +0000141 query_set = self.CustomQuery.convert_query(query_set)
showard7e67b432010-01-20 01:13:04 +0000142 query_set.query.add_custom_join(join_table,
143 full_join_condition,
144 join_type,
145 condition_values=join_condition_values,
146 alias=alias)
showard43a3d262008-11-12 18:17:05 +0000147
showard7e67b432010-01-20 01:13:04 +0000148 if exclude:
149 query_set = query_set.extra(where=[full_join_key + ' IS NULL'])
150
151 return query_set
152
153
154 def _info_for_many_to_one_join(self, field, join_to_query, alias):
155 """
156 @param field: the ForeignKey field on the related model
157 @param join_to_query: the query over the related model that we're
158 joining to
159 @param alias: alias of joined table
160 """
161 info = {}
162 rhs_table = join_to_query.model._meta.db_table
163 info['rhs_table'] = rhs_table
164 info['rhs_column'] = field.column
165 info['lhs_column'] = field.rel.get_related_field().column
166 rhs_where = join_to_query.query.where
167 rhs_where.relabel_aliases({rhs_table: alias})
Dale Curtis74a314b2011-06-23 14:55:46 -0700168 compiler = join_to_query.query.get_compiler(using=join_to_query.db)
169 initial_clause, values = compiler.as_sql()
170 all_clauses = (initial_clause,)
171 if hasattr(join_to_query.query, 'extra_where'):
172 all_clauses += join_to_query.query.extra_where
173 info['where_clause'] = (
174 ' AND '.join('(%s)' % clause for clause in all_clauses))
showard7e67b432010-01-20 01:13:04 +0000175 info['values'] = values
176 return info
177
178
179 def _info_for_many_to_many_join(self, m2m_field, join_to_query, alias,
180 m2m_is_on_this_model):
181 """
182 @param m2m_field: a Django field representing the M2M relationship.
183 It uses a pivot table with the following structure:
184 this model table <---> M2M pivot table <---> joined model table
185 @param join_to_query: the query over the related model that we're
186 joining to.
187 @param alias: alias of joined table
188 """
189 if m2m_is_on_this_model:
190 # referenced field on this model
191 lhs_id_field = self.model._meta.pk
192 # foreign key on the pivot table referencing lhs_id_field
193 m2m_lhs_column = m2m_field.m2m_column_name()
194 # foreign key on the pivot table referencing rhd_id_field
195 m2m_rhs_column = m2m_field.m2m_reverse_name()
196 # referenced field on related model
197 rhs_id_field = m2m_field.rel.get_related_field()
198 else:
199 lhs_id_field = m2m_field.rel.get_related_field()
200 m2m_lhs_column = m2m_field.m2m_reverse_name()
201 m2m_rhs_column = m2m_field.m2m_column_name()
202 rhs_id_field = join_to_query.model._meta.pk
203
204 info = {}
205 info['rhs_table'] = m2m_field.m2m_db_table()
206 info['rhs_column'] = m2m_lhs_column
207 info['lhs_column'] = lhs_id_field.column
208
209 # select the ID of related models relevant to this join. we can only do
210 # a single join, so we need to gather this information up front and
211 # include it in the join condition.
212 rhs_ids = join_to_query.values_list(rhs_id_field.attname, flat=True)
213 assert len(rhs_ids) == 1, ('Many-to-many custom field joins can only '
214 'match a single related object.')
215 rhs_id = rhs_ids[0]
216
217 info['where_clause'] = '%s.%s = %s' % (_quote_name(alias),
218 _quote_name(m2m_rhs_column),
219 rhs_id)
220 info['values'] = ()
221 return info
222
223
224 def join_custom_field(self, query_set, join_to_query, alias,
225 left_join=True):
226 """Join to a related model to create a custom field in the given query.
227
228 This method is used to construct a custom field on the given query based
229 on a many-valued relationsip. join_to_query should be a simple query
230 (no joins) on the related model which returns at most one related row
231 per instance of this model.
232
233 For many-to-one relationships, the joined table contains the matching
234 row from the related model it one is related, NULL otherwise.
235
236 For many-to-many relationships, the joined table contains the matching
237 row if it's related, NULL otherwise.
238 """
239 relationship_type, field = self.determine_relationship(
240 join_to_query.model)
241
242 if relationship_type == self.MANY_TO_ONE:
243 info = self._info_for_many_to_one_join(field, join_to_query, alias)
244 elif relationship_type == self.M2M_ON_RELATED_MODEL:
245 info = self._info_for_many_to_many_join(
246 m2m_field=field, join_to_query=join_to_query, alias=alias,
247 m2m_is_on_this_model=False)
248 elif relationship_type ==self.M2M_ON_THIS_MODEL:
249 info = self._info_for_many_to_many_join(
250 m2m_field=field, join_to_query=join_to_query, alias=alias,
251 m2m_is_on_this_model=True)
252
253 return self.add_join(query_set, info['rhs_table'], info['rhs_column'],
254 join_from_key=info['lhs_column'],
255 join_condition=info['where_clause'],
256 join_condition_values=info['values'],
257 alias=alias,
258 force_left_join=left_join)
259
260
261 def add_where(self, query_set, where, values=()):
262 query_set = query_set.all()
263 query_set.query.where.add(self._WhereClause(where, values),
264 django.db.models.sql.where.AND)
showardc4780402009-08-31 18:31:34 +0000265 return query_set
showard7c785282008-05-29 19:45:12 +0000266
267
showardeaccf8f2009-04-16 03:11:33 +0000268 def _get_quoted_field(self, table, field):
showarda5288b42009-07-28 20:06:08 +0000269 return _quote_name(table) + '.' + _quote_name(field)
showard5ef36e92008-07-02 16:37:09 +0000270
271
showard7c199df2008-10-03 10:17:15 +0000272 def get_key_on_this_table(self, key_field=None):
showard5ef36e92008-07-02 16:37:09 +0000273 if key_field is None:
274 # default to primary key
275 key_field = self.model._meta.pk.column
276 return self._get_quoted_field(self.model._meta.db_table, key_field)
277
278
showardeaccf8f2009-04-16 03:11:33 +0000279 def escape_user_sql(self, sql):
280 return sql.replace('%', '%%')
281
showard5ef36e92008-07-02 16:37:09 +0000282
showard0957a842009-05-11 19:25:08 +0000283 def _custom_select_query(self, query_set, selects):
Jakob Juelich7bef8412014-10-14 19:11:54 -0700284 """Execute a custom select query.
285
286 @param query_set: query set as returned by query_objects.
287 @param selects: Tables/Columns to select, e.g. tko_test_labels_list.id.
288
289 @returns: Result of the query as returned by cursor.fetchall().
290 """
Dale Curtis74a314b2011-06-23 14:55:46 -0700291 compiler = query_set.query.get_compiler(using=query_set.db)
292 sql, params = compiler.as_sql()
showarda5288b42009-07-28 20:06:08 +0000293 from_ = sql[sql.find(' FROM'):]
294
295 if query_set.query.distinct:
showard0957a842009-05-11 19:25:08 +0000296 distinct = 'DISTINCT '
297 else:
298 distinct = ''
showarda5288b42009-07-28 20:06:08 +0000299
300 sql_query = ('SELECT ' + distinct + ','.join(selects) + from_)
Jakob Juelich7bef8412014-10-14 19:11:54 -0700301 # Chose the connection that's responsible for this type of object
302 cursor = connections[query_set.db].cursor()
showard0957a842009-05-11 19:25:08 +0000303 cursor.execute(sql_query, params)
304 return cursor.fetchall()
305
306
showard68693f72009-05-20 00:31:53 +0000307 def _is_relation_to(self, field, model_class):
308 return field.rel and field.rel.to is model_class
showard0957a842009-05-11 19:25:08 +0000309
310
showard7e67b432010-01-20 01:13:04 +0000311 MANY_TO_ONE = object()
312 M2M_ON_RELATED_MODEL = object()
313 M2M_ON_THIS_MODEL = object()
314
315 def determine_relationship(self, related_model):
316 """
317 Determine the relationship between this model and related_model.
318
319 related_model must have some sort of many-valued relationship to this
320 manager's model.
321 @returns (relationship_type, field), where relationship_type is one of
322 MANY_TO_ONE, M2M_ON_RELATED_MODEL, M2M_ON_THIS_MODEL, and field
323 is the Django field object for the relationship.
324 """
325 # look for a foreign key field on related_model relating to this model
326 for field in related_model._meta.fields:
327 if self._is_relation_to(field, self.model):
328 return self.MANY_TO_ONE, field
329
330 # look for an M2M field on related_model relating to this model
331 for field in related_model._meta.many_to_many:
332 if self._is_relation_to(field, self.model):
333 return self.M2M_ON_RELATED_MODEL, field
334
335 # maybe this model has the many-to-many field
336 for field in self.model._meta.many_to_many:
337 if self._is_relation_to(field, related_model):
338 return self.M2M_ON_THIS_MODEL, field
339
340 raise ValueError('%s has no relation to %s' %
341 (related_model, self.model))
342
343
showard68693f72009-05-20 00:31:53 +0000344 def _get_pivot_iterator(self, base_objects_by_id, related_model):
showard0957a842009-05-11 19:25:08 +0000345 """
showard68693f72009-05-20 00:31:53 +0000346 Determine the relationship between this model and related_model, and
347 return a pivot iterator.
348 @param base_objects_by_id: dict of instances of this model indexed by
349 their IDs
350 @returns a pivot iterator, which yields a tuple (base_object,
351 related_object) for each relationship between a base object and a
352 related object. all base_object instances come from base_objects_by_id.
showard7e67b432010-01-20 01:13:04 +0000353 Note -- this depends on Django model internals.
showard0957a842009-05-11 19:25:08 +0000354 """
showard7e67b432010-01-20 01:13:04 +0000355 relationship_type, field = self.determine_relationship(related_model)
356 if relationship_type == self.MANY_TO_ONE:
357 return self._many_to_one_pivot(base_objects_by_id,
358 related_model, field)
359 elif relationship_type == self.M2M_ON_RELATED_MODEL:
360 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000361 base_objects_by_id, related_model, field.m2m_db_table(),
362 field.m2m_reverse_name(), field.m2m_column_name())
showard7e67b432010-01-20 01:13:04 +0000363 else:
364 assert relationship_type == self.M2M_ON_THIS_MODEL
365 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000366 base_objects_by_id, related_model, field.m2m_db_table(),
367 field.m2m_column_name(), field.m2m_reverse_name())
showard0957a842009-05-11 19:25:08 +0000368
showard0957a842009-05-11 19:25:08 +0000369
showard68693f72009-05-20 00:31:53 +0000370 def _many_to_one_pivot(self, base_objects_by_id, related_model,
371 foreign_key_field):
372 """
373 @returns a pivot iterator - see _get_pivot_iterator()
374 """
375 filter_data = {foreign_key_field.name + '__pk__in':
376 base_objects_by_id.keys()}
377 for related_object in related_model.objects.filter(**filter_data):
showarda5a72c92009-08-20 23:35:21 +0000378 # lookup base object in the dict, rather than grabbing it from the
379 # related object. we need to return instances from the dict, not
380 # fresh instances of the same models (and grabbing model instances
381 # from the related models incurs a DB query each time).
382 base_object_id = getattr(related_object, foreign_key_field.attname)
383 base_object = base_objects_by_id[base_object_id]
showard68693f72009-05-20 00:31:53 +0000384 yield base_object, related_object
385
386
387 def _query_pivot_table(self, base_objects_by_id, pivot_table,
Jakob Juelich7bef8412014-10-14 19:11:54 -0700388 pivot_from_field, pivot_to_field, related_model):
showard0957a842009-05-11 19:25:08 +0000389 """
390 @param id_list list of IDs of self.model objects to include
391 @param pivot_table the name of the pivot table
392 @param pivot_from_field a field name on pivot_table referencing
393 self.model
394 @param pivot_to_field a field name on pivot_table referencing the
395 related model.
Jakob Juelich7bef8412014-10-14 19:11:54 -0700396 @param related_model the related model
397
showard68693f72009-05-20 00:31:53 +0000398 @returns pivot list of IDs (base_id, related_id)
showard0957a842009-05-11 19:25:08 +0000399 """
400 query = """
401 SELECT %(from_field)s, %(to_field)s
402 FROM %(table)s
403 WHERE %(from_field)s IN (%(id_list)s)
404 """ % dict(from_field=pivot_from_field,
405 to_field=pivot_to_field,
406 table=pivot_table,
showard68693f72009-05-20 00:31:53 +0000407 id_list=','.join(str(id_) for id_
408 in base_objects_by_id.iterkeys()))
Jakob Juelich7bef8412014-10-14 19:11:54 -0700409
410 # Chose the connection that's responsible for this type of object
411 # The databases for related_model and the current model will always
412 # be the same, related_model is just easier to obtain here because
413 # self is only a ExtendedManager, not the object.
414 cursor = connections[related_model.objects.db].cursor()
showard0957a842009-05-11 19:25:08 +0000415 cursor.execute(query)
showard68693f72009-05-20 00:31:53 +0000416 return cursor.fetchall()
showard0957a842009-05-11 19:25:08 +0000417
418
showard68693f72009-05-20 00:31:53 +0000419 def _many_to_many_pivot(self, base_objects_by_id, related_model,
420 pivot_table, pivot_from_field, pivot_to_field):
421 """
422 @param pivot_table: see _query_pivot_table
423 @param pivot_from_field: see _query_pivot_table
424 @param pivot_to_field: see _query_pivot_table
425 @returns a pivot iterator - see _get_pivot_iterator()
426 """
427 id_pivot = self._query_pivot_table(base_objects_by_id, pivot_table,
Jakob Juelich7bef8412014-10-14 19:11:54 -0700428 pivot_from_field, pivot_to_field,
429 related_model)
showard68693f72009-05-20 00:31:53 +0000430
431 all_related_ids = list(set(related_id for base_id, related_id
432 in id_pivot))
433 related_objects_by_id = related_model.objects.in_bulk(all_related_ids)
434
435 for base_id, related_id in id_pivot:
436 yield base_objects_by_id[base_id], related_objects_by_id[related_id]
437
438
439 def populate_relationships(self, base_objects, related_model,
showard0957a842009-05-11 19:25:08 +0000440 related_list_name):
441 """
showard68693f72009-05-20 00:31:53 +0000442 For each instance of this model in base_objects, add a field named
443 related_list_name listing all the related objects of type related_model.
444 related_model must be in a many-to-one or many-to-many relationship with
445 this model.
446 @param base_objects - list of instances of this model
447 @param related_model - model class related to this model
448 @param related_list_name - attribute name in which to store the related
449 object list.
showard0957a842009-05-11 19:25:08 +0000450 """
showard68693f72009-05-20 00:31:53 +0000451 if not base_objects:
showard0957a842009-05-11 19:25:08 +0000452 # if we don't bail early, we'll get a SQL error later
453 return
showard0957a842009-05-11 19:25:08 +0000454
showard68693f72009-05-20 00:31:53 +0000455 base_objects_by_id = dict((base_object._get_pk_val(), base_object)
456 for base_object in base_objects)
457 pivot_iterator = self._get_pivot_iterator(base_objects_by_id,
458 related_model)
showard0957a842009-05-11 19:25:08 +0000459
showard68693f72009-05-20 00:31:53 +0000460 for base_object in base_objects:
461 setattr(base_object, related_list_name, [])
462
463 for base_object, related_object in pivot_iterator:
464 getattr(base_object, related_list_name).append(related_object)
showard0957a842009-05-11 19:25:08 +0000465
466
jamesrene3656232010-03-02 00:00:30 +0000467class ModelWithInvalidQuerySet(dbmodels.query.QuerySet):
468 """
469 QuerySet that handles delete() properly for models with an "invalid" bit
470 """
471 def delete(self):
472 for model in self:
473 model.delete()
474
475
476class ModelWithInvalidManager(ExtendedManager):
477 """
478 Manager for objects with an "invalid" bit
479 """
480 def get_query_set(self):
481 return ModelWithInvalidQuerySet(self.model)
482
483
484class ValidObjectsManager(ModelWithInvalidManager):
jadmanski0afbb632008-06-06 21:10:57 +0000485 """
486 Manager returning only objects with invalid=False.
487 """
488 def get_query_set(self):
489 queryset = super(ValidObjectsManager, self).get_query_set()
490 return queryset.filter(invalid=False)
showard7c785282008-05-29 19:45:12 +0000491
492
Prashanth B489b91d2014-03-15 12:17:16 -0700493class ModelExtensions(rdb_model_extensions.ModelValidators):
jadmanski0afbb632008-06-06 21:10:57 +0000494 """\
Prashanth B489b91d2014-03-15 12:17:16 -0700495 Mixin with convenience functions for models, built on top of
496 the model validators in rdb_model_extensions.
jadmanski0afbb632008-06-06 21:10:57 +0000497 """
498 # TODO: at least some of these functions really belong in a custom
499 # Manager class
showard7c785282008-05-29 19:45:12 +0000500
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700501
502 SERIALIZATION_LINKS_TO_FOLLOW = set()
503 """
504 To be able to send jobs and hosts to shards, it's necessary to find their
505 dependencies.
506 The most generic approach for this would be to traverse all relationships
507 to other objects recursively. This would list all objects that are related
508 in any way.
509 But this approach finds too many objects: If a host should be transferred,
510 all it's relationships would be traversed. This would find an acl group.
511 If then the acl group's relationships are traversed, the relationship
512 would be followed backwards and many other hosts would be found.
513
514 This mapping tells that algorithm which relations to follow explicitly.
515 """
516
Jakob Juelichf865d332014-09-29 10:47:49 -0700517
Fang Deng86248502014-12-18 16:38:00 -0800518 SERIALIZATION_LINKS_TO_KEEP = set()
519 """This set stores foreign keys which we don't want to follow, but
520 still want to include in the serialized dictionary. For
521 example, we follow the relationship `Host.hostattribute_set`,
522 but we do not want to follow `HostAttributes.host_id` back to
523 to Host, which would otherwise lead to a circle. However, we still
524 like to serialize HostAttribute.`host_id`."""
525
Jakob Juelichf865d332014-09-29 10:47:49 -0700526 SERIALIZATION_LOCAL_LINKS_TO_UPDATE = set()
527 """
528 On deserializion, if the object to persist already exists, local fields
529 will only be updated, if their name is in this set.
530 """
531
532
jadmanski0afbb632008-06-06 21:10:57 +0000533 @classmethod
534 def convert_human_readable_values(cls, data, to_human_readable=False):
535 """\
536 Performs conversions on user-supplied field data, to make it
537 easier for users to pass human-readable data.
showard7c785282008-05-29 19:45:12 +0000538
jadmanski0afbb632008-06-06 21:10:57 +0000539 For all fields that have choice sets, convert their values
540 from human-readable strings to enum values, if necessary. This
541 allows users to pass strings instead of the corresponding
542 integer values.
showard7c785282008-05-29 19:45:12 +0000543
jadmanski0afbb632008-06-06 21:10:57 +0000544 For all foreign key fields, call smart_get with the supplied
545 data. This allows the user to pass either an ID value or
546 the name of the object as a string.
showard7c785282008-05-29 19:45:12 +0000547
jadmanski0afbb632008-06-06 21:10:57 +0000548 If to_human_readable=True, perform the inverse - i.e. convert
549 numeric values to human readable values.
showard7c785282008-05-29 19:45:12 +0000550
jadmanski0afbb632008-06-06 21:10:57 +0000551 This method modifies data in-place.
552 """
553 field_dict = cls.get_field_dict()
554 for field_name in data:
showarde732ee72008-09-23 19:15:43 +0000555 if field_name not in field_dict or data[field_name] is None:
jadmanski0afbb632008-06-06 21:10:57 +0000556 continue
557 field_obj = field_dict[field_name]
558 # convert enum values
559 if field_obj.choices:
560 for choice_data in field_obj.choices:
561 # choice_data is (value, name)
562 if to_human_readable:
563 from_val, to_val = choice_data
564 else:
565 to_val, from_val = choice_data
566 if from_val == data[field_name]:
567 data[field_name] = to_val
568 break
569 # convert foreign key values
570 elif field_obj.rel:
showarda4ea5742009-02-17 20:56:23 +0000571 dest_obj = field_obj.rel.to.smart_get(data[field_name],
572 valid_only=False)
showardf8b19042009-05-12 17:22:49 +0000573 if to_human_readable:
Paul Pendlebury5a8c6ad2011-02-01 07:20:17 -0800574 # parameterized_jobs do not have a name_field
575 if (field_name != 'parameterized_job' and
576 dest_obj.name_field is not None):
showardf8b19042009-05-12 17:22:49 +0000577 data[field_name] = getattr(dest_obj,
578 dest_obj.name_field)
jadmanski0afbb632008-06-06 21:10:57 +0000579 else:
showardb0a73032009-03-27 18:35:41 +0000580 data[field_name] = dest_obj
showard7c785282008-05-29 19:45:12 +0000581
582
showard7c785282008-05-29 19:45:12 +0000583
584
Dale Curtis74a314b2011-06-23 14:55:46 -0700585 def _validate_unique(self):
jadmanski0afbb632008-06-06 21:10:57 +0000586 """\
587 Validate that unique fields are unique. Django manipulators do
588 this too, but they're a huge pain to use manually. Trust me.
589 """
590 errors = {}
591 cls = type(self)
592 field_dict = self.get_field_dict()
593 manager = cls.get_valid_manager()
594 for field_name, field_obj in field_dict.iteritems():
595 if not field_obj.unique:
596 continue
showard7c785282008-05-29 19:45:12 +0000597
jadmanski0afbb632008-06-06 21:10:57 +0000598 value = getattr(self, field_name)
showardbd18ab72009-09-18 21:20:27 +0000599 if value is None and field_obj.auto_created:
600 # don't bother checking autoincrement fields about to be
601 # generated
602 continue
603
jadmanski0afbb632008-06-06 21:10:57 +0000604 existing_objs = manager.filter(**{field_name : value})
605 num_existing = existing_objs.count()
showard7c785282008-05-29 19:45:12 +0000606
jadmanski0afbb632008-06-06 21:10:57 +0000607 if num_existing == 0:
608 continue
609 if num_existing == 1 and existing_objs[0].id == self.id:
610 continue
611 errors[field_name] = (
612 'This value must be unique (%s)' % (value))
613 return errors
showard7c785282008-05-29 19:45:12 +0000614
615
showarda5288b42009-07-28 20:06:08 +0000616 def _validate(self):
617 """
618 First coerces all fields on this instance to their proper Python types.
619 Then runs validation on every field. Returns a dictionary of
620 field_name -> error_list.
621
622 Based on validate() from django.db.models.Model in Django 0.96, which
623 was removed in Django 1.0. It should reappear in a later version. See:
624 http://code.djangoproject.com/ticket/6845
625 """
626 error_dict = {}
627 for f in self._meta.fields:
628 try:
629 python_value = f.to_python(
630 getattr(self, f.attname, f.get_default()))
631 except django.core.exceptions.ValidationError, e:
jamesren1e0a4ce2010-04-21 17:45:11 +0000632 error_dict[f.name] = str(e)
showarda5288b42009-07-28 20:06:08 +0000633 continue
634
635 if not f.blank and not python_value:
636 error_dict[f.name] = 'This field is required.'
637 continue
638
639 setattr(self, f.attname, python_value)
640
641 return error_dict
642
643
jadmanski0afbb632008-06-06 21:10:57 +0000644 def do_validate(self):
showarda5288b42009-07-28 20:06:08 +0000645 errors = self._validate()
Dale Curtis74a314b2011-06-23 14:55:46 -0700646 unique_errors = self._validate_unique()
jadmanski0afbb632008-06-06 21:10:57 +0000647 for field_name, error in unique_errors.iteritems():
648 errors.setdefault(field_name, error)
649 if errors:
650 raise ValidationError(errors)
showard7c785282008-05-29 19:45:12 +0000651
652
jadmanski0afbb632008-06-06 21:10:57 +0000653 # actually (externally) useful methods follow
showard7c785282008-05-29 19:45:12 +0000654
jadmanski0afbb632008-06-06 21:10:57 +0000655 @classmethod
656 def add_object(cls, data={}, **kwargs):
657 """\
658 Returns a new object created with the given data (a dictionary
659 mapping field names to values). Merges any extra keyword args
660 into data.
661 """
Prashanth B489b91d2014-03-15 12:17:16 -0700662 data = dict(data)
663 data.update(kwargs)
664 data = cls.prepare_data_args(data)
665 cls.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000666 data = cls.provide_default_values(data)
Prashanth B489b91d2014-03-15 12:17:16 -0700667
jadmanski0afbb632008-06-06 21:10:57 +0000668 obj = cls(**data)
669 obj.do_validate()
670 obj.save()
671 return obj
showard7c785282008-05-29 19:45:12 +0000672
673
jadmanski0afbb632008-06-06 21:10:57 +0000674 def update_object(self, data={}, **kwargs):
675 """\
676 Updates the object with the given data (a dictionary mapping
677 field names to values). Merges any extra keyword args into
678 data.
679 """
Prashanth B489b91d2014-03-15 12:17:16 -0700680 data = dict(data)
681 data.update(kwargs)
682 data = self.prepare_data_args(data)
683 self.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000684 for field_name, value in data.iteritems():
showardb0a73032009-03-27 18:35:41 +0000685 setattr(self, field_name, value)
jadmanski0afbb632008-06-06 21:10:57 +0000686 self.do_validate()
687 self.save()
showard7c785282008-05-29 19:45:12 +0000688
689
showard8bfb5cb2009-10-07 20:49:15 +0000690 # see query_objects()
691 _SPECIAL_FILTER_KEYS = ('query_start', 'query_limit', 'sort_by',
692 'extra_args', 'extra_where', 'no_distinct')
693
694
jadmanski0afbb632008-06-06 21:10:57 +0000695 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000696 def _extract_special_params(cls, filter_data):
697 """
698 @returns a tuple of dicts (special_params, regular_filters), where
699 special_params contains the parameters we handle specially and
700 regular_filters is the remaining data to be handled by Django.
701 """
702 regular_filters = dict(filter_data)
703 special_params = {}
704 for key in cls._SPECIAL_FILTER_KEYS:
705 if key in regular_filters:
706 special_params[key] = regular_filters.pop(key)
707 return special_params, regular_filters
708
709
710 @classmethod
711 def apply_presentation(cls, query, filter_data):
712 """
713 Apply presentation parameters -- sorting and paging -- to the given
714 query.
715 @returns new query with presentation applied
716 """
717 special_params, _ = cls._extract_special_params(filter_data)
718 sort_by = special_params.get('sort_by', None)
719 if sort_by:
720 assert isinstance(sort_by, list) or isinstance(sort_by, tuple)
showard8b0ea222009-12-23 19:23:03 +0000721 query = query.extra(order_by=sort_by)
showard8bfb5cb2009-10-07 20:49:15 +0000722
723 query_start = special_params.get('query_start', None)
724 query_limit = special_params.get('query_limit', None)
725 if query_start is not None:
726 if query_limit is None:
727 raise ValueError('Cannot pass query_start without query_limit')
728 # query_limit is passed as a page size
showard7074b742009-10-12 20:30:04 +0000729 query_limit += query_start
730 return query[query_start:query_limit]
showard8bfb5cb2009-10-07 20:49:15 +0000731
732
733 @classmethod
734 def query_objects(cls, filter_data, valid_only=True, initial_query=None,
735 apply_presentation=True):
jadmanski0afbb632008-06-06 21:10:57 +0000736 """\
737 Returns a QuerySet object for querying the given model_class
738 with the given filter_data. Optional special arguments in
739 filter_data include:
740 -query_start: index of first return to return
741 -query_limit: maximum number of results to return
742 -sort_by: list of fields to sort on. prefixing a '-' onto a
743 field name changes the sort to descending order.
744 -extra_args: keyword args to pass to query.extra() (see Django
745 DB layer documentation)
showarda5288b42009-07-28 20:06:08 +0000746 -extra_where: extra WHERE clause to append
showard8bfb5cb2009-10-07 20:49:15 +0000747 -no_distinct: if True, a DISTINCT will not be added to the SELECT
jadmanski0afbb632008-06-06 21:10:57 +0000748 """
showard8bfb5cb2009-10-07 20:49:15 +0000749 special_params, regular_filters = cls._extract_special_params(
750 filter_data)
showard7c785282008-05-29 19:45:12 +0000751
showard7ac7b7a2008-07-21 20:24:29 +0000752 if initial_query is None:
753 if valid_only:
754 initial_query = cls.get_valid_manager()
755 else:
756 initial_query = cls.objects
showard8bfb5cb2009-10-07 20:49:15 +0000757
758 query = initial_query.filter(**regular_filters)
759
760 use_distinct = not special_params.get('no_distinct', False)
showard7ac7b7a2008-07-21 20:24:29 +0000761 if use_distinct:
762 query = query.distinct()
showard7c785282008-05-29 19:45:12 +0000763
showard8bfb5cb2009-10-07 20:49:15 +0000764 extra_args = special_params.get('extra_args', {})
765 extra_where = special_params.get('extra_where', None)
766 if extra_where:
767 # escape %'s
768 extra_where = cls.objects.escape_user_sql(extra_where)
769 extra_args.setdefault('where', []).append(extra_where)
jadmanski0afbb632008-06-06 21:10:57 +0000770 if extra_args:
771 query = query.extra(**extra_args)
Jakob Juelich7bef8412014-10-14 19:11:54 -0700772 # TODO: Use readonly connection for these queries.
773 # This has been disabled, because it's not used anyway, as the
774 # configured readonly user is the same as the real user anyway.
showard7c785282008-05-29 19:45:12 +0000775
showard8bfb5cb2009-10-07 20:49:15 +0000776 if apply_presentation:
777 query = cls.apply_presentation(query, filter_data)
778
779 return query
showard7c785282008-05-29 19:45:12 +0000780
781
jadmanski0afbb632008-06-06 21:10:57 +0000782 @classmethod
showard585c2ab2008-07-23 19:29:49 +0000783 def query_count(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000784 """\
785 Like query_objects, but retreive only the count of results.
786 """
787 filter_data.pop('query_start', None)
788 filter_data.pop('query_limit', None)
showard585c2ab2008-07-23 19:29:49 +0000789 query = cls.query_objects(filter_data, initial_query=initial_query)
790 return query.count()
showard7c785282008-05-29 19:45:12 +0000791
792
jadmanski0afbb632008-06-06 21:10:57 +0000793 @classmethod
794 def clean_object_dicts(cls, field_dicts):
795 """\
796 Take a list of dicts corresponding to object (as returned by
797 query.values()) and clean the data to be more suitable for
798 returning to the user.
799 """
showarde732ee72008-09-23 19:15:43 +0000800 for field_dict in field_dicts:
801 cls.clean_foreign_keys(field_dict)
showard21baa452008-10-21 00:08:39 +0000802 cls._convert_booleans(field_dict)
showarde732ee72008-09-23 19:15:43 +0000803 cls.convert_human_readable_values(field_dict,
804 to_human_readable=True)
showard7c785282008-05-29 19:45:12 +0000805
806
jadmanski0afbb632008-06-06 21:10:57 +0000807 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000808 def list_objects(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000809 """\
810 Like query_objects, but return a list of dictionaries.
811 """
showard7ac7b7a2008-07-21 20:24:29 +0000812 query = cls.query_objects(filter_data, initial_query=initial_query)
showard8bfb5cb2009-10-07 20:49:15 +0000813 extra_fields = query.query.extra_select.keys()
814 field_dicts = [model_object.get_object_dict(extra_fields=extra_fields)
showarde732ee72008-09-23 19:15:43 +0000815 for model_object in query]
jadmanski0afbb632008-06-06 21:10:57 +0000816 return field_dicts
showard7c785282008-05-29 19:45:12 +0000817
818
jadmanski0afbb632008-06-06 21:10:57 +0000819 @classmethod
showarda4ea5742009-02-17 20:56:23 +0000820 def smart_get(cls, id_or_name, valid_only=True):
jadmanski0afbb632008-06-06 21:10:57 +0000821 """\
822 smart_get(integer) -> get object by ID
823 smart_get(string) -> get object by name_field
jadmanski0afbb632008-06-06 21:10:57 +0000824 """
showarda4ea5742009-02-17 20:56:23 +0000825 if valid_only:
826 manager = cls.get_valid_manager()
827 else:
828 manager = cls.objects
829
830 if isinstance(id_or_name, (int, long)):
831 return manager.get(pk=id_or_name)
jamesren3e9f6092010-03-11 21:32:10 +0000832 if isinstance(id_or_name, basestring) and hasattr(cls, 'name_field'):
showarda4ea5742009-02-17 20:56:23 +0000833 return manager.get(**{cls.name_field : id_or_name})
834 raise ValueError(
835 'Invalid positional argument: %s (%s)' % (id_or_name,
836 type(id_or_name)))
showard7c785282008-05-29 19:45:12 +0000837
838
showardbe3ec042008-11-12 18:16:07 +0000839 @classmethod
840 def smart_get_bulk(cls, id_or_name_list):
841 invalid_inputs = []
842 result_objects = []
843 for id_or_name in id_or_name_list:
844 try:
845 result_objects.append(cls.smart_get(id_or_name))
846 except cls.DoesNotExist:
847 invalid_inputs.append(id_or_name)
848 if invalid_inputs:
mbligh7a3ebe32008-12-01 17:10:33 +0000849 raise cls.DoesNotExist('The following %ss do not exist: %s'
850 % (cls.__name__.lower(),
851 ', '.join(invalid_inputs)))
showardbe3ec042008-11-12 18:16:07 +0000852 return result_objects
853
854
showard8bfb5cb2009-10-07 20:49:15 +0000855 def get_object_dict(self, extra_fields=None):
jadmanski0afbb632008-06-06 21:10:57 +0000856 """\
showard8bfb5cb2009-10-07 20:49:15 +0000857 Return a dictionary mapping fields to this object's values. @param
858 extra_fields: list of extra attribute names to include, in addition to
859 the fields defined on this object.
jadmanski0afbb632008-06-06 21:10:57 +0000860 """
showard8bfb5cb2009-10-07 20:49:15 +0000861 fields = self.get_field_dict().keys()
862 if extra_fields:
863 fields += extra_fields
jadmanski0afbb632008-06-06 21:10:57 +0000864 object_dict = dict((field_name, getattr(self, field_name))
showarde732ee72008-09-23 19:15:43 +0000865 for field_name in fields)
jadmanski0afbb632008-06-06 21:10:57 +0000866 self.clean_object_dicts([object_dict])
showardd3dc1992009-04-22 21:01:40 +0000867 self._postprocess_object_dict(object_dict)
jadmanski0afbb632008-06-06 21:10:57 +0000868 return object_dict
showard7c785282008-05-29 19:45:12 +0000869
870
showardd3dc1992009-04-22 21:01:40 +0000871 def _postprocess_object_dict(self, object_dict):
872 """For subclasses to override."""
873 pass
874
875
jadmanski0afbb632008-06-06 21:10:57 +0000876 @classmethod
877 def get_valid_manager(cls):
878 return cls.objects
showard7c785282008-05-29 19:45:12 +0000879
880
showard2bab8f42008-11-12 18:15:22 +0000881 def _record_attributes(self, attributes):
882 """
883 See on_attribute_changed.
884 """
885 assert not isinstance(attributes, basestring)
886 self._recorded_attributes = dict((attribute, getattr(self, attribute))
887 for attribute in attributes)
888
889
890 def _check_for_updated_attributes(self):
891 """
892 See on_attribute_changed.
893 """
894 for attribute, original_value in self._recorded_attributes.iteritems():
895 new_value = getattr(self, attribute)
896 if original_value != new_value:
897 self.on_attribute_changed(attribute, original_value)
898 self._record_attributes(self._recorded_attributes.keys())
899
900
901 def on_attribute_changed(self, attribute, old_value):
902 """
903 Called whenever an attribute is updated. To be overridden.
904
905 To use this method, you must:
906 * call _record_attributes() from __init__() (after making the super
907 call) with a list of attributes for which you want to be notified upon
908 change.
909 * call _check_for_updated_attributes() from save().
910 """
911 pass
912
913
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700914 def serialize(self, include_dependencies=True):
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700915 """Serializes the object with dependencies.
916
917 The variable SERIALIZATION_LINKS_TO_FOLLOW defines which dependencies
918 this function will serialize with the object.
919
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700920 @param include_dependencies: Whether or not to follow relations to
921 objects this object depends on.
922 This parameter is used when uploading
923 jobs from a shard to the master, as the
924 master already has all the dependent
925 objects.
926
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700927 @returns: Dictionary representation of the object.
928 """
929 serialized = {}
MK Ryu5cfd96a2015-01-30 15:31:23 -0800930 timer = autotest_stats.Timer('serialize_latency.%s' % (
931 type(self).__name__))
932 with timer.get_client('local'):
933 for field in self._meta.concrete_model._meta.local_fields:
934 if field.rel is None:
935 serialized[field.name] = field._get_val_from_obj(self)
MK Ryudf4d4232015-03-06 11:18:47 -0800936 elif field.name in self.SERIALIZATION_LINKS_TO_KEEP:
MK Ryu5cfd96a2015-01-30 15:31:23 -0800937 # attname will contain "_id" suffix for foreign keys,
938 # e.g. HostAttribute.host will be serialized as 'host_id'.
939 # Use it for easy deserialization.
940 serialized[field.attname] = field._get_val_from_obj(self)
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700941
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700942 if include_dependencies:
MK Ryu5cfd96a2015-01-30 15:31:23 -0800943 with timer.get_client('related'):
944 for link in self.SERIALIZATION_LINKS_TO_FOLLOW:
945 serialized[link] = self._serialize_relation(link)
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700946
947 return serialized
948
949
950 def _serialize_relation(self, link):
951 """Serializes dependent objects given the name of the relation.
952
953 @param link: Name of the relation to take objects from.
954
955 @returns For To-Many relationships a list of the serialized related
956 objects, for To-One relationships the serialized related object.
957 """
958 try:
959 attr = getattr(self, link)
960 except AttributeError:
961 # One-To-One relationships that point to None may raise this
962 return None
963
964 if attr is None:
965 return None
966 if hasattr(attr, 'all'):
967 return [obj.serialize() for obj in attr.all()]
968 return attr.serialize()
969
970
Jakob Juelichf88fa932014-09-03 17:58:04 -0700971 @classmethod
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700972 def _split_local_from_foreign_values(cls, data):
973 """This splits local from foreign values in a serialized object.
974
975 @param data: The serialized object.
976
977 @returns A tuple of two lists, both containing tuples in the form
978 (link_name, link_value). The first list contains all links
979 for local fields, the second one contains those for foreign
980 fields/objects.
981 """
982 links_to_local_values, links_to_related_values = [], []
983 for link, value in data.iteritems():
984 if link in cls.SERIALIZATION_LINKS_TO_FOLLOW:
985 # It's a foreign key
986 links_to_related_values.append((link, value))
987 else:
Fang Deng86248502014-12-18 16:38:00 -0800988 # It's a local attribute or a foreign key
989 # we don't want to follow.
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700990 links_to_local_values.append((link, value))
991 return links_to_local_values, links_to_related_values
992
993
Jakob Juelichf865d332014-09-29 10:47:49 -0700994 @classmethod
995 def _filter_update_allowed_fields(cls, data):
996 """Filters data and returns only files that updates are allowed on.
997
998 This is i.e. needed for syncing aborted bits from the master to shards.
999
1000 Local links are only allowed to be updated, if they are in
1001 SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1002 Overwriting existing values is allowed in order to be able to sync i.e.
1003 the aborted bit from the master to a shard.
1004
1005 The whitelisting mechanism is in place to prevent overwriting local
1006 status: If all fields were overwritten, jobs would be completely be
1007 set back to their original (unstarted) state.
1008
1009 @param data: List with tuples of the form (link_name, link_value), as
1010 returned by _split_local_from_foreign_values.
1011
1012 @returns List of the same format as data, but only containing data for
1013 fields that updates are allowed on.
1014 """
1015 return [pair for pair in data
1016 if pair[0] in cls.SERIALIZATION_LOCAL_LINKS_TO_UPDATE]
1017
1018
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001019 @classmethod
1020 def delete_matching_record(cls, **filter_args):
1021 """Delete records matching the filter.
1022
1023 @param filter_args: Arguments for the django filter
1024 used to locate the record to delete.
1025 """
1026 try:
1027 existing_record = cls.objects.get(**filter_args)
1028 except cls.DoesNotExist:
1029 return
1030 existing_record.delete()
1031
1032
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001033 def _deserialize_local(self, data):
1034 """Set local attributes from a list of tuples.
1035
1036 @param data: List of tuples like returned by
1037 _split_local_from_foreign_values.
1038 """
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001039 if not data:
1040 return
1041
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001042 for link, value in data:
1043 setattr(self, link, value)
1044 # Overwridden save() methods are prone to errors, so don't execute them.
1045 # This is because:
1046 # - the overwritten methods depend on ACL groups that don't yet exist
1047 # and don't handle errors
1048 # - the overwritten methods think this object already exists in the db
1049 # because the id is already set
1050 super(type(self), self).save()
1051
1052
1053 def _deserialize_relations(self, data):
1054 """Set foreign attributes from a list of tuples.
1055
1056 This deserialized the related objects using their own deserialize()
1057 function and then sets the relation.
1058
1059 @param data: List of tuples like returned by
1060 _split_local_from_foreign_values.
1061 """
1062 for link, value in data:
1063 self._deserialize_relation(link, value)
1064 # See comment in _deserialize_local
1065 super(type(self), self).save()
1066
1067
1068 @classmethod
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001069 def get_record(cls, data):
1070 """Retrieve a record with the data in the given input arg.
1071
1072 @param data: A dictionary containing the information to use in a query
1073 for data. If child models have different constraints of
1074 uniqueness they should override this model.
1075
1076 @return: An object with matching data.
1077
1078 @raises DoesNotExist: If a record with the given data doesn't exist.
1079 """
1080 return cls.objects.get(id=data['id'])
1081
1082
1083 @classmethod
Jakob Juelichf88fa932014-09-03 17:58:04 -07001084 def deserialize(cls, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001085 """Recursively deserializes and saves an object with it's dependencies.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001086
1087 This takes the result of the serialize method and creates objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001088 in the database that are just like the original.
1089
1090 If an object of the same type with the same id already exists, it's
Jakob Juelichf865d332014-09-29 10:47:49 -07001091 local values will be left untouched, unless they are explicitly
1092 whitelisted in SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1093
1094 Deserialize will always recursively propagate to all related objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001095 present in data though.
1096 I.e. this is necessary to add users to an already existing acl-group.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001097
1098 @param data: Representation of an object and its dependencies, as
1099 returned by serialize.
1100
1101 @returns: The object represented by data if it didn't exist before,
1102 otherwise the object that existed before and has the same type
1103 and id as the one described by data.
1104 """
1105 if data is None:
1106 return None
1107
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001108 local, related = cls._split_local_from_foreign_values(data)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001109 try:
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001110 instance = cls.get_record(data)
Jakob Juelichf865d332014-09-29 10:47:49 -07001111 local = cls._filter_update_allowed_fields(local)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001112 except cls.DoesNotExist:
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001113 instance = cls()
Jakob Juelichf88fa932014-09-03 17:58:04 -07001114
MK Ryu5cfd96a2015-01-30 15:31:23 -08001115 timer = autotest_stats.Timer('deserialize_latency.%s' % (
1116 type(instance).__name__))
1117 with timer.get_client('local'):
1118 instance._deserialize_local(local)
1119 with timer.get_client('related'):
1120 instance._deserialize_relations(related)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001121
1122 return instance
1123
1124
Jakob Juelicha94efe62014-09-18 16:02:49 -07001125 def sanity_check_update_from_shard(self, shard, updated_serialized,
1126 *args, **kwargs):
1127 """Check if an update sent from a shard is legitimate.
1128
1129 @raises error.UnallowedRecordsSentToMaster if an update is not
1130 legitimate.
1131 """
1132 raise NotImplementedError(
1133 'sanity_check_update_from_shard must be implemented by subclass %s '
1134 'for type %s' % type(self))
1135
1136
Prashanth Balasubramanian75be1d32014-11-25 18:03:09 -08001137 @transaction.commit_on_success
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001138 def update_from_serialized(self, serialized):
1139 """Updates local fields of an existing object from a serialized form.
1140
1141 This is different than the normal deserialize() in the way that it
1142 does update local values, which deserialize doesn't, but doesn't
1143 recursively propagate to related objects, which deserialize() does.
1144
1145 The use case of this function is to update job records on the master
1146 after the jobs have been executed on a slave, as the master is not
1147 interested in updates for users, labels, specialtasks, etc.
1148
1149 @param serialized: Representation of an object and its dependencies, as
1150 returned by serialize.
1151
1152 @raises ValueError: if serialized contains related objects, i.e. not
1153 only local fields.
1154 """
1155 local, related = (
1156 self._split_local_from_foreign_values(serialized))
1157 if related:
1158 raise ValueError('Serialized must not contain foreign '
1159 'objects: %s' % related)
1160
1161 self._deserialize_local(local)
1162
1163
Jakob Juelichf88fa932014-09-03 17:58:04 -07001164 def custom_deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001165 """Allows overriding the deserialization behaviour by subclasses."""
Jakob Juelichf88fa932014-09-03 17:58:04 -07001166 raise NotImplementedError(
1167 'custom_deserialize_relation must be implemented by subclass %s '
1168 'for relation %s' % (type(self), link))
1169
1170
1171 def _deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001172 """Deserializes related objects and sets references on this object.
1173
1174 Relations that point to a list of objects are handled automatically.
1175 For many-to-one or one-to-one relations custom_deserialize_relation
1176 must be overridden by the subclass.
1177
1178 Related objects are deserialized using their deserialize() method.
1179 Thereby they and their dependencies are created if they don't exist
1180 and saved to the database.
1181
1182 @param link: Name of the relation.
1183 @param data: Serialized representation of the related object(s).
1184 This means a list of dictionaries for to-many relations,
1185 just a dictionary for to-one relations.
1186 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001187 field = getattr(self, link)
1188
1189 if field and hasattr(field, 'all'):
1190 self._deserialize_2m_relation(link, data, field.model)
1191 else:
1192 self.custom_deserialize_relation(link, data)
1193
1194
1195 def _deserialize_2m_relation(self, link, data, related_class):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001196 """Deserialize related objects for one to-many relationship.
1197
1198 @param link: Name of the relation.
1199 @param data: Serialized representation of the related objects.
1200 This is a list with of dictionaries.
Fang Dengff361592015-02-02 15:27:34 -08001201 @param related_class: A class representing a django model, with which
1202 this class has a one-to-many relationship.
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001203 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001204 relation_set = getattr(self, link)
Fang Dengff361592015-02-02 15:27:34 -08001205 if related_class == self.get_attribute_model():
1206 # When deserializing a model together with
1207 # its attributes, clear all the exising attributes to ensure
1208 # db consistency. Note 'update' won't be sufficient, as we also
1209 # want to remove any attributes that no longer exist in |data|.
1210 #
1211 # core_filters is a dictionary of filters, defines how
1212 # RelatedMangager would query for the 1-to-many relationship. E.g.
1213 # Host.objects.get(
1214 # id=20).hostattribute_set.core_filters = {host_id:20}
1215 # We use it to delete objects related to the current object.
1216 related_class.objects.filter(**relation_set.core_filters).delete()
Jakob Juelichf88fa932014-09-03 17:58:04 -07001217 for serialized in data:
1218 relation_set.add(related_class.deserialize(serialized))
1219
1220
Fang Dengff361592015-02-02 15:27:34 -08001221 @classmethod
1222 def get_attribute_model(cls):
1223 """Return the attribute model.
1224
1225 Subclass with attribute-like model should override this to
1226 return the attribute model class. This method will be
1227 called by _deserialize_2m_relation to determine whether
1228 to clear the one-to-many relations first on deserialization of object.
1229 """
1230 return None
1231
1232
showard7c785282008-05-29 19:45:12 +00001233class ModelWithInvalid(ModelExtensions):
jadmanski0afbb632008-06-06 21:10:57 +00001234 """
1235 Overrides model methods save() and delete() to support invalidation in
1236 place of actual deletion. Subclasses must have a boolean "invalid"
1237 field.
1238 """
showard7c785282008-05-29 19:45:12 +00001239
showarda5288b42009-07-28 20:06:08 +00001240 def save(self, *args, **kwargs):
showardddb90992009-02-11 23:39:32 +00001241 first_time = (self.id is None)
1242 if first_time:
1243 # see if this object was previously added and invalidated
1244 my_name = getattr(self, self.name_field)
1245 filters = {self.name_field : my_name, 'invalid' : True}
1246 try:
1247 old_object = self.__class__.objects.get(**filters)
showardafd97de2009-10-01 18:45:09 +00001248 self.resurrect_object(old_object)
showardddb90992009-02-11 23:39:32 +00001249 except self.DoesNotExist:
1250 # no existing object
1251 pass
showard7c785282008-05-29 19:45:12 +00001252
showarda5288b42009-07-28 20:06:08 +00001253 super(ModelWithInvalid, self).save(*args, **kwargs)
showard7c785282008-05-29 19:45:12 +00001254
1255
showardafd97de2009-10-01 18:45:09 +00001256 def resurrect_object(self, old_object):
1257 """
1258 Called when self is about to be saved for the first time and is actually
1259 "undeleting" a previously deleted object. Can be overridden by
1260 subclasses to copy data as desired from the deleted entry (but this
1261 superclass implementation must normally be called).
1262 """
1263 self.id = old_object.id
1264
1265
jadmanski0afbb632008-06-06 21:10:57 +00001266 def clean_object(self):
1267 """
1268 This method is called when an object is marked invalid.
1269 Subclasses should override this to clean up relationships that
showardafd97de2009-10-01 18:45:09 +00001270 should no longer exist if the object were deleted.
1271 """
jadmanski0afbb632008-06-06 21:10:57 +00001272 pass
showard7c785282008-05-29 19:45:12 +00001273
1274
jadmanski0afbb632008-06-06 21:10:57 +00001275 def delete(self):
Dale Curtis74a314b2011-06-23 14:55:46 -07001276 self.invalid = self.invalid
jadmanski0afbb632008-06-06 21:10:57 +00001277 assert not self.invalid
1278 self.invalid = True
1279 self.save()
1280 self.clean_object()
showard7c785282008-05-29 19:45:12 +00001281
1282
jadmanski0afbb632008-06-06 21:10:57 +00001283 @classmethod
1284 def get_valid_manager(cls):
1285 return cls.valid_objects
showard7c785282008-05-29 19:45:12 +00001286
1287
jadmanski0afbb632008-06-06 21:10:57 +00001288 class Manipulator(object):
1289 """
1290 Force default manipulators to look only at valid objects -
1291 otherwise they will match against invalid objects when checking
1292 uniqueness.
1293 """
1294 @classmethod
1295 def _prepare(cls, model):
1296 super(ModelWithInvalid.Manipulator, cls)._prepare(model)
1297 cls.manager = model.valid_objects
showardf8b19042009-05-12 17:22:49 +00001298
1299
1300class ModelWithAttributes(object):
1301 """
1302 Mixin class for models that have an attribute model associated with them.
1303 The attribute model is assumed to have its value field named "value".
1304 """
1305
1306 def _get_attribute_model_and_args(self, attribute):
1307 """
1308 Subclasses should override this to return a tuple (attribute_model,
1309 keyword_args), where attribute_model is a model class and keyword_args
1310 is a dict of args to pass to attribute_model.objects.get() to get an
1311 instance of the given attribute on this object.
1312 """
Dale Curtis74a314b2011-06-23 14:55:46 -07001313 raise NotImplementedError
showardf8b19042009-05-12 17:22:49 +00001314
1315
1316 def set_attribute(self, attribute, value):
1317 attribute_model, get_args = self._get_attribute_model_and_args(
1318 attribute)
1319 attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
1320 attribute_object.value = value
1321 attribute_object.save()
1322
1323
1324 def delete_attribute(self, attribute):
1325 attribute_model, get_args = self._get_attribute_model_and_args(
1326 attribute)
1327 try:
1328 attribute_model.objects.get(**get_args).delete()
showard16245422009-09-08 16:28:15 +00001329 except attribute_model.DoesNotExist:
showardf8b19042009-05-12 17:22:49 +00001330 pass
1331
1332
1333 def set_or_delete_attribute(self, attribute, value):
1334 if value is None:
1335 self.delete_attribute(attribute)
1336 else:
1337 self.set_attribute(attribute, value)
showard26b7ec72009-12-21 22:43:57 +00001338
1339
1340class ModelWithHashManager(dbmodels.Manager):
1341 """Manager for use with the ModelWithHash abstract model class"""
1342
1343 def create(self, **kwargs):
1344 raise Exception('ModelWithHash manager should use get_or_create() '
1345 'instead of create()')
1346
1347
1348 def get_or_create(self, **kwargs):
1349 kwargs['the_hash'] = self.model._compute_hash(**kwargs)
1350 return super(ModelWithHashManager, self).get_or_create(**kwargs)
1351
1352
1353class ModelWithHash(dbmodels.Model):
1354 """Superclass with methods for dealing with a hash column"""
1355
1356 the_hash = dbmodels.CharField(max_length=40, unique=True)
1357
1358 objects = ModelWithHashManager()
1359
1360 class Meta:
1361 abstract = True
1362
1363
1364 @classmethod
1365 def _compute_hash(cls, **kwargs):
1366 raise NotImplementedError('Subclasses must override _compute_hash()')
1367
1368
1369 def save(self, force_insert=False, **kwargs):
1370 """Prevents saving the model in most cases
1371
1372 We want these models to be immutable, so the generic save() operation
1373 will not work. These models should be instantiated through their the
1374 model.objects.get_or_create() method instead.
1375
1376 The exception is that save(force_insert=True) will be allowed, since
1377 that creates a new row. However, the preferred way to make instances of
1378 these models is through the get_or_create() method.
1379 """
1380 if not force_insert:
1381 # Allow a forced insert to happen; if it's a duplicate, the unique
1382 # constraint will catch it later anyways
1383 raise Exception('ModelWithHash is immutable')
1384 super(ModelWithHash, self).save(force_insert=force_insert, **kwargs)