blob: 84033f1c5ae5bf8681af0fa51e86b644d8e2e731 [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 backend
7from django.db import connection
8from django.db import connections
9from django.db import models as dbmodels
10from django.db import transaction
showarda5288b42009-07-28 20:06:08 +000011from django.db.models.sql import query
showard7e67b432010-01-20 01:13:04 +000012import django.db.models.sql.where
showard7c785282008-05-29 19:45:12 +000013from django.utils import datastructures
MK Ryu5cfd96a2015-01-30 15:31:23 -080014from autotest_lib.client.common_lib.cros.graphite import autotest_stats
Prashanth B489b91d2014-03-15 12:17:16 -070015from autotest_lib.frontend.afe import rdb_model_extensions
showard7c785282008-05-29 19:45:12 +000016
Prashanth B489b91d2014-03-15 12:17:16 -070017
18class ValidationError(django.core.exceptions.ValidationError):
jadmanski0afbb632008-06-06 21:10:57 +000019 """\
showarda5288b42009-07-28 20:06:08 +000020 Data validation error in adding or updating an object. The associated
jadmanski0afbb632008-06-06 21:10:57 +000021 value is a dictionary mapping field names to error strings.
22 """
showard7c785282008-05-29 19:45:12 +000023
showarda5288b42009-07-28 20:06:08 +000024def _quote_name(name):
25 """Shorthand for connection.ops.quote_name()."""
26 return connection.ops.quote_name(name)
27
28
beepscc9fc702013-12-02 12:45:38 -080029class LeasedHostManager(dbmodels.Manager):
30 """Query manager for unleased, unlocked hosts.
31 """
32 def get_query_set(self):
33 return (super(LeasedHostManager, self).get_query_set().filter(
34 leased=0, locked=0))
35
36
showard7c785282008-05-29 19:45:12 +000037class ExtendedManager(dbmodels.Manager):
jadmanski0afbb632008-06-06 21:10:57 +000038 """\
39 Extended manager supporting subquery filtering.
40 """
showard7c785282008-05-29 19:45:12 +000041
showardf828c772010-01-25 21:49:42 +000042 class CustomQuery(query.Query):
showard7e67b432010-01-20 01:13:04 +000043 def __init__(self, *args, **kwargs):
showardf828c772010-01-25 21:49:42 +000044 super(ExtendedManager.CustomQuery, self).__init__(*args, **kwargs)
showard7e67b432010-01-20 01:13:04 +000045 self._custom_joins = []
46
47
showarda5288b42009-07-28 20:06:08 +000048 def clone(self, klass=None, **kwargs):
showardf828c772010-01-25 21:49:42 +000049 obj = super(ExtendedManager.CustomQuery, self).clone(klass)
showard7e67b432010-01-20 01:13:04 +000050 obj._custom_joins = list(self._custom_joins)
showarda5288b42009-07-28 20:06:08 +000051 return obj
showard08f981b2008-06-24 21:59:03 +000052
showard7e67b432010-01-20 01:13:04 +000053
54 def combine(self, rhs, connector):
showardf828c772010-01-25 21:49:42 +000055 super(ExtendedManager.CustomQuery, self).combine(rhs, connector)
showard7e67b432010-01-20 01:13:04 +000056 if hasattr(rhs, '_custom_joins'):
57 self._custom_joins.extend(rhs._custom_joins)
58
59
60 def add_custom_join(self, table, condition, join_type,
61 condition_values=(), alias=None):
62 if alias is None:
63 alias = table
64 join_dict = dict(table=table,
65 condition=condition,
66 condition_values=condition_values,
67 join_type=join_type,
68 alias=alias)
69 self._custom_joins.append(join_dict)
70
71
showard7e67b432010-01-20 01:13:04 +000072 @classmethod
73 def convert_query(self, query_set):
74 """
showardf828c772010-01-25 21:49:42 +000075 Convert the query set's "query" attribute to a CustomQuery.
showard7e67b432010-01-20 01:13:04 +000076 """
77 # Make a copy of the query set
78 query_set = query_set.all()
79 query_set.query = query_set.query.clone(
showardf828c772010-01-25 21:49:42 +000080 klass=ExtendedManager.CustomQuery,
showard7e67b432010-01-20 01:13:04 +000081 _custom_joins=[])
82 return query_set
showard43a3d262008-11-12 18:17:05 +000083
84
showard7e67b432010-01-20 01:13:04 +000085 class _WhereClause(object):
86 """Object allowing us to inject arbitrary SQL into Django queries.
showard43a3d262008-11-12 18:17:05 +000087
showard7e67b432010-01-20 01:13:04 +000088 By using this instead of extra(where=...), we can still freely combine
89 queries with & and |.
showarda5288b42009-07-28 20:06:08 +000090 """
showard7e67b432010-01-20 01:13:04 +000091 def __init__(self, clause, values=()):
92 self._clause = clause
93 self._values = values
showarda5288b42009-07-28 20:06:08 +000094
showard7e67b432010-01-20 01:13:04 +000095
Dale Curtis74a314b2011-06-23 14:55:46 -070096 def as_sql(self, qn=None, connection=None):
showard7e67b432010-01-20 01:13:04 +000097 return self._clause, self._values
98
99
100 def relabel_aliases(self, change_map):
101 return
showard43a3d262008-11-12 18:17:05 +0000102
103
showard8b0ea222009-12-23 19:23:03 +0000104 def add_join(self, query_set, join_table, join_key, join_condition='',
showard7e67b432010-01-20 01:13:04 +0000105 join_condition_values=(), join_from_key=None, alias=None,
106 suffix='', exclude=False, force_left_join=False):
107 """Add a join to query_set.
108
109 Join looks like this:
110 (INNER|LEFT) JOIN <join_table> AS <alias>
111 ON (<this table>.<join_from_key> = <join_table>.<join_key>
112 and <join_condition>)
113
showard0957a842009-05-11 19:25:08 +0000114 @param join_table table to join to
115 @param join_key field referencing back to this model to use for the join
116 @param join_condition extra condition for the ON clause of the join
showard7e67b432010-01-20 01:13:04 +0000117 @param join_condition_values values to substitute into join_condition
118 @param join_from_key column on this model to join from.
showard8b0ea222009-12-23 19:23:03 +0000119 @param alias alias to use for for join
120 @param suffix suffix to add to join_table for the join alias, if no
121 alias is provided
showard0957a842009-05-11 19:25:08 +0000122 @param exclude if true, exclude rows that match this join (will use a
showarda5288b42009-07-28 20:06:08 +0000123 LEFT OUTER JOIN and an appropriate WHERE condition)
showardc4780402009-08-31 18:31:34 +0000124 @param force_left_join - if true, a LEFT OUTER JOIN will be used
125 instead of an INNER JOIN regardless of other options
showard0957a842009-05-11 19:25:08 +0000126 """
showard7e67b432010-01-20 01:13:04 +0000127 join_from_table = query_set.model._meta.db_table
128 if join_from_key is None:
129 join_from_key = self.model._meta.pk.name
130 if alias is None:
131 alias = join_table + suffix
132 full_join_key = _quote_name(alias) + '.' + _quote_name(join_key)
133 full_join_condition = '%s = %s.%s' % (full_join_key,
134 _quote_name(join_from_table),
135 _quote_name(join_from_key))
showard43a3d262008-11-12 18:17:05 +0000136 if join_condition:
137 full_join_condition += ' AND (' + join_condition + ')'
138 if exclude or force_left_join:
showarda5288b42009-07-28 20:06:08 +0000139 join_type = query_set.query.LOUTER
showard43a3d262008-11-12 18:17:05 +0000140 else:
showarda5288b42009-07-28 20:06:08 +0000141 join_type = query_set.query.INNER
showard43a3d262008-11-12 18:17:05 +0000142
showardf828c772010-01-25 21:49:42 +0000143 query_set = self.CustomQuery.convert_query(query_set)
showard7e67b432010-01-20 01:13:04 +0000144 query_set.query.add_custom_join(join_table,
145 full_join_condition,
146 join_type,
147 condition_values=join_condition_values,
148 alias=alias)
showard43a3d262008-11-12 18:17:05 +0000149
showard7e67b432010-01-20 01:13:04 +0000150 if exclude:
151 query_set = query_set.extra(where=[full_join_key + ' IS NULL'])
152
153 return query_set
154
155
156 def _info_for_many_to_one_join(self, field, join_to_query, alias):
157 """
158 @param field: the ForeignKey field on the related model
159 @param join_to_query: the query over the related model that we're
160 joining to
161 @param alias: alias of joined table
162 """
163 info = {}
164 rhs_table = join_to_query.model._meta.db_table
165 info['rhs_table'] = rhs_table
166 info['rhs_column'] = field.column
167 info['lhs_column'] = field.rel.get_related_field().column
168 rhs_where = join_to_query.query.where
169 rhs_where.relabel_aliases({rhs_table: alias})
Dale Curtis74a314b2011-06-23 14:55:46 -0700170 compiler = join_to_query.query.get_compiler(using=join_to_query.db)
171 initial_clause, values = compiler.as_sql()
172 all_clauses = (initial_clause,)
173 if hasattr(join_to_query.query, 'extra_where'):
174 all_clauses += join_to_query.query.extra_where
175 info['where_clause'] = (
176 ' AND '.join('(%s)' % clause for clause in all_clauses))
showard7e67b432010-01-20 01:13:04 +0000177 info['values'] = values
178 return info
179
180
181 def _info_for_many_to_many_join(self, m2m_field, join_to_query, alias,
182 m2m_is_on_this_model):
183 """
184 @param m2m_field: a Django field representing the M2M relationship.
185 It uses a pivot table with the following structure:
186 this model table <---> M2M pivot table <---> joined model table
187 @param join_to_query: the query over the related model that we're
188 joining to.
189 @param alias: alias of joined table
190 """
191 if m2m_is_on_this_model:
192 # referenced field on this model
193 lhs_id_field = self.model._meta.pk
194 # foreign key on the pivot table referencing lhs_id_field
195 m2m_lhs_column = m2m_field.m2m_column_name()
196 # foreign key on the pivot table referencing rhd_id_field
197 m2m_rhs_column = m2m_field.m2m_reverse_name()
198 # referenced field on related model
199 rhs_id_field = m2m_field.rel.get_related_field()
200 else:
201 lhs_id_field = m2m_field.rel.get_related_field()
202 m2m_lhs_column = m2m_field.m2m_reverse_name()
203 m2m_rhs_column = m2m_field.m2m_column_name()
204 rhs_id_field = join_to_query.model._meta.pk
205
206 info = {}
207 info['rhs_table'] = m2m_field.m2m_db_table()
208 info['rhs_column'] = m2m_lhs_column
209 info['lhs_column'] = lhs_id_field.column
210
211 # select the ID of related models relevant to this join. we can only do
212 # a single join, so we need to gather this information up front and
213 # include it in the join condition.
214 rhs_ids = join_to_query.values_list(rhs_id_field.attname, flat=True)
215 assert len(rhs_ids) == 1, ('Many-to-many custom field joins can only '
216 'match a single related object.')
217 rhs_id = rhs_ids[0]
218
219 info['where_clause'] = '%s.%s = %s' % (_quote_name(alias),
220 _quote_name(m2m_rhs_column),
221 rhs_id)
222 info['values'] = ()
223 return info
224
225
226 def join_custom_field(self, query_set, join_to_query, alias,
227 left_join=True):
228 """Join to a related model to create a custom field in the given query.
229
230 This method is used to construct a custom field on the given query based
231 on a many-valued relationsip. join_to_query should be a simple query
232 (no joins) on the related model which returns at most one related row
233 per instance of this model.
234
235 For many-to-one relationships, the joined table contains the matching
236 row from the related model it one is related, NULL otherwise.
237
238 For many-to-many relationships, the joined table contains the matching
239 row if it's related, NULL otherwise.
240 """
241 relationship_type, field = self.determine_relationship(
242 join_to_query.model)
243
244 if relationship_type == self.MANY_TO_ONE:
245 info = self._info_for_many_to_one_join(field, join_to_query, alias)
246 elif relationship_type == self.M2M_ON_RELATED_MODEL:
247 info = self._info_for_many_to_many_join(
248 m2m_field=field, join_to_query=join_to_query, alias=alias,
249 m2m_is_on_this_model=False)
250 elif relationship_type ==self.M2M_ON_THIS_MODEL:
251 info = self._info_for_many_to_many_join(
252 m2m_field=field, join_to_query=join_to_query, alias=alias,
253 m2m_is_on_this_model=True)
254
255 return self.add_join(query_set, info['rhs_table'], info['rhs_column'],
256 join_from_key=info['lhs_column'],
257 join_condition=info['where_clause'],
258 join_condition_values=info['values'],
259 alias=alias,
260 force_left_join=left_join)
261
262
showardf828c772010-01-25 21:49:42 +0000263 def key_on_joined_table(self, join_to_query):
264 """Get a non-null column on the table joined for the given query.
265
266 This analyzes the join that would be produced if join_to_query were
267 passed to join_custom_field.
268 """
269 relationship_type, field = self.determine_relationship(
270 join_to_query.model)
271 if relationship_type == self.MANY_TO_ONE:
272 return join_to_query.model._meta.pk.column
273 return field.m2m_column_name() # any column on the M2M table will do
274
275
showard7e67b432010-01-20 01:13:04 +0000276 def add_where(self, query_set, where, values=()):
277 query_set = query_set.all()
278 query_set.query.where.add(self._WhereClause(where, values),
279 django.db.models.sql.where.AND)
showardc4780402009-08-31 18:31:34 +0000280 return query_set
showard7c785282008-05-29 19:45:12 +0000281
282
showardeaccf8f2009-04-16 03:11:33 +0000283 def _get_quoted_field(self, table, field):
showarda5288b42009-07-28 20:06:08 +0000284 return _quote_name(table) + '.' + _quote_name(field)
showard5ef36e92008-07-02 16:37:09 +0000285
286
showard7c199df2008-10-03 10:17:15 +0000287 def get_key_on_this_table(self, key_field=None):
showard5ef36e92008-07-02 16:37:09 +0000288 if key_field is None:
289 # default to primary key
290 key_field = self.model._meta.pk.column
291 return self._get_quoted_field(self.model._meta.db_table, key_field)
292
293
showardeaccf8f2009-04-16 03:11:33 +0000294 def escape_user_sql(self, sql):
295 return sql.replace('%', '%%')
296
showard5ef36e92008-07-02 16:37:09 +0000297
showard0957a842009-05-11 19:25:08 +0000298 def _custom_select_query(self, query_set, selects):
Jakob Juelich7bef8412014-10-14 19:11:54 -0700299 """Execute a custom select query.
300
301 @param query_set: query set as returned by query_objects.
302 @param selects: Tables/Columns to select, e.g. tko_test_labels_list.id.
303
304 @returns: Result of the query as returned by cursor.fetchall().
305 """
Dale Curtis74a314b2011-06-23 14:55:46 -0700306 compiler = query_set.query.get_compiler(using=query_set.db)
307 sql, params = compiler.as_sql()
showarda5288b42009-07-28 20:06:08 +0000308 from_ = sql[sql.find(' FROM'):]
309
310 if query_set.query.distinct:
showard0957a842009-05-11 19:25:08 +0000311 distinct = 'DISTINCT '
312 else:
313 distinct = ''
showarda5288b42009-07-28 20:06:08 +0000314
315 sql_query = ('SELECT ' + distinct + ','.join(selects) + from_)
Jakob Juelich7bef8412014-10-14 19:11:54 -0700316 # Chose the connection that's responsible for this type of object
317 cursor = connections[query_set.db].cursor()
showard0957a842009-05-11 19:25:08 +0000318 cursor.execute(sql_query, params)
319 return cursor.fetchall()
320
321
showard68693f72009-05-20 00:31:53 +0000322 def _is_relation_to(self, field, model_class):
323 return field.rel and field.rel.to is model_class
showard0957a842009-05-11 19:25:08 +0000324
325
showard7e67b432010-01-20 01:13:04 +0000326 MANY_TO_ONE = object()
327 M2M_ON_RELATED_MODEL = object()
328 M2M_ON_THIS_MODEL = object()
329
330 def determine_relationship(self, related_model):
331 """
332 Determine the relationship between this model and related_model.
333
334 related_model must have some sort of many-valued relationship to this
335 manager's model.
336 @returns (relationship_type, field), where relationship_type is one of
337 MANY_TO_ONE, M2M_ON_RELATED_MODEL, M2M_ON_THIS_MODEL, and field
338 is the Django field object for the relationship.
339 """
340 # look for a foreign key field on related_model relating to this model
341 for field in related_model._meta.fields:
342 if self._is_relation_to(field, self.model):
343 return self.MANY_TO_ONE, field
344
345 # look for an M2M field on related_model relating to this model
346 for field in related_model._meta.many_to_many:
347 if self._is_relation_to(field, self.model):
348 return self.M2M_ON_RELATED_MODEL, field
349
350 # maybe this model has the many-to-many field
351 for field in self.model._meta.many_to_many:
352 if self._is_relation_to(field, related_model):
353 return self.M2M_ON_THIS_MODEL, field
354
355 raise ValueError('%s has no relation to %s' %
356 (related_model, self.model))
357
358
showard68693f72009-05-20 00:31:53 +0000359 def _get_pivot_iterator(self, base_objects_by_id, related_model):
showard0957a842009-05-11 19:25:08 +0000360 """
showard68693f72009-05-20 00:31:53 +0000361 Determine the relationship between this model and related_model, and
362 return a pivot iterator.
363 @param base_objects_by_id: dict of instances of this model indexed by
364 their IDs
365 @returns a pivot iterator, which yields a tuple (base_object,
366 related_object) for each relationship between a base object and a
367 related object. all base_object instances come from base_objects_by_id.
showard7e67b432010-01-20 01:13:04 +0000368 Note -- this depends on Django model internals.
showard0957a842009-05-11 19:25:08 +0000369 """
showard7e67b432010-01-20 01:13:04 +0000370 relationship_type, field = self.determine_relationship(related_model)
371 if relationship_type == self.MANY_TO_ONE:
372 return self._many_to_one_pivot(base_objects_by_id,
373 related_model, field)
374 elif relationship_type == self.M2M_ON_RELATED_MODEL:
375 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000376 base_objects_by_id, related_model, field.m2m_db_table(),
377 field.m2m_reverse_name(), field.m2m_column_name())
showard7e67b432010-01-20 01:13:04 +0000378 else:
379 assert relationship_type == self.M2M_ON_THIS_MODEL
380 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000381 base_objects_by_id, related_model, field.m2m_db_table(),
382 field.m2m_column_name(), field.m2m_reverse_name())
showard0957a842009-05-11 19:25:08 +0000383
showard0957a842009-05-11 19:25:08 +0000384
showard68693f72009-05-20 00:31:53 +0000385 def _many_to_one_pivot(self, base_objects_by_id, related_model,
386 foreign_key_field):
387 """
388 @returns a pivot iterator - see _get_pivot_iterator()
389 """
390 filter_data = {foreign_key_field.name + '__pk__in':
391 base_objects_by_id.keys()}
392 for related_object in related_model.objects.filter(**filter_data):
showarda5a72c92009-08-20 23:35:21 +0000393 # lookup base object in the dict, rather than grabbing it from the
394 # related object. we need to return instances from the dict, not
395 # fresh instances of the same models (and grabbing model instances
396 # from the related models incurs a DB query each time).
397 base_object_id = getattr(related_object, foreign_key_field.attname)
398 base_object = base_objects_by_id[base_object_id]
showard68693f72009-05-20 00:31:53 +0000399 yield base_object, related_object
400
401
402 def _query_pivot_table(self, base_objects_by_id, pivot_table,
Jakob Juelich7bef8412014-10-14 19:11:54 -0700403 pivot_from_field, pivot_to_field, related_model):
showard0957a842009-05-11 19:25:08 +0000404 """
405 @param id_list list of IDs of self.model objects to include
406 @param pivot_table the name of the pivot table
407 @param pivot_from_field a field name on pivot_table referencing
408 self.model
409 @param pivot_to_field a field name on pivot_table referencing the
410 related model.
Jakob Juelich7bef8412014-10-14 19:11:54 -0700411 @param related_model the related model
412
showard68693f72009-05-20 00:31:53 +0000413 @returns pivot list of IDs (base_id, related_id)
showard0957a842009-05-11 19:25:08 +0000414 """
415 query = """
416 SELECT %(from_field)s, %(to_field)s
417 FROM %(table)s
418 WHERE %(from_field)s IN (%(id_list)s)
419 """ % dict(from_field=pivot_from_field,
420 to_field=pivot_to_field,
421 table=pivot_table,
showard68693f72009-05-20 00:31:53 +0000422 id_list=','.join(str(id_) for id_
423 in base_objects_by_id.iterkeys()))
Jakob Juelich7bef8412014-10-14 19:11:54 -0700424
425 # Chose the connection that's responsible for this type of object
426 # The databases for related_model and the current model will always
427 # be the same, related_model is just easier to obtain here because
428 # self is only a ExtendedManager, not the object.
429 cursor = connections[related_model.objects.db].cursor()
showard0957a842009-05-11 19:25:08 +0000430 cursor.execute(query)
showard68693f72009-05-20 00:31:53 +0000431 return cursor.fetchall()
showard0957a842009-05-11 19:25:08 +0000432
433
showard68693f72009-05-20 00:31:53 +0000434 def _many_to_many_pivot(self, base_objects_by_id, related_model,
435 pivot_table, pivot_from_field, pivot_to_field):
436 """
437 @param pivot_table: see _query_pivot_table
438 @param pivot_from_field: see _query_pivot_table
439 @param pivot_to_field: see _query_pivot_table
440 @returns a pivot iterator - see _get_pivot_iterator()
441 """
442 id_pivot = self._query_pivot_table(base_objects_by_id, pivot_table,
Jakob Juelich7bef8412014-10-14 19:11:54 -0700443 pivot_from_field, pivot_to_field,
444 related_model)
showard68693f72009-05-20 00:31:53 +0000445
446 all_related_ids = list(set(related_id for base_id, related_id
447 in id_pivot))
448 related_objects_by_id = related_model.objects.in_bulk(all_related_ids)
449
450 for base_id, related_id in id_pivot:
451 yield base_objects_by_id[base_id], related_objects_by_id[related_id]
452
453
454 def populate_relationships(self, base_objects, related_model,
showard0957a842009-05-11 19:25:08 +0000455 related_list_name):
456 """
showard68693f72009-05-20 00:31:53 +0000457 For each instance of this model in base_objects, add a field named
458 related_list_name listing all the related objects of type related_model.
459 related_model must be in a many-to-one or many-to-many relationship with
460 this model.
461 @param base_objects - list of instances of this model
462 @param related_model - model class related to this model
463 @param related_list_name - attribute name in which to store the related
464 object list.
showard0957a842009-05-11 19:25:08 +0000465 """
showard68693f72009-05-20 00:31:53 +0000466 if not base_objects:
showard0957a842009-05-11 19:25:08 +0000467 # if we don't bail early, we'll get a SQL error later
468 return
showard0957a842009-05-11 19:25:08 +0000469
showard68693f72009-05-20 00:31:53 +0000470 base_objects_by_id = dict((base_object._get_pk_val(), base_object)
471 for base_object in base_objects)
472 pivot_iterator = self._get_pivot_iterator(base_objects_by_id,
473 related_model)
showard0957a842009-05-11 19:25:08 +0000474
showard68693f72009-05-20 00:31:53 +0000475 for base_object in base_objects:
476 setattr(base_object, related_list_name, [])
477
478 for base_object, related_object in pivot_iterator:
479 getattr(base_object, related_list_name).append(related_object)
showard0957a842009-05-11 19:25:08 +0000480
481
jamesrene3656232010-03-02 00:00:30 +0000482class ModelWithInvalidQuerySet(dbmodels.query.QuerySet):
483 """
484 QuerySet that handles delete() properly for models with an "invalid" bit
485 """
486 def delete(self):
487 for model in self:
488 model.delete()
489
490
491class ModelWithInvalidManager(ExtendedManager):
492 """
493 Manager for objects with an "invalid" bit
494 """
495 def get_query_set(self):
496 return ModelWithInvalidQuerySet(self.model)
497
498
499class ValidObjectsManager(ModelWithInvalidManager):
jadmanski0afbb632008-06-06 21:10:57 +0000500 """
501 Manager returning only objects with invalid=False.
502 """
503 def get_query_set(self):
504 queryset = super(ValidObjectsManager, self).get_query_set()
505 return queryset.filter(invalid=False)
showard7c785282008-05-29 19:45:12 +0000506
507
Prashanth B489b91d2014-03-15 12:17:16 -0700508class ModelExtensions(rdb_model_extensions.ModelValidators):
jadmanski0afbb632008-06-06 21:10:57 +0000509 """\
Prashanth B489b91d2014-03-15 12:17:16 -0700510 Mixin with convenience functions for models, built on top of
511 the model validators in rdb_model_extensions.
jadmanski0afbb632008-06-06 21:10:57 +0000512 """
513 # TODO: at least some of these functions really belong in a custom
514 # Manager class
showard7c785282008-05-29 19:45:12 +0000515
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700516
517 SERIALIZATION_LINKS_TO_FOLLOW = set()
518 """
519 To be able to send jobs and hosts to shards, it's necessary to find their
520 dependencies.
521 The most generic approach for this would be to traverse all relationships
522 to other objects recursively. This would list all objects that are related
523 in any way.
524 But this approach finds too many objects: If a host should be transferred,
525 all it's relationships would be traversed. This would find an acl group.
526 If then the acl group's relationships are traversed, the relationship
527 would be followed backwards and many other hosts would be found.
528
529 This mapping tells that algorithm which relations to follow explicitly.
530 """
531
Jakob Juelichf865d332014-09-29 10:47:49 -0700532
Fang Deng86248502014-12-18 16:38:00 -0800533 SERIALIZATION_LINKS_TO_KEEP = set()
534 """This set stores foreign keys which we don't want to follow, but
535 still want to include in the serialized dictionary. For
536 example, we follow the relationship `Host.hostattribute_set`,
537 but we do not want to follow `HostAttributes.host_id` back to
538 to Host, which would otherwise lead to a circle. However, we still
539 like to serialize HostAttribute.`host_id`."""
540
Jakob Juelichf865d332014-09-29 10:47:49 -0700541 SERIALIZATION_LOCAL_LINKS_TO_UPDATE = set()
542 """
543 On deserializion, if the object to persist already exists, local fields
544 will only be updated, if their name is in this set.
545 """
546
547
jadmanski0afbb632008-06-06 21:10:57 +0000548 @classmethod
549 def convert_human_readable_values(cls, data, to_human_readable=False):
550 """\
551 Performs conversions on user-supplied field data, to make it
552 easier for users to pass human-readable data.
showard7c785282008-05-29 19:45:12 +0000553
jadmanski0afbb632008-06-06 21:10:57 +0000554 For all fields that have choice sets, convert their values
555 from human-readable strings to enum values, if necessary. This
556 allows users to pass strings instead of the corresponding
557 integer values.
showard7c785282008-05-29 19:45:12 +0000558
jadmanski0afbb632008-06-06 21:10:57 +0000559 For all foreign key fields, call smart_get with the supplied
560 data. This allows the user to pass either an ID value or
561 the name of the object as a string.
showard7c785282008-05-29 19:45:12 +0000562
jadmanski0afbb632008-06-06 21:10:57 +0000563 If to_human_readable=True, perform the inverse - i.e. convert
564 numeric values to human readable values.
showard7c785282008-05-29 19:45:12 +0000565
jadmanski0afbb632008-06-06 21:10:57 +0000566 This method modifies data in-place.
567 """
568 field_dict = cls.get_field_dict()
569 for field_name in data:
showarde732ee72008-09-23 19:15:43 +0000570 if field_name not in field_dict or data[field_name] is None:
jadmanski0afbb632008-06-06 21:10:57 +0000571 continue
572 field_obj = field_dict[field_name]
573 # convert enum values
574 if field_obj.choices:
575 for choice_data in field_obj.choices:
576 # choice_data is (value, name)
577 if to_human_readable:
578 from_val, to_val = choice_data
579 else:
580 to_val, from_val = choice_data
581 if from_val == data[field_name]:
582 data[field_name] = to_val
583 break
584 # convert foreign key values
585 elif field_obj.rel:
showarda4ea5742009-02-17 20:56:23 +0000586 dest_obj = field_obj.rel.to.smart_get(data[field_name],
587 valid_only=False)
showardf8b19042009-05-12 17:22:49 +0000588 if to_human_readable:
Paul Pendlebury5a8c6ad2011-02-01 07:20:17 -0800589 # parameterized_jobs do not have a name_field
590 if (field_name != 'parameterized_job' and
591 dest_obj.name_field is not None):
showardf8b19042009-05-12 17:22:49 +0000592 data[field_name] = getattr(dest_obj,
593 dest_obj.name_field)
jadmanski0afbb632008-06-06 21:10:57 +0000594 else:
showardb0a73032009-03-27 18:35:41 +0000595 data[field_name] = dest_obj
showard7c785282008-05-29 19:45:12 +0000596
597
showard7c785282008-05-29 19:45:12 +0000598
599
Dale Curtis74a314b2011-06-23 14:55:46 -0700600 def _validate_unique(self):
jadmanski0afbb632008-06-06 21:10:57 +0000601 """\
602 Validate that unique fields are unique. Django manipulators do
603 this too, but they're a huge pain to use manually. Trust me.
604 """
605 errors = {}
606 cls = type(self)
607 field_dict = self.get_field_dict()
608 manager = cls.get_valid_manager()
609 for field_name, field_obj in field_dict.iteritems():
610 if not field_obj.unique:
611 continue
showard7c785282008-05-29 19:45:12 +0000612
jadmanski0afbb632008-06-06 21:10:57 +0000613 value = getattr(self, field_name)
showardbd18ab72009-09-18 21:20:27 +0000614 if value is None and field_obj.auto_created:
615 # don't bother checking autoincrement fields about to be
616 # generated
617 continue
618
jadmanski0afbb632008-06-06 21:10:57 +0000619 existing_objs = manager.filter(**{field_name : value})
620 num_existing = existing_objs.count()
showard7c785282008-05-29 19:45:12 +0000621
jadmanski0afbb632008-06-06 21:10:57 +0000622 if num_existing == 0:
623 continue
624 if num_existing == 1 and existing_objs[0].id == self.id:
625 continue
626 errors[field_name] = (
627 'This value must be unique (%s)' % (value))
628 return errors
showard7c785282008-05-29 19:45:12 +0000629
630
showarda5288b42009-07-28 20:06:08 +0000631 def _validate(self):
632 """
633 First coerces all fields on this instance to their proper Python types.
634 Then runs validation on every field. Returns a dictionary of
635 field_name -> error_list.
636
637 Based on validate() from django.db.models.Model in Django 0.96, which
638 was removed in Django 1.0. It should reappear in a later version. See:
639 http://code.djangoproject.com/ticket/6845
640 """
641 error_dict = {}
642 for f in self._meta.fields:
643 try:
644 python_value = f.to_python(
645 getattr(self, f.attname, f.get_default()))
646 except django.core.exceptions.ValidationError, e:
jamesren1e0a4ce2010-04-21 17:45:11 +0000647 error_dict[f.name] = str(e)
showarda5288b42009-07-28 20:06:08 +0000648 continue
649
650 if not f.blank and not python_value:
651 error_dict[f.name] = 'This field is required.'
652 continue
653
654 setattr(self, f.attname, python_value)
655
656 return error_dict
657
658
jadmanski0afbb632008-06-06 21:10:57 +0000659 def do_validate(self):
showarda5288b42009-07-28 20:06:08 +0000660 errors = self._validate()
Dale Curtis74a314b2011-06-23 14:55:46 -0700661 unique_errors = self._validate_unique()
jadmanski0afbb632008-06-06 21:10:57 +0000662 for field_name, error in unique_errors.iteritems():
663 errors.setdefault(field_name, error)
664 if errors:
665 raise ValidationError(errors)
showard7c785282008-05-29 19:45:12 +0000666
667
jadmanski0afbb632008-06-06 21:10:57 +0000668 # actually (externally) useful methods follow
showard7c785282008-05-29 19:45:12 +0000669
jadmanski0afbb632008-06-06 21:10:57 +0000670 @classmethod
671 def add_object(cls, data={}, **kwargs):
672 """\
673 Returns a new object created with the given data (a dictionary
674 mapping field names to values). Merges any extra keyword args
675 into data.
676 """
Prashanth B489b91d2014-03-15 12:17:16 -0700677 data = dict(data)
678 data.update(kwargs)
679 data = cls.prepare_data_args(data)
680 cls.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000681 data = cls.provide_default_values(data)
Prashanth B489b91d2014-03-15 12:17:16 -0700682
jadmanski0afbb632008-06-06 21:10:57 +0000683 obj = cls(**data)
684 obj.do_validate()
685 obj.save()
686 return obj
showard7c785282008-05-29 19:45:12 +0000687
688
jadmanski0afbb632008-06-06 21:10:57 +0000689 def update_object(self, data={}, **kwargs):
690 """\
691 Updates the object with the given data (a dictionary mapping
692 field names to values). Merges any extra keyword args into
693 data.
694 """
Prashanth B489b91d2014-03-15 12:17:16 -0700695 data = dict(data)
696 data.update(kwargs)
697 data = self.prepare_data_args(data)
698 self.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000699 for field_name, value in data.iteritems():
showardb0a73032009-03-27 18:35:41 +0000700 setattr(self, field_name, value)
jadmanski0afbb632008-06-06 21:10:57 +0000701 self.do_validate()
702 self.save()
showard7c785282008-05-29 19:45:12 +0000703
704
showard8bfb5cb2009-10-07 20:49:15 +0000705 # see query_objects()
706 _SPECIAL_FILTER_KEYS = ('query_start', 'query_limit', 'sort_by',
707 'extra_args', 'extra_where', 'no_distinct')
708
709
jadmanski0afbb632008-06-06 21:10:57 +0000710 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000711 def _extract_special_params(cls, filter_data):
712 """
713 @returns a tuple of dicts (special_params, regular_filters), where
714 special_params contains the parameters we handle specially and
715 regular_filters is the remaining data to be handled by Django.
716 """
717 regular_filters = dict(filter_data)
718 special_params = {}
719 for key in cls._SPECIAL_FILTER_KEYS:
720 if key in regular_filters:
721 special_params[key] = regular_filters.pop(key)
722 return special_params, regular_filters
723
724
725 @classmethod
726 def apply_presentation(cls, query, filter_data):
727 """
728 Apply presentation parameters -- sorting and paging -- to the given
729 query.
730 @returns new query with presentation applied
731 """
732 special_params, _ = cls._extract_special_params(filter_data)
733 sort_by = special_params.get('sort_by', None)
734 if sort_by:
735 assert isinstance(sort_by, list) or isinstance(sort_by, tuple)
showard8b0ea222009-12-23 19:23:03 +0000736 query = query.extra(order_by=sort_by)
showard8bfb5cb2009-10-07 20:49:15 +0000737
738 query_start = special_params.get('query_start', None)
739 query_limit = special_params.get('query_limit', None)
740 if query_start is not None:
741 if query_limit is None:
742 raise ValueError('Cannot pass query_start without query_limit')
743 # query_limit is passed as a page size
showard7074b742009-10-12 20:30:04 +0000744 query_limit += query_start
745 return query[query_start:query_limit]
showard8bfb5cb2009-10-07 20:49:15 +0000746
747
748 @classmethod
749 def query_objects(cls, filter_data, valid_only=True, initial_query=None,
750 apply_presentation=True):
jadmanski0afbb632008-06-06 21:10:57 +0000751 """\
752 Returns a QuerySet object for querying the given model_class
753 with the given filter_data. Optional special arguments in
754 filter_data include:
755 -query_start: index of first return to return
756 -query_limit: maximum number of results to return
757 -sort_by: list of fields to sort on. prefixing a '-' onto a
758 field name changes the sort to descending order.
759 -extra_args: keyword args to pass to query.extra() (see Django
760 DB layer documentation)
showarda5288b42009-07-28 20:06:08 +0000761 -extra_where: extra WHERE clause to append
showard8bfb5cb2009-10-07 20:49:15 +0000762 -no_distinct: if True, a DISTINCT will not be added to the SELECT
jadmanski0afbb632008-06-06 21:10:57 +0000763 """
showard8bfb5cb2009-10-07 20:49:15 +0000764 special_params, regular_filters = cls._extract_special_params(
765 filter_data)
showard7c785282008-05-29 19:45:12 +0000766
showard7ac7b7a2008-07-21 20:24:29 +0000767 if initial_query is None:
768 if valid_only:
769 initial_query = cls.get_valid_manager()
770 else:
771 initial_query = cls.objects
showard8bfb5cb2009-10-07 20:49:15 +0000772
773 query = initial_query.filter(**regular_filters)
774
775 use_distinct = not special_params.get('no_distinct', False)
showard7ac7b7a2008-07-21 20:24:29 +0000776 if use_distinct:
777 query = query.distinct()
showard7c785282008-05-29 19:45:12 +0000778
showard8bfb5cb2009-10-07 20:49:15 +0000779 extra_args = special_params.get('extra_args', {})
780 extra_where = special_params.get('extra_where', None)
781 if extra_where:
782 # escape %'s
783 extra_where = cls.objects.escape_user_sql(extra_where)
784 extra_args.setdefault('where', []).append(extra_where)
jadmanski0afbb632008-06-06 21:10:57 +0000785 if extra_args:
786 query = query.extra(**extra_args)
Jakob Juelich7bef8412014-10-14 19:11:54 -0700787 # TODO: Use readonly connection for these queries.
788 # This has been disabled, because it's not used anyway, as the
789 # configured readonly user is the same as the real user anyway.
showard7c785282008-05-29 19:45:12 +0000790
showard8bfb5cb2009-10-07 20:49:15 +0000791 if apply_presentation:
792 query = cls.apply_presentation(query, filter_data)
793
794 return query
showard7c785282008-05-29 19:45:12 +0000795
796
jadmanski0afbb632008-06-06 21:10:57 +0000797 @classmethod
showard585c2ab2008-07-23 19:29:49 +0000798 def query_count(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000799 """\
800 Like query_objects, but retreive only the count of results.
801 """
802 filter_data.pop('query_start', None)
803 filter_data.pop('query_limit', None)
showard585c2ab2008-07-23 19:29:49 +0000804 query = cls.query_objects(filter_data, initial_query=initial_query)
805 return query.count()
showard7c785282008-05-29 19:45:12 +0000806
807
jadmanski0afbb632008-06-06 21:10:57 +0000808 @classmethod
809 def clean_object_dicts(cls, field_dicts):
810 """\
811 Take a list of dicts corresponding to object (as returned by
812 query.values()) and clean the data to be more suitable for
813 returning to the user.
814 """
showarde732ee72008-09-23 19:15:43 +0000815 for field_dict in field_dicts:
816 cls.clean_foreign_keys(field_dict)
showard21baa452008-10-21 00:08:39 +0000817 cls._convert_booleans(field_dict)
showarde732ee72008-09-23 19:15:43 +0000818 cls.convert_human_readable_values(field_dict,
819 to_human_readable=True)
showard7c785282008-05-29 19:45:12 +0000820
821
jadmanski0afbb632008-06-06 21:10:57 +0000822 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000823 def list_objects(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000824 """\
825 Like query_objects, but return a list of dictionaries.
826 """
showard7ac7b7a2008-07-21 20:24:29 +0000827 query = cls.query_objects(filter_data, initial_query=initial_query)
showard8bfb5cb2009-10-07 20:49:15 +0000828 extra_fields = query.query.extra_select.keys()
829 field_dicts = [model_object.get_object_dict(extra_fields=extra_fields)
showarde732ee72008-09-23 19:15:43 +0000830 for model_object in query]
jadmanski0afbb632008-06-06 21:10:57 +0000831 return field_dicts
showard7c785282008-05-29 19:45:12 +0000832
833
jadmanski0afbb632008-06-06 21:10:57 +0000834 @classmethod
showarda4ea5742009-02-17 20:56:23 +0000835 def smart_get(cls, id_or_name, valid_only=True):
jadmanski0afbb632008-06-06 21:10:57 +0000836 """\
837 smart_get(integer) -> get object by ID
838 smart_get(string) -> get object by name_field
jadmanski0afbb632008-06-06 21:10:57 +0000839 """
showarda4ea5742009-02-17 20:56:23 +0000840 if valid_only:
841 manager = cls.get_valid_manager()
842 else:
843 manager = cls.objects
844
845 if isinstance(id_or_name, (int, long)):
846 return manager.get(pk=id_or_name)
jamesren3e9f6092010-03-11 21:32:10 +0000847 if isinstance(id_or_name, basestring) and hasattr(cls, 'name_field'):
showarda4ea5742009-02-17 20:56:23 +0000848 return manager.get(**{cls.name_field : id_or_name})
849 raise ValueError(
850 'Invalid positional argument: %s (%s)' % (id_or_name,
851 type(id_or_name)))
showard7c785282008-05-29 19:45:12 +0000852
853
showardbe3ec042008-11-12 18:16:07 +0000854 @classmethod
855 def smart_get_bulk(cls, id_or_name_list):
856 invalid_inputs = []
857 result_objects = []
858 for id_or_name in id_or_name_list:
859 try:
860 result_objects.append(cls.smart_get(id_or_name))
861 except cls.DoesNotExist:
862 invalid_inputs.append(id_or_name)
863 if invalid_inputs:
mbligh7a3ebe32008-12-01 17:10:33 +0000864 raise cls.DoesNotExist('The following %ss do not exist: %s'
865 % (cls.__name__.lower(),
866 ', '.join(invalid_inputs)))
showardbe3ec042008-11-12 18:16:07 +0000867 return result_objects
868
869
showard8bfb5cb2009-10-07 20:49:15 +0000870 def get_object_dict(self, extra_fields=None):
jadmanski0afbb632008-06-06 21:10:57 +0000871 """\
showard8bfb5cb2009-10-07 20:49:15 +0000872 Return a dictionary mapping fields to this object's values. @param
873 extra_fields: list of extra attribute names to include, in addition to
874 the fields defined on this object.
jadmanski0afbb632008-06-06 21:10:57 +0000875 """
showard8bfb5cb2009-10-07 20:49:15 +0000876 fields = self.get_field_dict().keys()
877 if extra_fields:
878 fields += extra_fields
jadmanski0afbb632008-06-06 21:10:57 +0000879 object_dict = dict((field_name, getattr(self, field_name))
showarde732ee72008-09-23 19:15:43 +0000880 for field_name in fields)
jadmanski0afbb632008-06-06 21:10:57 +0000881 self.clean_object_dicts([object_dict])
showardd3dc1992009-04-22 21:01:40 +0000882 self._postprocess_object_dict(object_dict)
jadmanski0afbb632008-06-06 21:10:57 +0000883 return object_dict
showard7c785282008-05-29 19:45:12 +0000884
885
showardd3dc1992009-04-22 21:01:40 +0000886 def _postprocess_object_dict(self, object_dict):
887 """For subclasses to override."""
888 pass
889
890
jadmanski0afbb632008-06-06 21:10:57 +0000891 @classmethod
892 def get_valid_manager(cls):
893 return cls.objects
showard7c785282008-05-29 19:45:12 +0000894
895
showard2bab8f42008-11-12 18:15:22 +0000896 def _record_attributes(self, attributes):
897 """
898 See on_attribute_changed.
899 """
900 assert not isinstance(attributes, basestring)
901 self._recorded_attributes = dict((attribute, getattr(self, attribute))
902 for attribute in attributes)
903
904
905 def _check_for_updated_attributes(self):
906 """
907 See on_attribute_changed.
908 """
909 for attribute, original_value in self._recorded_attributes.iteritems():
910 new_value = getattr(self, attribute)
911 if original_value != new_value:
912 self.on_attribute_changed(attribute, original_value)
913 self._record_attributes(self._recorded_attributes.keys())
914
915
916 def on_attribute_changed(self, attribute, old_value):
917 """
918 Called whenever an attribute is updated. To be overridden.
919
920 To use this method, you must:
921 * call _record_attributes() from __init__() (after making the super
922 call) with a list of attributes for which you want to be notified upon
923 change.
924 * call _check_for_updated_attributes() from save().
925 """
926 pass
927
928
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700929 def serialize(self, include_dependencies=True):
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700930 """Serializes the object with dependencies.
931
932 The variable SERIALIZATION_LINKS_TO_FOLLOW defines which dependencies
933 this function will serialize with the object.
934
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700935 @param include_dependencies: Whether or not to follow relations to
936 objects this object depends on.
937 This parameter is used when uploading
938 jobs from a shard to the master, as the
939 master already has all the dependent
940 objects.
941
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700942 @returns: Dictionary representation of the object.
943 """
944 serialized = {}
MK Ryu5cfd96a2015-01-30 15:31:23 -0800945 timer = autotest_stats.Timer('serialize_latency.%s' % (
946 type(self).__name__))
947 with timer.get_client('local'):
948 for field in self._meta.concrete_model._meta.local_fields:
949 if field.rel is None:
950 serialized[field.name] = field._get_val_from_obj(self)
MK Ryudf4d4232015-03-06 11:18:47 -0800951 elif field.name in self.SERIALIZATION_LINKS_TO_KEEP:
MK Ryu5cfd96a2015-01-30 15:31:23 -0800952 # attname will contain "_id" suffix for foreign keys,
953 # e.g. HostAttribute.host will be serialized as 'host_id'.
954 # Use it for easy deserialization.
955 serialized[field.attname] = field._get_val_from_obj(self)
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700956
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700957 if include_dependencies:
MK Ryu5cfd96a2015-01-30 15:31:23 -0800958 with timer.get_client('related'):
959 for link in self.SERIALIZATION_LINKS_TO_FOLLOW:
960 serialized[link] = self._serialize_relation(link)
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700961
962 return serialized
963
964
965 def _serialize_relation(self, link):
966 """Serializes dependent objects given the name of the relation.
967
968 @param link: Name of the relation to take objects from.
969
970 @returns For To-Many relationships a list of the serialized related
971 objects, for To-One relationships the serialized related object.
972 """
973 try:
974 attr = getattr(self, link)
975 except AttributeError:
976 # One-To-One relationships that point to None may raise this
977 return None
978
979 if attr is None:
980 return None
981 if hasattr(attr, 'all'):
982 return [obj.serialize() for obj in attr.all()]
983 return attr.serialize()
984
985
Jakob Juelichf88fa932014-09-03 17:58:04 -0700986 @classmethod
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700987 def _split_local_from_foreign_values(cls, data):
988 """This splits local from foreign values in a serialized object.
989
990 @param data: The serialized object.
991
992 @returns A tuple of two lists, both containing tuples in the form
993 (link_name, link_value). The first list contains all links
994 for local fields, the second one contains those for foreign
995 fields/objects.
996 """
997 links_to_local_values, links_to_related_values = [], []
998 for link, value in data.iteritems():
999 if link in cls.SERIALIZATION_LINKS_TO_FOLLOW:
1000 # It's a foreign key
1001 links_to_related_values.append((link, value))
1002 else:
Fang Deng86248502014-12-18 16:38:00 -08001003 # It's a local attribute or a foreign key
1004 # we don't want to follow.
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001005 links_to_local_values.append((link, value))
1006 return links_to_local_values, links_to_related_values
1007
1008
Jakob Juelichf865d332014-09-29 10:47:49 -07001009 @classmethod
1010 def _filter_update_allowed_fields(cls, data):
1011 """Filters data and returns only files that updates are allowed on.
1012
1013 This is i.e. needed for syncing aborted bits from the master to shards.
1014
1015 Local links are only allowed to be updated, if they are in
1016 SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1017 Overwriting existing values is allowed in order to be able to sync i.e.
1018 the aborted bit from the master to a shard.
1019
1020 The whitelisting mechanism is in place to prevent overwriting local
1021 status: If all fields were overwritten, jobs would be completely be
1022 set back to their original (unstarted) state.
1023
1024 @param data: List with tuples of the form (link_name, link_value), as
1025 returned by _split_local_from_foreign_values.
1026
1027 @returns List of the same format as data, but only containing data for
1028 fields that updates are allowed on.
1029 """
1030 return [pair for pair in data
1031 if pair[0] in cls.SERIALIZATION_LOCAL_LINKS_TO_UPDATE]
1032
1033
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001034 @classmethod
1035 def delete_matching_record(cls, **filter_args):
1036 """Delete records matching the filter.
1037
1038 @param filter_args: Arguments for the django filter
1039 used to locate the record to delete.
1040 """
1041 try:
1042 existing_record = cls.objects.get(**filter_args)
1043 except cls.DoesNotExist:
1044 return
1045 existing_record.delete()
1046
1047
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001048 def _deserialize_local(self, data):
1049 """Set local attributes from a list of tuples.
1050
1051 @param data: List of tuples like returned by
1052 _split_local_from_foreign_values.
1053 """
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001054 if not data:
1055 return
1056
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001057 for link, value in data:
1058 setattr(self, link, value)
1059 # Overwridden save() methods are prone to errors, so don't execute them.
1060 # This is because:
1061 # - the overwritten methods depend on ACL groups that don't yet exist
1062 # and don't handle errors
1063 # - the overwritten methods think this object already exists in the db
1064 # because the id is already set
1065 super(type(self), self).save()
1066
1067
1068 def _deserialize_relations(self, data):
1069 """Set foreign attributes from a list of tuples.
1070
1071 This deserialized the related objects using their own deserialize()
1072 function and then sets the relation.
1073
1074 @param data: List of tuples like returned by
1075 _split_local_from_foreign_values.
1076 """
1077 for link, value in data:
1078 self._deserialize_relation(link, value)
1079 # See comment in _deserialize_local
1080 super(type(self), self).save()
1081
1082
1083 @classmethod
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001084 def get_record(cls, data):
1085 """Retrieve a record with the data in the given input arg.
1086
1087 @param data: A dictionary containing the information to use in a query
1088 for data. If child models have different constraints of
1089 uniqueness they should override this model.
1090
1091 @return: An object with matching data.
1092
1093 @raises DoesNotExist: If a record with the given data doesn't exist.
1094 """
1095 return cls.objects.get(id=data['id'])
1096
1097
1098 @classmethod
Jakob Juelichf88fa932014-09-03 17:58:04 -07001099 def deserialize(cls, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001100 """Recursively deserializes and saves an object with it's dependencies.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001101
1102 This takes the result of the serialize method and creates objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001103 in the database that are just like the original.
1104
1105 If an object of the same type with the same id already exists, it's
Jakob Juelichf865d332014-09-29 10:47:49 -07001106 local values will be left untouched, unless they are explicitly
1107 whitelisted in SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1108
1109 Deserialize will always recursively propagate to all related objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001110 present in data though.
1111 I.e. this is necessary to add users to an already existing acl-group.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001112
1113 @param data: Representation of an object and its dependencies, as
1114 returned by serialize.
1115
1116 @returns: The object represented by data if it didn't exist before,
1117 otherwise the object that existed before and has the same type
1118 and id as the one described by data.
1119 """
1120 if data is None:
1121 return None
1122
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001123 local, related = cls._split_local_from_foreign_values(data)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001124 try:
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001125 instance = cls.get_record(data)
Jakob Juelichf865d332014-09-29 10:47:49 -07001126 local = cls._filter_update_allowed_fields(local)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001127 except cls.DoesNotExist:
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001128 instance = cls()
Jakob Juelichf88fa932014-09-03 17:58:04 -07001129
MK Ryu5cfd96a2015-01-30 15:31:23 -08001130 timer = autotest_stats.Timer('deserialize_latency.%s' % (
1131 type(instance).__name__))
1132 with timer.get_client('local'):
1133 instance._deserialize_local(local)
1134 with timer.get_client('related'):
1135 instance._deserialize_relations(related)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001136
1137 return instance
1138
1139
Jakob Juelicha94efe62014-09-18 16:02:49 -07001140 def sanity_check_update_from_shard(self, shard, updated_serialized,
1141 *args, **kwargs):
1142 """Check if an update sent from a shard is legitimate.
1143
1144 @raises error.UnallowedRecordsSentToMaster if an update is not
1145 legitimate.
1146 """
1147 raise NotImplementedError(
1148 'sanity_check_update_from_shard must be implemented by subclass %s '
1149 'for type %s' % type(self))
1150
1151
Prashanth Balasubramanian75be1d32014-11-25 18:03:09 -08001152 @transaction.commit_on_success
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001153 def update_from_serialized(self, serialized):
1154 """Updates local fields of an existing object from a serialized form.
1155
1156 This is different than the normal deserialize() in the way that it
1157 does update local values, which deserialize doesn't, but doesn't
1158 recursively propagate to related objects, which deserialize() does.
1159
1160 The use case of this function is to update job records on the master
1161 after the jobs have been executed on a slave, as the master is not
1162 interested in updates for users, labels, specialtasks, etc.
1163
1164 @param serialized: Representation of an object and its dependencies, as
1165 returned by serialize.
1166
1167 @raises ValueError: if serialized contains related objects, i.e. not
1168 only local fields.
1169 """
1170 local, related = (
1171 self._split_local_from_foreign_values(serialized))
1172 if related:
1173 raise ValueError('Serialized must not contain foreign '
1174 'objects: %s' % related)
1175
1176 self._deserialize_local(local)
1177
1178
Jakob Juelichf88fa932014-09-03 17:58:04 -07001179 def custom_deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001180 """Allows overriding the deserialization behaviour by subclasses."""
Jakob Juelichf88fa932014-09-03 17:58:04 -07001181 raise NotImplementedError(
1182 'custom_deserialize_relation must be implemented by subclass %s '
1183 'for relation %s' % (type(self), link))
1184
1185
1186 def _deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001187 """Deserializes related objects and sets references on this object.
1188
1189 Relations that point to a list of objects are handled automatically.
1190 For many-to-one or one-to-one relations custom_deserialize_relation
1191 must be overridden by the subclass.
1192
1193 Related objects are deserialized using their deserialize() method.
1194 Thereby they and their dependencies are created if they don't exist
1195 and saved to the database.
1196
1197 @param link: Name of the relation.
1198 @param data: Serialized representation of the related object(s).
1199 This means a list of dictionaries for to-many relations,
1200 just a dictionary for to-one relations.
1201 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001202 field = getattr(self, link)
1203
1204 if field and hasattr(field, 'all'):
1205 self._deserialize_2m_relation(link, data, field.model)
1206 else:
1207 self.custom_deserialize_relation(link, data)
1208
1209
1210 def _deserialize_2m_relation(self, link, data, related_class):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001211 """Deserialize related objects for one to-many relationship.
1212
1213 @param link: Name of the relation.
1214 @param data: Serialized representation of the related objects.
1215 This is a list with of dictionaries.
Fang Dengff361592015-02-02 15:27:34 -08001216 @param related_class: A class representing a django model, with which
1217 this class has a one-to-many relationship.
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001218 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001219 relation_set = getattr(self, link)
Fang Dengff361592015-02-02 15:27:34 -08001220 if related_class == self.get_attribute_model():
1221 # When deserializing a model together with
1222 # its attributes, clear all the exising attributes to ensure
1223 # db consistency. Note 'update' won't be sufficient, as we also
1224 # want to remove any attributes that no longer exist in |data|.
1225 #
1226 # core_filters is a dictionary of filters, defines how
1227 # RelatedMangager would query for the 1-to-many relationship. E.g.
1228 # Host.objects.get(
1229 # id=20).hostattribute_set.core_filters = {host_id:20}
1230 # We use it to delete objects related to the current object.
1231 related_class.objects.filter(**relation_set.core_filters).delete()
Jakob Juelichf88fa932014-09-03 17:58:04 -07001232 for serialized in data:
1233 relation_set.add(related_class.deserialize(serialized))
1234
1235
Fang Dengff361592015-02-02 15:27:34 -08001236 @classmethod
1237 def get_attribute_model(cls):
1238 """Return the attribute model.
1239
1240 Subclass with attribute-like model should override this to
1241 return the attribute model class. This method will be
1242 called by _deserialize_2m_relation to determine whether
1243 to clear the one-to-many relations first on deserialization of object.
1244 """
1245 return None
1246
1247
showard7c785282008-05-29 19:45:12 +00001248class ModelWithInvalid(ModelExtensions):
jadmanski0afbb632008-06-06 21:10:57 +00001249 """
1250 Overrides model methods save() and delete() to support invalidation in
1251 place of actual deletion. Subclasses must have a boolean "invalid"
1252 field.
1253 """
showard7c785282008-05-29 19:45:12 +00001254
showarda5288b42009-07-28 20:06:08 +00001255 def save(self, *args, **kwargs):
showardddb90992009-02-11 23:39:32 +00001256 first_time = (self.id is None)
1257 if first_time:
1258 # see if this object was previously added and invalidated
1259 my_name = getattr(self, self.name_field)
1260 filters = {self.name_field : my_name, 'invalid' : True}
1261 try:
1262 old_object = self.__class__.objects.get(**filters)
showardafd97de2009-10-01 18:45:09 +00001263 self.resurrect_object(old_object)
showardddb90992009-02-11 23:39:32 +00001264 except self.DoesNotExist:
1265 # no existing object
1266 pass
showard7c785282008-05-29 19:45:12 +00001267
showarda5288b42009-07-28 20:06:08 +00001268 super(ModelWithInvalid, self).save(*args, **kwargs)
showard7c785282008-05-29 19:45:12 +00001269
1270
showardafd97de2009-10-01 18:45:09 +00001271 def resurrect_object(self, old_object):
1272 """
1273 Called when self is about to be saved for the first time and is actually
1274 "undeleting" a previously deleted object. Can be overridden by
1275 subclasses to copy data as desired from the deleted entry (but this
1276 superclass implementation must normally be called).
1277 """
1278 self.id = old_object.id
1279
1280
jadmanski0afbb632008-06-06 21:10:57 +00001281 def clean_object(self):
1282 """
1283 This method is called when an object is marked invalid.
1284 Subclasses should override this to clean up relationships that
showardafd97de2009-10-01 18:45:09 +00001285 should no longer exist if the object were deleted.
1286 """
jadmanski0afbb632008-06-06 21:10:57 +00001287 pass
showard7c785282008-05-29 19:45:12 +00001288
1289
jadmanski0afbb632008-06-06 21:10:57 +00001290 def delete(self):
Dale Curtis74a314b2011-06-23 14:55:46 -07001291 self.invalid = self.invalid
jadmanski0afbb632008-06-06 21:10:57 +00001292 assert not self.invalid
1293 self.invalid = True
1294 self.save()
1295 self.clean_object()
showard7c785282008-05-29 19:45:12 +00001296
1297
jadmanski0afbb632008-06-06 21:10:57 +00001298 @classmethod
1299 def get_valid_manager(cls):
1300 return cls.valid_objects
showard7c785282008-05-29 19:45:12 +00001301
1302
jadmanski0afbb632008-06-06 21:10:57 +00001303 class Manipulator(object):
1304 """
1305 Force default manipulators to look only at valid objects -
1306 otherwise they will match against invalid objects when checking
1307 uniqueness.
1308 """
1309 @classmethod
1310 def _prepare(cls, model):
1311 super(ModelWithInvalid.Manipulator, cls)._prepare(model)
1312 cls.manager = model.valid_objects
showardf8b19042009-05-12 17:22:49 +00001313
1314
1315class ModelWithAttributes(object):
1316 """
1317 Mixin class for models that have an attribute model associated with them.
1318 The attribute model is assumed to have its value field named "value".
1319 """
1320
1321 def _get_attribute_model_and_args(self, attribute):
1322 """
1323 Subclasses should override this to return a tuple (attribute_model,
1324 keyword_args), where attribute_model is a model class and keyword_args
1325 is a dict of args to pass to attribute_model.objects.get() to get an
1326 instance of the given attribute on this object.
1327 """
Dale Curtis74a314b2011-06-23 14:55:46 -07001328 raise NotImplementedError
showardf8b19042009-05-12 17:22:49 +00001329
1330
1331 def set_attribute(self, attribute, value):
1332 attribute_model, get_args = self._get_attribute_model_and_args(
1333 attribute)
1334 attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
1335 attribute_object.value = value
1336 attribute_object.save()
1337
1338
1339 def delete_attribute(self, attribute):
1340 attribute_model, get_args = self._get_attribute_model_and_args(
1341 attribute)
1342 try:
1343 attribute_model.objects.get(**get_args).delete()
showard16245422009-09-08 16:28:15 +00001344 except attribute_model.DoesNotExist:
showardf8b19042009-05-12 17:22:49 +00001345 pass
1346
1347
1348 def set_or_delete_attribute(self, attribute, value):
1349 if value is None:
1350 self.delete_attribute(attribute)
1351 else:
1352 self.set_attribute(attribute, value)
showard26b7ec72009-12-21 22:43:57 +00001353
1354
1355class ModelWithHashManager(dbmodels.Manager):
1356 """Manager for use with the ModelWithHash abstract model class"""
1357
1358 def create(self, **kwargs):
1359 raise Exception('ModelWithHash manager should use get_or_create() '
1360 'instead of create()')
1361
1362
1363 def get_or_create(self, **kwargs):
1364 kwargs['the_hash'] = self.model._compute_hash(**kwargs)
1365 return super(ModelWithHashManager, self).get_or_create(**kwargs)
1366
1367
1368class ModelWithHash(dbmodels.Model):
1369 """Superclass with methods for dealing with a hash column"""
1370
1371 the_hash = dbmodels.CharField(max_length=40, unique=True)
1372
1373 objects = ModelWithHashManager()
1374
1375 class Meta:
1376 abstract = True
1377
1378
1379 @classmethod
1380 def _compute_hash(cls, **kwargs):
1381 raise NotImplementedError('Subclasses must override _compute_hash()')
1382
1383
1384 def save(self, force_insert=False, **kwargs):
1385 """Prevents saving the model in most cases
1386
1387 We want these models to be immutable, so the generic save() operation
1388 will not work. These models should be instantiated through their the
1389 model.objects.get_or_create() method instead.
1390
1391 The exception is that save(force_insert=True) will be allowed, since
1392 that creates a new row. However, the preferred way to make instances of
1393 these models is through the get_or_create() method.
1394 """
1395 if not force_insert:
1396 # Allow a forced insert to happen; if it's a duplicate, the unique
1397 # constraint will catch it later anyways
1398 raise Exception('ModelWithHash is immutable')
1399 super(ModelWithHash, self).save(force_insert=force_insert, **kwargs)