blob: fd90f252cbf8e1b79d90251f5ef25b71a58457e1 [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 re
Michael Liang8864e862014-07-22 08:36:05 -07006import time
showarda5288b42009-07-28 20:06:08 +00007import django.core.exceptions
showard7c785282008-05-29 19:45:12 +00008from django.db import models as dbmodels, backend, connection
showarda5288b42009-07-28 20:06:08 +00009from django.db.models.sql import query
showard7e67b432010-01-20 01:13:04 +000010import django.db.models.sql.where
showard7c785282008-05-29 19:45:12 +000011from django.utils import datastructures
Michael Liang8864e862014-07-22 08:36:05 -070012from autotest_lib.client.common_lib.cros.graphite import es_utils
Prashanth B489b91d2014-03-15 12:17:16 -070013from autotest_lib.frontend.afe import rdb_model_extensions
showard56e93772008-10-06 10:06:22 +000014from autotest_lib.frontend.afe import readonly_connection
showard7c785282008-05-29 19:45:12 +000015
Prashanth B489b91d2014-03-15 12:17:16 -070016
17class ValidationError(django.core.exceptions.ValidationError):
jadmanski0afbb632008-06-06 21:10:57 +000018 """\
showarda5288b42009-07-28 20:06:08 +000019 Data validation error in adding or updating an object. The associated
jadmanski0afbb632008-06-06 21:10:57 +000020 value is a dictionary mapping field names to error strings.
21 """
showard7c785282008-05-29 19:45:12 +000022
23
showard09096d82008-07-07 23:20:49 +000024def _wrap_with_readonly(method):
mbligh1ef218d2009-08-03 16:57:56 +000025 def wrapper_method(*args, **kwargs):
26 readonly_connection.connection().set_django_connection()
27 try:
28 return method(*args, **kwargs)
29 finally:
30 readonly_connection.connection().unset_django_connection()
31 wrapper_method.__name__ = method.__name__
32 return wrapper_method
showard09096d82008-07-07 23:20:49 +000033
34
showarda5288b42009-07-28 20:06:08 +000035def _quote_name(name):
36 """Shorthand for connection.ops.quote_name()."""
37 return connection.ops.quote_name(name)
38
39
showard09096d82008-07-07 23:20:49 +000040def _wrap_generator_with_readonly(generator):
41 """
42 We have to wrap generators specially. Assume it performs
43 the query on the first call to next().
44 """
45 def wrapper_generator(*args, **kwargs):
46 generator_obj = generator(*args, **kwargs)
showard56e93772008-10-06 10:06:22 +000047 readonly_connection.connection().set_django_connection()
showard09096d82008-07-07 23:20:49 +000048 try:
49 first_value = generator_obj.next()
50 finally:
showard56e93772008-10-06 10:06:22 +000051 readonly_connection.connection().unset_django_connection()
showard09096d82008-07-07 23:20:49 +000052 yield first_value
53
54 while True:
55 yield generator_obj.next()
56
57 wrapper_generator.__name__ = generator.__name__
58 return wrapper_generator
59
60
61def _make_queryset_readonly(queryset):
62 """
63 Wrap all methods that do database queries with a readonly connection.
64 """
65 db_query_methods = ['count', 'get', 'get_or_create', 'latest', 'in_bulk',
66 'delete']
67 for method_name in db_query_methods:
68 method = getattr(queryset, method_name)
69 wrapped_method = _wrap_with_readonly(method)
70 setattr(queryset, method_name, wrapped_method)
71
72 queryset.iterator = _wrap_generator_with_readonly(queryset.iterator)
73
74
75class ReadonlyQuerySet(dbmodels.query.QuerySet):
76 """
77 QuerySet object that performs all database queries with the read-only
78 connection.
79 """
showarda5288b42009-07-28 20:06:08 +000080 def __init__(self, model=None, *args, **kwargs):
81 super(ReadonlyQuerySet, self).__init__(model, *args, **kwargs)
showard09096d82008-07-07 23:20:49 +000082 _make_queryset_readonly(self)
83
84
85 def values(self, *fields):
showarda5288b42009-07-28 20:06:08 +000086 return self._clone(klass=ReadonlyValuesQuerySet,
87 setup=True, _fields=fields)
showard09096d82008-07-07 23:20:49 +000088
89
90class ReadonlyValuesQuerySet(dbmodels.query.ValuesQuerySet):
showarda5288b42009-07-28 20:06:08 +000091 def __init__(self, model=None, *args, **kwargs):
92 super(ReadonlyValuesQuerySet, self).__init__(model, *args, **kwargs)
showard09096d82008-07-07 23:20:49 +000093 _make_queryset_readonly(self)
94
95
beepscc9fc702013-12-02 12:45:38 -080096class LeasedHostManager(dbmodels.Manager):
97 """Query manager for unleased, unlocked hosts.
98 """
99 def get_query_set(self):
100 return (super(LeasedHostManager, self).get_query_set().filter(
101 leased=0, locked=0))
102
103
showard7c785282008-05-29 19:45:12 +0000104class ExtendedManager(dbmodels.Manager):
jadmanski0afbb632008-06-06 21:10:57 +0000105 """\
106 Extended manager supporting subquery filtering.
107 """
showard7c785282008-05-29 19:45:12 +0000108
showardf828c772010-01-25 21:49:42 +0000109 class CustomQuery(query.Query):
showard7e67b432010-01-20 01:13:04 +0000110 def __init__(self, *args, **kwargs):
showardf828c772010-01-25 21:49:42 +0000111 super(ExtendedManager.CustomQuery, self).__init__(*args, **kwargs)
showard7e67b432010-01-20 01:13:04 +0000112 self._custom_joins = []
113
114
showarda5288b42009-07-28 20:06:08 +0000115 def clone(self, klass=None, **kwargs):
showardf828c772010-01-25 21:49:42 +0000116 obj = super(ExtendedManager.CustomQuery, self).clone(klass)
showard7e67b432010-01-20 01:13:04 +0000117 obj._custom_joins = list(self._custom_joins)
showarda5288b42009-07-28 20:06:08 +0000118 return obj
showard08f981b2008-06-24 21:59:03 +0000119
showard7e67b432010-01-20 01:13:04 +0000120
121 def combine(self, rhs, connector):
showardf828c772010-01-25 21:49:42 +0000122 super(ExtendedManager.CustomQuery, self).combine(rhs, connector)
showard7e67b432010-01-20 01:13:04 +0000123 if hasattr(rhs, '_custom_joins'):
124 self._custom_joins.extend(rhs._custom_joins)
125
126
127 def add_custom_join(self, table, condition, join_type,
128 condition_values=(), alias=None):
129 if alias is None:
130 alias = table
131 join_dict = dict(table=table,
132 condition=condition,
133 condition_values=condition_values,
134 join_type=join_type,
135 alias=alias)
136 self._custom_joins.append(join_dict)
137
138
showard7e67b432010-01-20 01:13:04 +0000139 @classmethod
140 def convert_query(self, query_set):
141 """
showardf828c772010-01-25 21:49:42 +0000142 Convert the query set's "query" attribute to a CustomQuery.
showard7e67b432010-01-20 01:13:04 +0000143 """
144 # Make a copy of the query set
145 query_set = query_set.all()
146 query_set.query = query_set.query.clone(
showardf828c772010-01-25 21:49:42 +0000147 klass=ExtendedManager.CustomQuery,
showard7e67b432010-01-20 01:13:04 +0000148 _custom_joins=[])
149 return query_set
showard43a3d262008-11-12 18:17:05 +0000150
151
showard7e67b432010-01-20 01:13:04 +0000152 class _WhereClause(object):
153 """Object allowing us to inject arbitrary SQL into Django queries.
showard43a3d262008-11-12 18:17:05 +0000154
showard7e67b432010-01-20 01:13:04 +0000155 By using this instead of extra(where=...), we can still freely combine
156 queries with & and |.
showarda5288b42009-07-28 20:06:08 +0000157 """
showard7e67b432010-01-20 01:13:04 +0000158 def __init__(self, clause, values=()):
159 self._clause = clause
160 self._values = values
showarda5288b42009-07-28 20:06:08 +0000161
showard7e67b432010-01-20 01:13:04 +0000162
Dale Curtis74a314b2011-06-23 14:55:46 -0700163 def as_sql(self, qn=None, connection=None):
showard7e67b432010-01-20 01:13:04 +0000164 return self._clause, self._values
165
166
167 def relabel_aliases(self, change_map):
168 return
showard43a3d262008-11-12 18:17:05 +0000169
170
showard8b0ea222009-12-23 19:23:03 +0000171 def add_join(self, query_set, join_table, join_key, join_condition='',
showard7e67b432010-01-20 01:13:04 +0000172 join_condition_values=(), join_from_key=None, alias=None,
173 suffix='', exclude=False, force_left_join=False):
174 """Add a join to query_set.
175
176 Join looks like this:
177 (INNER|LEFT) JOIN <join_table> AS <alias>
178 ON (<this table>.<join_from_key> = <join_table>.<join_key>
179 and <join_condition>)
180
showard0957a842009-05-11 19:25:08 +0000181 @param join_table table to join to
182 @param join_key field referencing back to this model to use for the join
183 @param join_condition extra condition for the ON clause of the join
showard7e67b432010-01-20 01:13:04 +0000184 @param join_condition_values values to substitute into join_condition
185 @param join_from_key column on this model to join from.
showard8b0ea222009-12-23 19:23:03 +0000186 @param alias alias to use for for join
187 @param suffix suffix to add to join_table for the join alias, if no
188 alias is provided
showard0957a842009-05-11 19:25:08 +0000189 @param exclude if true, exclude rows that match this join (will use a
showarda5288b42009-07-28 20:06:08 +0000190 LEFT OUTER JOIN and an appropriate WHERE condition)
showardc4780402009-08-31 18:31:34 +0000191 @param force_left_join - if true, a LEFT OUTER JOIN will be used
192 instead of an INNER JOIN regardless of other options
showard0957a842009-05-11 19:25:08 +0000193 """
showard7e67b432010-01-20 01:13:04 +0000194 join_from_table = query_set.model._meta.db_table
195 if join_from_key is None:
196 join_from_key = self.model._meta.pk.name
197 if alias is None:
198 alias = join_table + suffix
199 full_join_key = _quote_name(alias) + '.' + _quote_name(join_key)
200 full_join_condition = '%s = %s.%s' % (full_join_key,
201 _quote_name(join_from_table),
202 _quote_name(join_from_key))
showard43a3d262008-11-12 18:17:05 +0000203 if join_condition:
204 full_join_condition += ' AND (' + join_condition + ')'
205 if exclude or force_left_join:
showarda5288b42009-07-28 20:06:08 +0000206 join_type = query_set.query.LOUTER
showard43a3d262008-11-12 18:17:05 +0000207 else:
showarda5288b42009-07-28 20:06:08 +0000208 join_type = query_set.query.INNER
showard43a3d262008-11-12 18:17:05 +0000209
showardf828c772010-01-25 21:49:42 +0000210 query_set = self.CustomQuery.convert_query(query_set)
showard7e67b432010-01-20 01:13:04 +0000211 query_set.query.add_custom_join(join_table,
212 full_join_condition,
213 join_type,
214 condition_values=join_condition_values,
215 alias=alias)
showard43a3d262008-11-12 18:17:05 +0000216
showard7e67b432010-01-20 01:13:04 +0000217 if exclude:
218 query_set = query_set.extra(where=[full_join_key + ' IS NULL'])
219
220 return query_set
221
222
223 def _info_for_many_to_one_join(self, field, join_to_query, alias):
224 """
225 @param field: the ForeignKey field on the related model
226 @param join_to_query: the query over the related model that we're
227 joining to
228 @param alias: alias of joined table
229 """
230 info = {}
231 rhs_table = join_to_query.model._meta.db_table
232 info['rhs_table'] = rhs_table
233 info['rhs_column'] = field.column
234 info['lhs_column'] = field.rel.get_related_field().column
235 rhs_where = join_to_query.query.where
236 rhs_where.relabel_aliases({rhs_table: alias})
Dale Curtis74a314b2011-06-23 14:55:46 -0700237 compiler = join_to_query.query.get_compiler(using=join_to_query.db)
238 initial_clause, values = compiler.as_sql()
239 all_clauses = (initial_clause,)
240 if hasattr(join_to_query.query, 'extra_where'):
241 all_clauses += join_to_query.query.extra_where
242 info['where_clause'] = (
243 ' AND '.join('(%s)' % clause for clause in all_clauses))
showard7e67b432010-01-20 01:13:04 +0000244 info['values'] = values
245 return info
246
247
248 def _info_for_many_to_many_join(self, m2m_field, join_to_query, alias,
249 m2m_is_on_this_model):
250 """
251 @param m2m_field: a Django field representing the M2M relationship.
252 It uses a pivot table with the following structure:
253 this model table <---> M2M pivot table <---> joined model table
254 @param join_to_query: the query over the related model that we're
255 joining to.
256 @param alias: alias of joined table
257 """
258 if m2m_is_on_this_model:
259 # referenced field on this model
260 lhs_id_field = self.model._meta.pk
261 # foreign key on the pivot table referencing lhs_id_field
262 m2m_lhs_column = m2m_field.m2m_column_name()
263 # foreign key on the pivot table referencing rhd_id_field
264 m2m_rhs_column = m2m_field.m2m_reverse_name()
265 # referenced field on related model
266 rhs_id_field = m2m_field.rel.get_related_field()
267 else:
268 lhs_id_field = m2m_field.rel.get_related_field()
269 m2m_lhs_column = m2m_field.m2m_reverse_name()
270 m2m_rhs_column = m2m_field.m2m_column_name()
271 rhs_id_field = join_to_query.model._meta.pk
272
273 info = {}
274 info['rhs_table'] = m2m_field.m2m_db_table()
275 info['rhs_column'] = m2m_lhs_column
276 info['lhs_column'] = lhs_id_field.column
277
278 # select the ID of related models relevant to this join. we can only do
279 # a single join, so we need to gather this information up front and
280 # include it in the join condition.
281 rhs_ids = join_to_query.values_list(rhs_id_field.attname, flat=True)
282 assert len(rhs_ids) == 1, ('Many-to-many custom field joins can only '
283 'match a single related object.')
284 rhs_id = rhs_ids[0]
285
286 info['where_clause'] = '%s.%s = %s' % (_quote_name(alias),
287 _quote_name(m2m_rhs_column),
288 rhs_id)
289 info['values'] = ()
290 return info
291
292
293 def join_custom_field(self, query_set, join_to_query, alias,
294 left_join=True):
295 """Join to a related model to create a custom field in the given query.
296
297 This method is used to construct a custom field on the given query based
298 on a many-valued relationsip. join_to_query should be a simple query
299 (no joins) on the related model which returns at most one related row
300 per instance of this model.
301
302 For many-to-one relationships, the joined table contains the matching
303 row from the related model it one is related, NULL otherwise.
304
305 For many-to-many relationships, the joined table contains the matching
306 row if it's related, NULL otherwise.
307 """
308 relationship_type, field = self.determine_relationship(
309 join_to_query.model)
310
311 if relationship_type == self.MANY_TO_ONE:
312 info = self._info_for_many_to_one_join(field, join_to_query, alias)
313 elif relationship_type == self.M2M_ON_RELATED_MODEL:
314 info = self._info_for_many_to_many_join(
315 m2m_field=field, join_to_query=join_to_query, alias=alias,
316 m2m_is_on_this_model=False)
317 elif relationship_type ==self.M2M_ON_THIS_MODEL:
318 info = self._info_for_many_to_many_join(
319 m2m_field=field, join_to_query=join_to_query, alias=alias,
320 m2m_is_on_this_model=True)
321
322 return self.add_join(query_set, info['rhs_table'], info['rhs_column'],
323 join_from_key=info['lhs_column'],
324 join_condition=info['where_clause'],
325 join_condition_values=info['values'],
326 alias=alias,
327 force_left_join=left_join)
328
329
showardf828c772010-01-25 21:49:42 +0000330 def key_on_joined_table(self, join_to_query):
331 """Get a non-null column on the table joined for the given query.
332
333 This analyzes the join that would be produced if join_to_query were
334 passed to join_custom_field.
335 """
336 relationship_type, field = self.determine_relationship(
337 join_to_query.model)
338 if relationship_type == self.MANY_TO_ONE:
339 return join_to_query.model._meta.pk.column
340 return field.m2m_column_name() # any column on the M2M table will do
341
342
showard7e67b432010-01-20 01:13:04 +0000343 def add_where(self, query_set, where, values=()):
344 query_set = query_set.all()
345 query_set.query.where.add(self._WhereClause(where, values),
346 django.db.models.sql.where.AND)
showardc4780402009-08-31 18:31:34 +0000347 return query_set
showard7c785282008-05-29 19:45:12 +0000348
349
showardeaccf8f2009-04-16 03:11:33 +0000350 def _get_quoted_field(self, table, field):
showarda5288b42009-07-28 20:06:08 +0000351 return _quote_name(table) + '.' + _quote_name(field)
showard5ef36e92008-07-02 16:37:09 +0000352
353
showard7c199df2008-10-03 10:17:15 +0000354 def get_key_on_this_table(self, key_field=None):
showard5ef36e92008-07-02 16:37:09 +0000355 if key_field is None:
356 # default to primary key
357 key_field = self.model._meta.pk.column
358 return self._get_quoted_field(self.model._meta.db_table, key_field)
359
360
showardeaccf8f2009-04-16 03:11:33 +0000361 def escape_user_sql(self, sql):
362 return sql.replace('%', '%%')
363
showard5ef36e92008-07-02 16:37:09 +0000364
showard0957a842009-05-11 19:25:08 +0000365 def _custom_select_query(self, query_set, selects):
Dale Curtis74a314b2011-06-23 14:55:46 -0700366 compiler = query_set.query.get_compiler(using=query_set.db)
367 sql, params = compiler.as_sql()
showarda5288b42009-07-28 20:06:08 +0000368 from_ = sql[sql.find(' FROM'):]
369
370 if query_set.query.distinct:
showard0957a842009-05-11 19:25:08 +0000371 distinct = 'DISTINCT '
372 else:
373 distinct = ''
showarda5288b42009-07-28 20:06:08 +0000374
375 sql_query = ('SELECT ' + distinct + ','.join(selects) + from_)
showard0957a842009-05-11 19:25:08 +0000376 cursor = readonly_connection.connection().cursor()
377 cursor.execute(sql_query, params)
378 return cursor.fetchall()
379
380
showard68693f72009-05-20 00:31:53 +0000381 def _is_relation_to(self, field, model_class):
382 return field.rel and field.rel.to is model_class
showard0957a842009-05-11 19:25:08 +0000383
384
showard7e67b432010-01-20 01:13:04 +0000385 MANY_TO_ONE = object()
386 M2M_ON_RELATED_MODEL = object()
387 M2M_ON_THIS_MODEL = object()
388
389 def determine_relationship(self, related_model):
390 """
391 Determine the relationship between this model and related_model.
392
393 related_model must have some sort of many-valued relationship to this
394 manager's model.
395 @returns (relationship_type, field), where relationship_type is one of
396 MANY_TO_ONE, M2M_ON_RELATED_MODEL, M2M_ON_THIS_MODEL, and field
397 is the Django field object for the relationship.
398 """
399 # look for a foreign key field on related_model relating to this model
400 for field in related_model._meta.fields:
401 if self._is_relation_to(field, self.model):
402 return self.MANY_TO_ONE, field
403
404 # look for an M2M field on related_model relating to this model
405 for field in related_model._meta.many_to_many:
406 if self._is_relation_to(field, self.model):
407 return self.M2M_ON_RELATED_MODEL, field
408
409 # maybe this model has the many-to-many field
410 for field in self.model._meta.many_to_many:
411 if self._is_relation_to(field, related_model):
412 return self.M2M_ON_THIS_MODEL, field
413
414 raise ValueError('%s has no relation to %s' %
415 (related_model, self.model))
416
417
showard68693f72009-05-20 00:31:53 +0000418 def _get_pivot_iterator(self, base_objects_by_id, related_model):
showard0957a842009-05-11 19:25:08 +0000419 """
showard68693f72009-05-20 00:31:53 +0000420 Determine the relationship between this model and related_model, and
421 return a pivot iterator.
422 @param base_objects_by_id: dict of instances of this model indexed by
423 their IDs
424 @returns a pivot iterator, which yields a tuple (base_object,
425 related_object) for each relationship between a base object and a
426 related object. all base_object instances come from base_objects_by_id.
showard7e67b432010-01-20 01:13:04 +0000427 Note -- this depends on Django model internals.
showard0957a842009-05-11 19:25:08 +0000428 """
showard7e67b432010-01-20 01:13:04 +0000429 relationship_type, field = self.determine_relationship(related_model)
430 if relationship_type == self.MANY_TO_ONE:
431 return self._many_to_one_pivot(base_objects_by_id,
432 related_model, field)
433 elif relationship_type == self.M2M_ON_RELATED_MODEL:
434 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000435 base_objects_by_id, related_model, field.m2m_db_table(),
436 field.m2m_reverse_name(), field.m2m_column_name())
showard7e67b432010-01-20 01:13:04 +0000437 else:
438 assert relationship_type == self.M2M_ON_THIS_MODEL
439 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000440 base_objects_by_id, related_model, field.m2m_db_table(),
441 field.m2m_column_name(), field.m2m_reverse_name())
showard0957a842009-05-11 19:25:08 +0000442
showard0957a842009-05-11 19:25:08 +0000443
showard68693f72009-05-20 00:31:53 +0000444 def _many_to_one_pivot(self, base_objects_by_id, related_model,
445 foreign_key_field):
446 """
447 @returns a pivot iterator - see _get_pivot_iterator()
448 """
449 filter_data = {foreign_key_field.name + '__pk__in':
450 base_objects_by_id.keys()}
451 for related_object in related_model.objects.filter(**filter_data):
showarda5a72c92009-08-20 23:35:21 +0000452 # lookup base object in the dict, rather than grabbing it from the
453 # related object. we need to return instances from the dict, not
454 # fresh instances of the same models (and grabbing model instances
455 # from the related models incurs a DB query each time).
456 base_object_id = getattr(related_object, foreign_key_field.attname)
457 base_object = base_objects_by_id[base_object_id]
showard68693f72009-05-20 00:31:53 +0000458 yield base_object, related_object
459
460
461 def _query_pivot_table(self, base_objects_by_id, pivot_table,
462 pivot_from_field, pivot_to_field):
showard0957a842009-05-11 19:25:08 +0000463 """
464 @param id_list list of IDs of self.model objects to include
465 @param pivot_table the name of the pivot table
466 @param pivot_from_field a field name on pivot_table referencing
467 self.model
468 @param pivot_to_field a field name on pivot_table referencing the
469 related model.
showard68693f72009-05-20 00:31:53 +0000470 @returns pivot list of IDs (base_id, related_id)
showard0957a842009-05-11 19:25:08 +0000471 """
472 query = """
473 SELECT %(from_field)s, %(to_field)s
474 FROM %(table)s
475 WHERE %(from_field)s IN (%(id_list)s)
476 """ % dict(from_field=pivot_from_field,
477 to_field=pivot_to_field,
478 table=pivot_table,
showard68693f72009-05-20 00:31:53 +0000479 id_list=','.join(str(id_) for id_
480 in base_objects_by_id.iterkeys()))
showard0957a842009-05-11 19:25:08 +0000481 cursor = readonly_connection.connection().cursor()
482 cursor.execute(query)
showard68693f72009-05-20 00:31:53 +0000483 return cursor.fetchall()
showard0957a842009-05-11 19:25:08 +0000484
485
showard68693f72009-05-20 00:31:53 +0000486 def _many_to_many_pivot(self, base_objects_by_id, related_model,
487 pivot_table, pivot_from_field, pivot_to_field):
488 """
489 @param pivot_table: see _query_pivot_table
490 @param pivot_from_field: see _query_pivot_table
491 @param pivot_to_field: see _query_pivot_table
492 @returns a pivot iterator - see _get_pivot_iterator()
493 """
494 id_pivot = self._query_pivot_table(base_objects_by_id, pivot_table,
495 pivot_from_field, pivot_to_field)
496
497 all_related_ids = list(set(related_id for base_id, related_id
498 in id_pivot))
499 related_objects_by_id = related_model.objects.in_bulk(all_related_ids)
500
501 for base_id, related_id in id_pivot:
502 yield base_objects_by_id[base_id], related_objects_by_id[related_id]
503
504
505 def populate_relationships(self, base_objects, related_model,
showard0957a842009-05-11 19:25:08 +0000506 related_list_name):
507 """
showard68693f72009-05-20 00:31:53 +0000508 For each instance of this model in base_objects, add a field named
509 related_list_name listing all the related objects of type related_model.
510 related_model must be in a many-to-one or many-to-many relationship with
511 this model.
512 @param base_objects - list of instances of this model
513 @param related_model - model class related to this model
514 @param related_list_name - attribute name in which to store the related
515 object list.
showard0957a842009-05-11 19:25:08 +0000516 """
showard68693f72009-05-20 00:31:53 +0000517 if not base_objects:
showard0957a842009-05-11 19:25:08 +0000518 # if we don't bail early, we'll get a SQL error later
519 return
showard0957a842009-05-11 19:25:08 +0000520
showard68693f72009-05-20 00:31:53 +0000521 base_objects_by_id = dict((base_object._get_pk_val(), base_object)
522 for base_object in base_objects)
523 pivot_iterator = self._get_pivot_iterator(base_objects_by_id,
524 related_model)
showard0957a842009-05-11 19:25:08 +0000525
showard68693f72009-05-20 00:31:53 +0000526 for base_object in base_objects:
527 setattr(base_object, related_list_name, [])
528
529 for base_object, related_object in pivot_iterator:
530 getattr(base_object, related_list_name).append(related_object)
showard0957a842009-05-11 19:25:08 +0000531
532
jamesrene3656232010-03-02 00:00:30 +0000533class ModelWithInvalidQuerySet(dbmodels.query.QuerySet):
534 """
535 QuerySet that handles delete() properly for models with an "invalid" bit
536 """
537 def delete(self):
538 for model in self:
539 model.delete()
540
541
542class ModelWithInvalidManager(ExtendedManager):
543 """
544 Manager for objects with an "invalid" bit
545 """
546 def get_query_set(self):
547 return ModelWithInvalidQuerySet(self.model)
548
549
550class ValidObjectsManager(ModelWithInvalidManager):
jadmanski0afbb632008-06-06 21:10:57 +0000551 """
552 Manager returning only objects with invalid=False.
553 """
554 def get_query_set(self):
555 queryset = super(ValidObjectsManager, self).get_query_set()
556 return queryset.filter(invalid=False)
showard7c785282008-05-29 19:45:12 +0000557
558
Prashanth B489b91d2014-03-15 12:17:16 -0700559class ModelExtensions(rdb_model_extensions.ModelValidators):
jadmanski0afbb632008-06-06 21:10:57 +0000560 """\
Prashanth B489b91d2014-03-15 12:17:16 -0700561 Mixin with convenience functions for models, built on top of
562 the model validators in rdb_model_extensions.
jadmanski0afbb632008-06-06 21:10:57 +0000563 """
564 # TODO: at least some of these functions really belong in a custom
565 # Manager class
showard7c785282008-05-29 19:45:12 +0000566
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700567
568 SERIALIZATION_LINKS_TO_FOLLOW = set()
569 """
570 To be able to send jobs and hosts to shards, it's necessary to find their
571 dependencies.
572 The most generic approach for this would be to traverse all relationships
573 to other objects recursively. This would list all objects that are related
574 in any way.
575 But this approach finds too many objects: If a host should be transferred,
576 all it's relationships would be traversed. This would find an acl group.
577 If then the acl group's relationships are traversed, the relationship
578 would be followed backwards and many other hosts would be found.
579
580 This mapping tells that algorithm which relations to follow explicitly.
581 """
582
Jakob Juelichf865d332014-09-29 10:47:49 -0700583
584 SERIALIZATION_LOCAL_LINKS_TO_UPDATE = set()
585 """
586 On deserializion, if the object to persist already exists, local fields
587 will only be updated, if their name is in this set.
588 """
589
590
jadmanski0afbb632008-06-06 21:10:57 +0000591 @classmethod
592 def convert_human_readable_values(cls, data, to_human_readable=False):
593 """\
594 Performs conversions on user-supplied field data, to make it
595 easier for users to pass human-readable data.
showard7c785282008-05-29 19:45:12 +0000596
jadmanski0afbb632008-06-06 21:10:57 +0000597 For all fields that have choice sets, convert their values
598 from human-readable strings to enum values, if necessary. This
599 allows users to pass strings instead of the corresponding
600 integer values.
showard7c785282008-05-29 19:45:12 +0000601
jadmanski0afbb632008-06-06 21:10:57 +0000602 For all foreign key fields, call smart_get with the supplied
603 data. This allows the user to pass either an ID value or
604 the name of the object as a string.
showard7c785282008-05-29 19:45:12 +0000605
jadmanski0afbb632008-06-06 21:10:57 +0000606 If to_human_readable=True, perform the inverse - i.e. convert
607 numeric values to human readable values.
showard7c785282008-05-29 19:45:12 +0000608
jadmanski0afbb632008-06-06 21:10:57 +0000609 This method modifies data in-place.
610 """
611 field_dict = cls.get_field_dict()
612 for field_name in data:
showarde732ee72008-09-23 19:15:43 +0000613 if field_name not in field_dict or data[field_name] is None:
jadmanski0afbb632008-06-06 21:10:57 +0000614 continue
615 field_obj = field_dict[field_name]
616 # convert enum values
617 if field_obj.choices:
618 for choice_data in field_obj.choices:
619 # choice_data is (value, name)
620 if to_human_readable:
621 from_val, to_val = choice_data
622 else:
623 to_val, from_val = choice_data
624 if from_val == data[field_name]:
625 data[field_name] = to_val
626 break
627 # convert foreign key values
628 elif field_obj.rel:
showarda4ea5742009-02-17 20:56:23 +0000629 dest_obj = field_obj.rel.to.smart_get(data[field_name],
630 valid_only=False)
showardf8b19042009-05-12 17:22:49 +0000631 if to_human_readable:
Paul Pendlebury5a8c6ad2011-02-01 07:20:17 -0800632 # parameterized_jobs do not have a name_field
633 if (field_name != 'parameterized_job' and
634 dest_obj.name_field is not None):
showardf8b19042009-05-12 17:22:49 +0000635 data[field_name] = getattr(dest_obj,
636 dest_obj.name_field)
jadmanski0afbb632008-06-06 21:10:57 +0000637 else:
showardb0a73032009-03-27 18:35:41 +0000638 data[field_name] = dest_obj
showard7c785282008-05-29 19:45:12 +0000639
640
showard7c785282008-05-29 19:45:12 +0000641
642
Dale Curtis74a314b2011-06-23 14:55:46 -0700643 def _validate_unique(self):
jadmanski0afbb632008-06-06 21:10:57 +0000644 """\
645 Validate that unique fields are unique. Django manipulators do
646 this too, but they're a huge pain to use manually. Trust me.
647 """
648 errors = {}
649 cls = type(self)
650 field_dict = self.get_field_dict()
651 manager = cls.get_valid_manager()
652 for field_name, field_obj in field_dict.iteritems():
653 if not field_obj.unique:
654 continue
showard7c785282008-05-29 19:45:12 +0000655
jadmanski0afbb632008-06-06 21:10:57 +0000656 value = getattr(self, field_name)
showardbd18ab72009-09-18 21:20:27 +0000657 if value is None and field_obj.auto_created:
658 # don't bother checking autoincrement fields about to be
659 # generated
660 continue
661
jadmanski0afbb632008-06-06 21:10:57 +0000662 existing_objs = manager.filter(**{field_name : value})
663 num_existing = existing_objs.count()
showard7c785282008-05-29 19:45:12 +0000664
jadmanski0afbb632008-06-06 21:10:57 +0000665 if num_existing == 0:
666 continue
667 if num_existing == 1 and existing_objs[0].id == self.id:
668 continue
669 errors[field_name] = (
670 'This value must be unique (%s)' % (value))
671 return errors
showard7c785282008-05-29 19:45:12 +0000672
673
showarda5288b42009-07-28 20:06:08 +0000674 def _validate(self):
675 """
676 First coerces all fields on this instance to their proper Python types.
677 Then runs validation on every field. Returns a dictionary of
678 field_name -> error_list.
679
680 Based on validate() from django.db.models.Model in Django 0.96, which
681 was removed in Django 1.0. It should reappear in a later version. See:
682 http://code.djangoproject.com/ticket/6845
683 """
684 error_dict = {}
685 for f in self._meta.fields:
686 try:
687 python_value = f.to_python(
688 getattr(self, f.attname, f.get_default()))
689 except django.core.exceptions.ValidationError, e:
jamesren1e0a4ce2010-04-21 17:45:11 +0000690 error_dict[f.name] = str(e)
showarda5288b42009-07-28 20:06:08 +0000691 continue
692
693 if not f.blank and not python_value:
694 error_dict[f.name] = 'This field is required.'
695 continue
696
697 setattr(self, f.attname, python_value)
698
699 return error_dict
700
701
jadmanski0afbb632008-06-06 21:10:57 +0000702 def do_validate(self):
showarda5288b42009-07-28 20:06:08 +0000703 errors = self._validate()
Dale Curtis74a314b2011-06-23 14:55:46 -0700704 unique_errors = self._validate_unique()
jadmanski0afbb632008-06-06 21:10:57 +0000705 for field_name, error in unique_errors.iteritems():
706 errors.setdefault(field_name, error)
707 if errors:
708 raise ValidationError(errors)
showard7c785282008-05-29 19:45:12 +0000709
710
jadmanski0afbb632008-06-06 21:10:57 +0000711 # actually (externally) useful methods follow
showard7c785282008-05-29 19:45:12 +0000712
jadmanski0afbb632008-06-06 21:10:57 +0000713 @classmethod
714 def add_object(cls, data={}, **kwargs):
715 """\
716 Returns a new object created with the given data (a dictionary
717 mapping field names to values). Merges any extra keyword args
718 into data.
719 """
Prashanth B489b91d2014-03-15 12:17:16 -0700720 data = dict(data)
721 data.update(kwargs)
722 data = cls.prepare_data_args(data)
723 cls.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000724 data = cls.provide_default_values(data)
Prashanth B489b91d2014-03-15 12:17:16 -0700725
jadmanski0afbb632008-06-06 21:10:57 +0000726 obj = cls(**data)
727 obj.do_validate()
728 obj.save()
729 return obj
showard7c785282008-05-29 19:45:12 +0000730
Michael Liang8864e862014-07-22 08:36:05 -0700731 def record_state(self, type_str, state, value):
732 """Record metadata in elasticsearch.
733
734 @param type_str: sets the _type field in elasticsearch db.
735 @param state: string representing what state we are recording,
736 e.g. 'locked'
737 @param value: value of the state, e.g. True
738 """
739 metadata = {
Dan Shie4cb9e22014-08-29 11:56:04 -0700740 state: value,
Michael Liang8864e862014-07-22 08:36:05 -0700741 'hostname': self.hostname,
742 }
743 es_utils.ESMetadata().post(type_str=type_str, metadata=metadata)
744
showard7c785282008-05-29 19:45:12 +0000745
jadmanski0afbb632008-06-06 21:10:57 +0000746 def update_object(self, data={}, **kwargs):
747 """\
748 Updates the object with the given data (a dictionary mapping
749 field names to values). Merges any extra keyword args into
750 data.
751 """
Prashanth B489b91d2014-03-15 12:17:16 -0700752 data = dict(data)
753 data.update(kwargs)
754 data = self.prepare_data_args(data)
755 self.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000756 for field_name, value in data.iteritems():
showardb0a73032009-03-27 18:35:41 +0000757 setattr(self, field_name, value)
Michael Liang8864e862014-07-22 08:36:05 -0700758 # Other fields such as label (when updated) are not sent over
759 # the es because it doesn't contribute to putting together host
760 # host history. Locks are important in host history because if
761 # a device is locked then we don't really care what state it is in.
762 if field_name == 'locked':
763 self.record_state('lock_history', 'locked', value)
jadmanski0afbb632008-06-06 21:10:57 +0000764 self.do_validate()
765 self.save()
showard7c785282008-05-29 19:45:12 +0000766
767
showard8bfb5cb2009-10-07 20:49:15 +0000768 # see query_objects()
769 _SPECIAL_FILTER_KEYS = ('query_start', 'query_limit', 'sort_by',
770 'extra_args', 'extra_where', 'no_distinct')
771
772
jadmanski0afbb632008-06-06 21:10:57 +0000773 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000774 def _extract_special_params(cls, filter_data):
775 """
776 @returns a tuple of dicts (special_params, regular_filters), where
777 special_params contains the parameters we handle specially and
778 regular_filters is the remaining data to be handled by Django.
779 """
780 regular_filters = dict(filter_data)
781 special_params = {}
782 for key in cls._SPECIAL_FILTER_KEYS:
783 if key in regular_filters:
784 special_params[key] = regular_filters.pop(key)
785 return special_params, regular_filters
786
787
788 @classmethod
789 def apply_presentation(cls, query, filter_data):
790 """
791 Apply presentation parameters -- sorting and paging -- to the given
792 query.
793 @returns new query with presentation applied
794 """
795 special_params, _ = cls._extract_special_params(filter_data)
796 sort_by = special_params.get('sort_by', None)
797 if sort_by:
798 assert isinstance(sort_by, list) or isinstance(sort_by, tuple)
showard8b0ea222009-12-23 19:23:03 +0000799 query = query.extra(order_by=sort_by)
showard8bfb5cb2009-10-07 20:49:15 +0000800
801 query_start = special_params.get('query_start', None)
802 query_limit = special_params.get('query_limit', None)
803 if query_start is not None:
804 if query_limit is None:
805 raise ValueError('Cannot pass query_start without query_limit')
806 # query_limit is passed as a page size
showard7074b742009-10-12 20:30:04 +0000807 query_limit += query_start
808 return query[query_start:query_limit]
showard8bfb5cb2009-10-07 20:49:15 +0000809
810
811 @classmethod
812 def query_objects(cls, filter_data, valid_only=True, initial_query=None,
813 apply_presentation=True):
jadmanski0afbb632008-06-06 21:10:57 +0000814 """\
815 Returns a QuerySet object for querying the given model_class
816 with the given filter_data. Optional special arguments in
817 filter_data include:
818 -query_start: index of first return to return
819 -query_limit: maximum number of results to return
820 -sort_by: list of fields to sort on. prefixing a '-' onto a
821 field name changes the sort to descending order.
822 -extra_args: keyword args to pass to query.extra() (see Django
823 DB layer documentation)
showarda5288b42009-07-28 20:06:08 +0000824 -extra_where: extra WHERE clause to append
showard8bfb5cb2009-10-07 20:49:15 +0000825 -no_distinct: if True, a DISTINCT will not be added to the SELECT
jadmanski0afbb632008-06-06 21:10:57 +0000826 """
showard8bfb5cb2009-10-07 20:49:15 +0000827 special_params, regular_filters = cls._extract_special_params(
828 filter_data)
showard7c785282008-05-29 19:45:12 +0000829
showard7ac7b7a2008-07-21 20:24:29 +0000830 if initial_query is None:
831 if valid_only:
832 initial_query = cls.get_valid_manager()
833 else:
834 initial_query = cls.objects
showard8bfb5cb2009-10-07 20:49:15 +0000835
836 query = initial_query.filter(**regular_filters)
837
838 use_distinct = not special_params.get('no_distinct', False)
showard7ac7b7a2008-07-21 20:24:29 +0000839 if use_distinct:
840 query = query.distinct()
showard7c785282008-05-29 19:45:12 +0000841
showard8bfb5cb2009-10-07 20:49:15 +0000842 extra_args = special_params.get('extra_args', {})
843 extra_where = special_params.get('extra_where', None)
844 if extra_where:
845 # escape %'s
846 extra_where = cls.objects.escape_user_sql(extra_where)
847 extra_args.setdefault('where', []).append(extra_where)
jadmanski0afbb632008-06-06 21:10:57 +0000848 if extra_args:
849 query = query.extra(**extra_args)
showard09096d82008-07-07 23:20:49 +0000850 query = query._clone(klass=ReadonlyQuerySet)
showard7c785282008-05-29 19:45:12 +0000851
showard8bfb5cb2009-10-07 20:49:15 +0000852 if apply_presentation:
853 query = cls.apply_presentation(query, filter_data)
854
855 return query
showard7c785282008-05-29 19:45:12 +0000856
857
jadmanski0afbb632008-06-06 21:10:57 +0000858 @classmethod
showard585c2ab2008-07-23 19:29:49 +0000859 def query_count(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000860 """\
861 Like query_objects, but retreive only the count of results.
862 """
863 filter_data.pop('query_start', None)
864 filter_data.pop('query_limit', None)
showard585c2ab2008-07-23 19:29:49 +0000865 query = cls.query_objects(filter_data, initial_query=initial_query)
866 return query.count()
showard7c785282008-05-29 19:45:12 +0000867
868
jadmanski0afbb632008-06-06 21:10:57 +0000869 @classmethod
870 def clean_object_dicts(cls, field_dicts):
871 """\
872 Take a list of dicts corresponding to object (as returned by
873 query.values()) and clean the data to be more suitable for
874 returning to the user.
875 """
showarde732ee72008-09-23 19:15:43 +0000876 for field_dict in field_dicts:
877 cls.clean_foreign_keys(field_dict)
showard21baa452008-10-21 00:08:39 +0000878 cls._convert_booleans(field_dict)
showarde732ee72008-09-23 19:15:43 +0000879 cls.convert_human_readable_values(field_dict,
880 to_human_readable=True)
showard7c785282008-05-29 19:45:12 +0000881
882
jadmanski0afbb632008-06-06 21:10:57 +0000883 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000884 def list_objects(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000885 """\
886 Like query_objects, but return a list of dictionaries.
887 """
showard7ac7b7a2008-07-21 20:24:29 +0000888 query = cls.query_objects(filter_data, initial_query=initial_query)
showard8bfb5cb2009-10-07 20:49:15 +0000889 extra_fields = query.query.extra_select.keys()
890 field_dicts = [model_object.get_object_dict(extra_fields=extra_fields)
showarde732ee72008-09-23 19:15:43 +0000891 for model_object in query]
jadmanski0afbb632008-06-06 21:10:57 +0000892 return field_dicts
showard7c785282008-05-29 19:45:12 +0000893
894
jadmanski0afbb632008-06-06 21:10:57 +0000895 @classmethod
showarda4ea5742009-02-17 20:56:23 +0000896 def smart_get(cls, id_or_name, valid_only=True):
jadmanski0afbb632008-06-06 21:10:57 +0000897 """\
898 smart_get(integer) -> get object by ID
899 smart_get(string) -> get object by name_field
jadmanski0afbb632008-06-06 21:10:57 +0000900 """
showarda4ea5742009-02-17 20:56:23 +0000901 if valid_only:
902 manager = cls.get_valid_manager()
903 else:
904 manager = cls.objects
905
906 if isinstance(id_or_name, (int, long)):
907 return manager.get(pk=id_or_name)
jamesren3e9f6092010-03-11 21:32:10 +0000908 if isinstance(id_or_name, basestring) and hasattr(cls, 'name_field'):
showarda4ea5742009-02-17 20:56:23 +0000909 return manager.get(**{cls.name_field : id_or_name})
910 raise ValueError(
911 'Invalid positional argument: %s (%s)' % (id_or_name,
912 type(id_or_name)))
showard7c785282008-05-29 19:45:12 +0000913
914
showardbe3ec042008-11-12 18:16:07 +0000915 @classmethod
916 def smart_get_bulk(cls, id_or_name_list):
917 invalid_inputs = []
918 result_objects = []
919 for id_or_name in id_or_name_list:
920 try:
921 result_objects.append(cls.smart_get(id_or_name))
922 except cls.DoesNotExist:
923 invalid_inputs.append(id_or_name)
924 if invalid_inputs:
mbligh7a3ebe32008-12-01 17:10:33 +0000925 raise cls.DoesNotExist('The following %ss do not exist: %s'
926 % (cls.__name__.lower(),
927 ', '.join(invalid_inputs)))
showardbe3ec042008-11-12 18:16:07 +0000928 return result_objects
929
930
showard8bfb5cb2009-10-07 20:49:15 +0000931 def get_object_dict(self, extra_fields=None):
jadmanski0afbb632008-06-06 21:10:57 +0000932 """\
showard8bfb5cb2009-10-07 20:49:15 +0000933 Return a dictionary mapping fields to this object's values. @param
934 extra_fields: list of extra attribute names to include, in addition to
935 the fields defined on this object.
jadmanski0afbb632008-06-06 21:10:57 +0000936 """
showard8bfb5cb2009-10-07 20:49:15 +0000937 fields = self.get_field_dict().keys()
938 if extra_fields:
939 fields += extra_fields
jadmanski0afbb632008-06-06 21:10:57 +0000940 object_dict = dict((field_name, getattr(self, field_name))
showarde732ee72008-09-23 19:15:43 +0000941 for field_name in fields)
jadmanski0afbb632008-06-06 21:10:57 +0000942 self.clean_object_dicts([object_dict])
showardd3dc1992009-04-22 21:01:40 +0000943 self._postprocess_object_dict(object_dict)
jadmanski0afbb632008-06-06 21:10:57 +0000944 return object_dict
showard7c785282008-05-29 19:45:12 +0000945
946
showardd3dc1992009-04-22 21:01:40 +0000947 def _postprocess_object_dict(self, object_dict):
948 """For subclasses to override."""
949 pass
950
951
jadmanski0afbb632008-06-06 21:10:57 +0000952 @classmethod
953 def get_valid_manager(cls):
954 return cls.objects
showard7c785282008-05-29 19:45:12 +0000955
956
showard2bab8f42008-11-12 18:15:22 +0000957 def _record_attributes(self, attributes):
958 """
959 See on_attribute_changed.
960 """
961 assert not isinstance(attributes, basestring)
962 self._recorded_attributes = dict((attribute, getattr(self, attribute))
963 for attribute in attributes)
964
965
966 def _check_for_updated_attributes(self):
967 """
968 See on_attribute_changed.
969 """
970 for attribute, original_value in self._recorded_attributes.iteritems():
971 new_value = getattr(self, attribute)
972 if original_value != new_value:
973 self.on_attribute_changed(attribute, original_value)
974 self._record_attributes(self._recorded_attributes.keys())
975
976
977 def on_attribute_changed(self, attribute, old_value):
978 """
979 Called whenever an attribute is updated. To be overridden.
980
981 To use this method, you must:
982 * call _record_attributes() from __init__() (after making the super
983 call) with a list of attributes for which you want to be notified upon
984 change.
985 * call _check_for_updated_attributes() from save().
986 """
987 pass
988
989
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700990 def serialize(self, include_dependencies=True):
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700991 """Serializes the object with dependencies.
992
993 The variable SERIALIZATION_LINKS_TO_FOLLOW defines which dependencies
994 this function will serialize with the object.
995
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700996 @param include_dependencies: Whether or not to follow relations to
997 objects this object depends on.
998 This parameter is used when uploading
999 jobs from a shard to the master, as the
1000 master already has all the dependent
1001 objects.
1002
Jakob Juelich3bb7c802014-09-02 16:31:11 -07001003 @returns: Dictionary representation of the object.
1004 """
1005 serialized = {}
1006 for field in self._meta.concrete_model._meta.local_fields:
1007 if field.rel is None:
1008 serialized[field.name] = field._get_val_from_obj(self)
1009
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001010 if include_dependencies:
1011 for link in self.SERIALIZATION_LINKS_TO_FOLLOW:
1012 serialized[link] = self._serialize_relation(link)
Jakob Juelich3bb7c802014-09-02 16:31:11 -07001013
1014 return serialized
1015
1016
1017 def _serialize_relation(self, link):
1018 """Serializes dependent objects given the name of the relation.
1019
1020 @param link: Name of the relation to take objects from.
1021
1022 @returns For To-Many relationships a list of the serialized related
1023 objects, for To-One relationships the serialized related object.
1024 """
1025 try:
1026 attr = getattr(self, link)
1027 except AttributeError:
1028 # One-To-One relationships that point to None may raise this
1029 return None
1030
1031 if attr is None:
1032 return None
1033 if hasattr(attr, 'all'):
1034 return [obj.serialize() for obj in attr.all()]
1035 return attr.serialize()
1036
1037
Jakob Juelichf88fa932014-09-03 17:58:04 -07001038 @classmethod
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001039 def _split_local_from_foreign_values(cls, data):
1040 """This splits local from foreign values in a serialized object.
1041
1042 @param data: The serialized object.
1043
1044 @returns A tuple of two lists, both containing tuples in the form
1045 (link_name, link_value). The first list contains all links
1046 for local fields, the second one contains those for foreign
1047 fields/objects.
1048 """
1049 links_to_local_values, links_to_related_values = [], []
1050 for link, value in data.iteritems():
1051 if link in cls.SERIALIZATION_LINKS_TO_FOLLOW:
1052 # It's a foreign key
1053 links_to_related_values.append((link, value))
1054 else:
1055 # It's a local attribute
1056 links_to_local_values.append((link, value))
1057 return links_to_local_values, links_to_related_values
1058
1059
Jakob Juelichf865d332014-09-29 10:47:49 -07001060 @classmethod
1061 def _filter_update_allowed_fields(cls, data):
1062 """Filters data and returns only files that updates are allowed on.
1063
1064 This is i.e. needed for syncing aborted bits from the master to shards.
1065
1066 Local links are only allowed to be updated, if they are in
1067 SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1068 Overwriting existing values is allowed in order to be able to sync i.e.
1069 the aborted bit from the master to a shard.
1070
1071 The whitelisting mechanism is in place to prevent overwriting local
1072 status: If all fields were overwritten, jobs would be completely be
1073 set back to their original (unstarted) state.
1074
1075 @param data: List with tuples of the form (link_name, link_value), as
1076 returned by _split_local_from_foreign_values.
1077
1078 @returns List of the same format as data, but only containing data for
1079 fields that updates are allowed on.
1080 """
1081 return [pair for pair in data
1082 if pair[0] in cls.SERIALIZATION_LOCAL_LINKS_TO_UPDATE]
1083
1084
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001085 def _deserialize_local(self, data):
1086 """Set local attributes from a list of tuples.
1087
1088 @param data: List of tuples like returned by
1089 _split_local_from_foreign_values.
1090 """
1091 for link, value in data:
1092 setattr(self, link, value)
1093 # Overwridden save() methods are prone to errors, so don't execute them.
1094 # This is because:
1095 # - the overwritten methods depend on ACL groups that don't yet exist
1096 # and don't handle errors
1097 # - the overwritten methods think this object already exists in the db
1098 # because the id is already set
1099 super(type(self), self).save()
1100
1101
1102 def _deserialize_relations(self, data):
1103 """Set foreign attributes from a list of tuples.
1104
1105 This deserialized the related objects using their own deserialize()
1106 function and then sets the relation.
1107
1108 @param data: List of tuples like returned by
1109 _split_local_from_foreign_values.
1110 """
1111 for link, value in data:
1112 self._deserialize_relation(link, value)
1113 # See comment in _deserialize_local
1114 super(type(self), self).save()
1115
1116
1117 @classmethod
Jakob Juelichf88fa932014-09-03 17:58:04 -07001118 def deserialize(cls, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001119 """Recursively deserializes and saves an object with it's dependencies.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001120
1121 This takes the result of the serialize method and creates objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001122 in the database that are just like the original.
1123
1124 If an object of the same type with the same id already exists, it's
Jakob Juelichf865d332014-09-29 10:47:49 -07001125 local values will be left untouched, unless they are explicitly
1126 whitelisted in SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1127
1128 Deserialize will always recursively propagate to all related objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001129 present in data though.
1130 I.e. this is necessary to add users to an already existing acl-group.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001131
1132 @param data: Representation of an object and its dependencies, as
1133 returned by serialize.
1134
1135 @returns: The object represented by data if it didn't exist before,
1136 otherwise the object that existed before and has the same type
1137 and id as the one described by data.
1138 """
1139 if data is None:
1140 return None
1141
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001142 local, related = cls._split_local_from_foreign_values(data)
1143
Jakob Juelichf88fa932014-09-03 17:58:04 -07001144 try:
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001145 instance = cls.objects.get(id=data['id'])
Jakob Juelichf865d332014-09-29 10:47:49 -07001146 local = cls._filter_update_allowed_fields(local)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001147 except cls.DoesNotExist:
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001148 instance = cls()
Jakob Juelichf88fa932014-09-03 17:58:04 -07001149
Jakob Juelichf865d332014-09-29 10:47:49 -07001150 instance._deserialize_local(local)
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001151 instance._deserialize_relations(related)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001152
1153 return instance
1154
1155
Jakob Juelicha94efe62014-09-18 16:02:49 -07001156 def sanity_check_update_from_shard(self, shard, updated_serialized,
1157 *args, **kwargs):
1158 """Check if an update sent from a shard is legitimate.
1159
1160 @raises error.UnallowedRecordsSentToMaster if an update is not
1161 legitimate.
1162 """
1163 raise NotImplementedError(
1164 'sanity_check_update_from_shard must be implemented by subclass %s '
1165 'for type %s' % type(self))
1166
1167
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001168 def update_from_serialized(self, serialized):
1169 """Updates local fields of an existing object from a serialized form.
1170
1171 This is different than the normal deserialize() in the way that it
1172 does update local values, which deserialize doesn't, but doesn't
1173 recursively propagate to related objects, which deserialize() does.
1174
1175 The use case of this function is to update job records on the master
1176 after the jobs have been executed on a slave, as the master is not
1177 interested in updates for users, labels, specialtasks, etc.
1178
1179 @param serialized: Representation of an object and its dependencies, as
1180 returned by serialize.
1181
1182 @raises ValueError: if serialized contains related objects, i.e. not
1183 only local fields.
1184 """
1185 local, related = (
1186 self._split_local_from_foreign_values(serialized))
1187 if related:
1188 raise ValueError('Serialized must not contain foreign '
1189 'objects: %s' % related)
1190
1191 self._deserialize_local(local)
1192
1193
Jakob Juelichf88fa932014-09-03 17:58:04 -07001194 def custom_deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001195 """Allows overriding the deserialization behaviour by subclasses."""
Jakob Juelichf88fa932014-09-03 17:58:04 -07001196 raise NotImplementedError(
1197 'custom_deserialize_relation must be implemented by subclass %s '
1198 'for relation %s' % (type(self), link))
1199
1200
1201 def _deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001202 """Deserializes related objects and sets references on this object.
1203
1204 Relations that point to a list of objects are handled automatically.
1205 For many-to-one or one-to-one relations custom_deserialize_relation
1206 must be overridden by the subclass.
1207
1208 Related objects are deserialized using their deserialize() method.
1209 Thereby they and their dependencies are created if they don't exist
1210 and saved to the database.
1211
1212 @param link: Name of the relation.
1213 @param data: Serialized representation of the related object(s).
1214 This means a list of dictionaries for to-many relations,
1215 just a dictionary for to-one relations.
1216 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001217 field = getattr(self, link)
1218
1219 if field and hasattr(field, 'all'):
1220 self._deserialize_2m_relation(link, data, field.model)
1221 else:
1222 self.custom_deserialize_relation(link, data)
1223
1224
1225 def _deserialize_2m_relation(self, link, data, related_class):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001226 """Deserialize related objects for one to-many relationship.
1227
1228 @param link: Name of the relation.
1229 @param data: Serialized representation of the related objects.
1230 This is a list with of dictionaries.
1231 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001232 relation_set = getattr(self, link)
1233 for serialized in data:
1234 relation_set.add(related_class.deserialize(serialized))
1235
1236
showard7c785282008-05-29 19:45:12 +00001237class ModelWithInvalid(ModelExtensions):
jadmanski0afbb632008-06-06 21:10:57 +00001238 """
1239 Overrides model methods save() and delete() to support invalidation in
1240 place of actual deletion. Subclasses must have a boolean "invalid"
1241 field.
1242 """
showard7c785282008-05-29 19:45:12 +00001243
showarda5288b42009-07-28 20:06:08 +00001244 def save(self, *args, **kwargs):
showardddb90992009-02-11 23:39:32 +00001245 first_time = (self.id is None)
1246 if first_time:
1247 # see if this object was previously added and invalidated
1248 my_name = getattr(self, self.name_field)
1249 filters = {self.name_field : my_name, 'invalid' : True}
1250 try:
1251 old_object = self.__class__.objects.get(**filters)
showardafd97de2009-10-01 18:45:09 +00001252 self.resurrect_object(old_object)
showardddb90992009-02-11 23:39:32 +00001253 except self.DoesNotExist:
1254 # no existing object
1255 pass
showard7c785282008-05-29 19:45:12 +00001256
showarda5288b42009-07-28 20:06:08 +00001257 super(ModelWithInvalid, self).save(*args, **kwargs)
showard7c785282008-05-29 19:45:12 +00001258
1259
showardafd97de2009-10-01 18:45:09 +00001260 def resurrect_object(self, old_object):
1261 """
1262 Called when self is about to be saved for the first time and is actually
1263 "undeleting" a previously deleted object. Can be overridden by
1264 subclasses to copy data as desired from the deleted entry (but this
1265 superclass implementation must normally be called).
1266 """
1267 self.id = old_object.id
1268
1269
jadmanski0afbb632008-06-06 21:10:57 +00001270 def clean_object(self):
1271 """
1272 This method is called when an object is marked invalid.
1273 Subclasses should override this to clean up relationships that
showardafd97de2009-10-01 18:45:09 +00001274 should no longer exist if the object were deleted.
1275 """
jadmanski0afbb632008-06-06 21:10:57 +00001276 pass
showard7c785282008-05-29 19:45:12 +00001277
1278
jadmanski0afbb632008-06-06 21:10:57 +00001279 def delete(self):
Dale Curtis74a314b2011-06-23 14:55:46 -07001280 self.invalid = self.invalid
jadmanski0afbb632008-06-06 21:10:57 +00001281 assert not self.invalid
1282 self.invalid = True
1283 self.save()
1284 self.clean_object()
showard7c785282008-05-29 19:45:12 +00001285
1286
jadmanski0afbb632008-06-06 21:10:57 +00001287 @classmethod
1288 def get_valid_manager(cls):
1289 return cls.valid_objects
showard7c785282008-05-29 19:45:12 +00001290
1291
jadmanski0afbb632008-06-06 21:10:57 +00001292 class Manipulator(object):
1293 """
1294 Force default manipulators to look only at valid objects -
1295 otherwise they will match against invalid objects when checking
1296 uniqueness.
1297 """
1298 @classmethod
1299 def _prepare(cls, model):
1300 super(ModelWithInvalid.Manipulator, cls)._prepare(model)
1301 cls.manager = model.valid_objects
showardf8b19042009-05-12 17:22:49 +00001302
1303
1304class ModelWithAttributes(object):
1305 """
1306 Mixin class for models that have an attribute model associated with them.
1307 The attribute model is assumed to have its value field named "value".
1308 """
1309
1310 def _get_attribute_model_and_args(self, attribute):
1311 """
1312 Subclasses should override this to return a tuple (attribute_model,
1313 keyword_args), where attribute_model is a model class and keyword_args
1314 is a dict of args to pass to attribute_model.objects.get() to get an
1315 instance of the given attribute on this object.
1316 """
Dale Curtis74a314b2011-06-23 14:55:46 -07001317 raise NotImplementedError
showardf8b19042009-05-12 17:22:49 +00001318
1319
1320 def set_attribute(self, attribute, value):
1321 attribute_model, get_args = self._get_attribute_model_and_args(
1322 attribute)
1323 attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
1324 attribute_object.value = value
1325 attribute_object.save()
1326
1327
1328 def delete_attribute(self, attribute):
1329 attribute_model, get_args = self._get_attribute_model_and_args(
1330 attribute)
1331 try:
1332 attribute_model.objects.get(**get_args).delete()
showard16245422009-09-08 16:28:15 +00001333 except attribute_model.DoesNotExist:
showardf8b19042009-05-12 17:22:49 +00001334 pass
1335
1336
1337 def set_or_delete_attribute(self, attribute, value):
1338 if value is None:
1339 self.delete_attribute(attribute)
1340 else:
1341 self.set_attribute(attribute, value)
showard26b7ec72009-12-21 22:43:57 +00001342
1343
1344class ModelWithHashManager(dbmodels.Manager):
1345 """Manager for use with the ModelWithHash abstract model class"""
1346
1347 def create(self, **kwargs):
1348 raise Exception('ModelWithHash manager should use get_or_create() '
1349 'instead of create()')
1350
1351
1352 def get_or_create(self, **kwargs):
1353 kwargs['the_hash'] = self.model._compute_hash(**kwargs)
1354 return super(ModelWithHashManager, self).get_or_create(**kwargs)
1355
1356
1357class ModelWithHash(dbmodels.Model):
1358 """Superclass with methods for dealing with a hash column"""
1359
1360 the_hash = dbmodels.CharField(max_length=40, unique=True)
1361
1362 objects = ModelWithHashManager()
1363
1364 class Meta:
1365 abstract = True
1366
1367
1368 @classmethod
1369 def _compute_hash(cls, **kwargs):
1370 raise NotImplementedError('Subclasses must override _compute_hash()')
1371
1372
1373 def save(self, force_insert=False, **kwargs):
1374 """Prevents saving the model in most cases
1375
1376 We want these models to be immutable, so the generic save() operation
1377 will not work. These models should be instantiated through their the
1378 model.objects.get_or_create() method instead.
1379
1380 The exception is that save(force_insert=True) will be allowed, since
1381 that creates a new row. However, the preferred way to make instances of
1382 these models is through the get_or_create() method.
1383 """
1384 if not force_insert:
1385 # Allow a forced insert to happen; if it's a duplicate, the unique
1386 # constraint will catch it later anyways
1387 raise Exception('ModelWithHash is immutable')
1388 super(ModelWithHash, self).save(force_insert=force_insert, **kwargs)