blob: 6329bd0e383382a8cf7c1370002fda03384edb16 [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
jadmanski0afbb632008-06-06 21:10:57 +0000583 @classmethod
584 def convert_human_readable_values(cls, data, to_human_readable=False):
585 """\
586 Performs conversions on user-supplied field data, to make it
587 easier for users to pass human-readable data.
showard7c785282008-05-29 19:45:12 +0000588
jadmanski0afbb632008-06-06 21:10:57 +0000589 For all fields that have choice sets, convert their values
590 from human-readable strings to enum values, if necessary. This
591 allows users to pass strings instead of the corresponding
592 integer values.
showard7c785282008-05-29 19:45:12 +0000593
jadmanski0afbb632008-06-06 21:10:57 +0000594 For all foreign key fields, call smart_get with the supplied
595 data. This allows the user to pass either an ID value or
596 the name of the object as a string.
showard7c785282008-05-29 19:45:12 +0000597
jadmanski0afbb632008-06-06 21:10:57 +0000598 If to_human_readable=True, perform the inverse - i.e. convert
599 numeric values to human readable values.
showard7c785282008-05-29 19:45:12 +0000600
jadmanski0afbb632008-06-06 21:10:57 +0000601 This method modifies data in-place.
602 """
603 field_dict = cls.get_field_dict()
604 for field_name in data:
showarde732ee72008-09-23 19:15:43 +0000605 if field_name not in field_dict or data[field_name] is None:
jadmanski0afbb632008-06-06 21:10:57 +0000606 continue
607 field_obj = field_dict[field_name]
608 # convert enum values
609 if field_obj.choices:
610 for choice_data in field_obj.choices:
611 # choice_data is (value, name)
612 if to_human_readable:
613 from_val, to_val = choice_data
614 else:
615 to_val, from_val = choice_data
616 if from_val == data[field_name]:
617 data[field_name] = to_val
618 break
619 # convert foreign key values
620 elif field_obj.rel:
showarda4ea5742009-02-17 20:56:23 +0000621 dest_obj = field_obj.rel.to.smart_get(data[field_name],
622 valid_only=False)
showardf8b19042009-05-12 17:22:49 +0000623 if to_human_readable:
Paul Pendlebury5a8c6ad2011-02-01 07:20:17 -0800624 # parameterized_jobs do not have a name_field
625 if (field_name != 'parameterized_job' and
626 dest_obj.name_field is not None):
showardf8b19042009-05-12 17:22:49 +0000627 data[field_name] = getattr(dest_obj,
628 dest_obj.name_field)
jadmanski0afbb632008-06-06 21:10:57 +0000629 else:
showardb0a73032009-03-27 18:35:41 +0000630 data[field_name] = dest_obj
showard7c785282008-05-29 19:45:12 +0000631
632
showard7c785282008-05-29 19:45:12 +0000633
634
Dale Curtis74a314b2011-06-23 14:55:46 -0700635 def _validate_unique(self):
jadmanski0afbb632008-06-06 21:10:57 +0000636 """\
637 Validate that unique fields are unique. Django manipulators do
638 this too, but they're a huge pain to use manually. Trust me.
639 """
640 errors = {}
641 cls = type(self)
642 field_dict = self.get_field_dict()
643 manager = cls.get_valid_manager()
644 for field_name, field_obj in field_dict.iteritems():
645 if not field_obj.unique:
646 continue
showard7c785282008-05-29 19:45:12 +0000647
jadmanski0afbb632008-06-06 21:10:57 +0000648 value = getattr(self, field_name)
showardbd18ab72009-09-18 21:20:27 +0000649 if value is None and field_obj.auto_created:
650 # don't bother checking autoincrement fields about to be
651 # generated
652 continue
653
jadmanski0afbb632008-06-06 21:10:57 +0000654 existing_objs = manager.filter(**{field_name : value})
655 num_existing = existing_objs.count()
showard7c785282008-05-29 19:45:12 +0000656
jadmanski0afbb632008-06-06 21:10:57 +0000657 if num_existing == 0:
658 continue
659 if num_existing == 1 and existing_objs[0].id == self.id:
660 continue
661 errors[field_name] = (
662 'This value must be unique (%s)' % (value))
663 return errors
showard7c785282008-05-29 19:45:12 +0000664
665
showarda5288b42009-07-28 20:06:08 +0000666 def _validate(self):
667 """
668 First coerces all fields on this instance to their proper Python types.
669 Then runs validation on every field. Returns a dictionary of
670 field_name -> error_list.
671
672 Based on validate() from django.db.models.Model in Django 0.96, which
673 was removed in Django 1.0. It should reappear in a later version. See:
674 http://code.djangoproject.com/ticket/6845
675 """
676 error_dict = {}
677 for f in self._meta.fields:
678 try:
679 python_value = f.to_python(
680 getattr(self, f.attname, f.get_default()))
681 except django.core.exceptions.ValidationError, e:
jamesren1e0a4ce2010-04-21 17:45:11 +0000682 error_dict[f.name] = str(e)
showarda5288b42009-07-28 20:06:08 +0000683 continue
684
685 if not f.blank and not python_value:
686 error_dict[f.name] = 'This field is required.'
687 continue
688
689 setattr(self, f.attname, python_value)
690
691 return error_dict
692
693
jadmanski0afbb632008-06-06 21:10:57 +0000694 def do_validate(self):
showarda5288b42009-07-28 20:06:08 +0000695 errors = self._validate()
Dale Curtis74a314b2011-06-23 14:55:46 -0700696 unique_errors = self._validate_unique()
jadmanski0afbb632008-06-06 21:10:57 +0000697 for field_name, error in unique_errors.iteritems():
698 errors.setdefault(field_name, error)
699 if errors:
700 raise ValidationError(errors)
showard7c785282008-05-29 19:45:12 +0000701
702
jadmanski0afbb632008-06-06 21:10:57 +0000703 # actually (externally) useful methods follow
showard7c785282008-05-29 19:45:12 +0000704
jadmanski0afbb632008-06-06 21:10:57 +0000705 @classmethod
706 def add_object(cls, data={}, **kwargs):
707 """\
708 Returns a new object created with the given data (a dictionary
709 mapping field names to values). Merges any extra keyword args
710 into data.
711 """
Prashanth B489b91d2014-03-15 12:17:16 -0700712 data = dict(data)
713 data.update(kwargs)
714 data = cls.prepare_data_args(data)
715 cls.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000716 data = cls.provide_default_values(data)
Prashanth B489b91d2014-03-15 12:17:16 -0700717
jadmanski0afbb632008-06-06 21:10:57 +0000718 obj = cls(**data)
719 obj.do_validate()
720 obj.save()
721 return obj
showard7c785282008-05-29 19:45:12 +0000722
Michael Liang8864e862014-07-22 08:36:05 -0700723 def record_state(self, type_str, state, value):
724 """Record metadata in elasticsearch.
725
726 @param type_str: sets the _type field in elasticsearch db.
727 @param state: string representing what state we are recording,
728 e.g. 'locked'
729 @param value: value of the state, e.g. True
730 """
731 metadata = {
Dan Shie4cb9e22014-08-29 11:56:04 -0700732 state: value,
Michael Liang8864e862014-07-22 08:36:05 -0700733 'hostname': self.hostname,
734 }
735 es_utils.ESMetadata().post(type_str=type_str, metadata=metadata)
736
showard7c785282008-05-29 19:45:12 +0000737
jadmanski0afbb632008-06-06 21:10:57 +0000738 def update_object(self, data={}, **kwargs):
739 """\
740 Updates the object with the given data (a dictionary mapping
741 field names to values). Merges any extra keyword args into
742 data.
743 """
Prashanth B489b91d2014-03-15 12:17:16 -0700744 data = dict(data)
745 data.update(kwargs)
746 data = self.prepare_data_args(data)
747 self.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000748 for field_name, value in data.iteritems():
showardb0a73032009-03-27 18:35:41 +0000749 setattr(self, field_name, value)
Michael Liang8864e862014-07-22 08:36:05 -0700750 # Other fields such as label (when updated) are not sent over
751 # the es because it doesn't contribute to putting together host
752 # host history. Locks are important in host history because if
753 # a device is locked then we don't really care what state it is in.
754 if field_name == 'locked':
755 self.record_state('lock_history', 'locked', value)
jadmanski0afbb632008-06-06 21:10:57 +0000756 self.do_validate()
757 self.save()
showard7c785282008-05-29 19:45:12 +0000758
759
showard8bfb5cb2009-10-07 20:49:15 +0000760 # see query_objects()
761 _SPECIAL_FILTER_KEYS = ('query_start', 'query_limit', 'sort_by',
762 'extra_args', 'extra_where', 'no_distinct')
763
764
jadmanski0afbb632008-06-06 21:10:57 +0000765 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000766 def _extract_special_params(cls, filter_data):
767 """
768 @returns a tuple of dicts (special_params, regular_filters), where
769 special_params contains the parameters we handle specially and
770 regular_filters is the remaining data to be handled by Django.
771 """
772 regular_filters = dict(filter_data)
773 special_params = {}
774 for key in cls._SPECIAL_FILTER_KEYS:
775 if key in regular_filters:
776 special_params[key] = regular_filters.pop(key)
777 return special_params, regular_filters
778
779
780 @classmethod
781 def apply_presentation(cls, query, filter_data):
782 """
783 Apply presentation parameters -- sorting and paging -- to the given
784 query.
785 @returns new query with presentation applied
786 """
787 special_params, _ = cls._extract_special_params(filter_data)
788 sort_by = special_params.get('sort_by', None)
789 if sort_by:
790 assert isinstance(sort_by, list) or isinstance(sort_by, tuple)
showard8b0ea222009-12-23 19:23:03 +0000791 query = query.extra(order_by=sort_by)
showard8bfb5cb2009-10-07 20:49:15 +0000792
793 query_start = special_params.get('query_start', None)
794 query_limit = special_params.get('query_limit', None)
795 if query_start is not None:
796 if query_limit is None:
797 raise ValueError('Cannot pass query_start without query_limit')
798 # query_limit is passed as a page size
showard7074b742009-10-12 20:30:04 +0000799 query_limit += query_start
800 return query[query_start:query_limit]
showard8bfb5cb2009-10-07 20:49:15 +0000801
802
803 @classmethod
804 def query_objects(cls, filter_data, valid_only=True, initial_query=None,
805 apply_presentation=True):
jadmanski0afbb632008-06-06 21:10:57 +0000806 """\
807 Returns a QuerySet object for querying the given model_class
808 with the given filter_data. Optional special arguments in
809 filter_data include:
810 -query_start: index of first return to return
811 -query_limit: maximum number of results to return
812 -sort_by: list of fields to sort on. prefixing a '-' onto a
813 field name changes the sort to descending order.
814 -extra_args: keyword args to pass to query.extra() (see Django
815 DB layer documentation)
showarda5288b42009-07-28 20:06:08 +0000816 -extra_where: extra WHERE clause to append
showard8bfb5cb2009-10-07 20:49:15 +0000817 -no_distinct: if True, a DISTINCT will not be added to the SELECT
jadmanski0afbb632008-06-06 21:10:57 +0000818 """
showard8bfb5cb2009-10-07 20:49:15 +0000819 special_params, regular_filters = cls._extract_special_params(
820 filter_data)
showard7c785282008-05-29 19:45:12 +0000821
showard7ac7b7a2008-07-21 20:24:29 +0000822 if initial_query is None:
823 if valid_only:
824 initial_query = cls.get_valid_manager()
825 else:
826 initial_query = cls.objects
showard8bfb5cb2009-10-07 20:49:15 +0000827
828 query = initial_query.filter(**regular_filters)
829
830 use_distinct = not special_params.get('no_distinct', False)
showard7ac7b7a2008-07-21 20:24:29 +0000831 if use_distinct:
832 query = query.distinct()
showard7c785282008-05-29 19:45:12 +0000833
showard8bfb5cb2009-10-07 20:49:15 +0000834 extra_args = special_params.get('extra_args', {})
835 extra_where = special_params.get('extra_where', None)
836 if extra_where:
837 # escape %'s
838 extra_where = cls.objects.escape_user_sql(extra_where)
839 extra_args.setdefault('where', []).append(extra_where)
jadmanski0afbb632008-06-06 21:10:57 +0000840 if extra_args:
841 query = query.extra(**extra_args)
showard09096d82008-07-07 23:20:49 +0000842 query = query._clone(klass=ReadonlyQuerySet)
showard7c785282008-05-29 19:45:12 +0000843
showard8bfb5cb2009-10-07 20:49:15 +0000844 if apply_presentation:
845 query = cls.apply_presentation(query, filter_data)
846
847 return query
showard7c785282008-05-29 19:45:12 +0000848
849
jadmanski0afbb632008-06-06 21:10:57 +0000850 @classmethod
showard585c2ab2008-07-23 19:29:49 +0000851 def query_count(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000852 """\
853 Like query_objects, but retreive only the count of results.
854 """
855 filter_data.pop('query_start', None)
856 filter_data.pop('query_limit', None)
showard585c2ab2008-07-23 19:29:49 +0000857 query = cls.query_objects(filter_data, initial_query=initial_query)
858 return query.count()
showard7c785282008-05-29 19:45:12 +0000859
860
jadmanski0afbb632008-06-06 21:10:57 +0000861 @classmethod
862 def clean_object_dicts(cls, field_dicts):
863 """\
864 Take a list of dicts corresponding to object (as returned by
865 query.values()) and clean the data to be more suitable for
866 returning to the user.
867 """
showarde732ee72008-09-23 19:15:43 +0000868 for field_dict in field_dicts:
869 cls.clean_foreign_keys(field_dict)
showard21baa452008-10-21 00:08:39 +0000870 cls._convert_booleans(field_dict)
showarde732ee72008-09-23 19:15:43 +0000871 cls.convert_human_readable_values(field_dict,
872 to_human_readable=True)
showard7c785282008-05-29 19:45:12 +0000873
874
jadmanski0afbb632008-06-06 21:10:57 +0000875 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000876 def list_objects(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000877 """\
878 Like query_objects, but return a list of dictionaries.
879 """
showard7ac7b7a2008-07-21 20:24:29 +0000880 query = cls.query_objects(filter_data, initial_query=initial_query)
showard8bfb5cb2009-10-07 20:49:15 +0000881 extra_fields = query.query.extra_select.keys()
882 field_dicts = [model_object.get_object_dict(extra_fields=extra_fields)
showarde732ee72008-09-23 19:15:43 +0000883 for model_object in query]
jadmanski0afbb632008-06-06 21:10:57 +0000884 return field_dicts
showard7c785282008-05-29 19:45:12 +0000885
886
jadmanski0afbb632008-06-06 21:10:57 +0000887 @classmethod
showarda4ea5742009-02-17 20:56:23 +0000888 def smart_get(cls, id_or_name, valid_only=True):
jadmanski0afbb632008-06-06 21:10:57 +0000889 """\
890 smart_get(integer) -> get object by ID
891 smart_get(string) -> get object by name_field
jadmanski0afbb632008-06-06 21:10:57 +0000892 """
showarda4ea5742009-02-17 20:56:23 +0000893 if valid_only:
894 manager = cls.get_valid_manager()
895 else:
896 manager = cls.objects
897
898 if isinstance(id_or_name, (int, long)):
899 return manager.get(pk=id_or_name)
jamesren3e9f6092010-03-11 21:32:10 +0000900 if isinstance(id_or_name, basestring) and hasattr(cls, 'name_field'):
showarda4ea5742009-02-17 20:56:23 +0000901 return manager.get(**{cls.name_field : id_or_name})
902 raise ValueError(
903 'Invalid positional argument: %s (%s)' % (id_or_name,
904 type(id_or_name)))
showard7c785282008-05-29 19:45:12 +0000905
906
showardbe3ec042008-11-12 18:16:07 +0000907 @classmethod
908 def smart_get_bulk(cls, id_or_name_list):
909 invalid_inputs = []
910 result_objects = []
911 for id_or_name in id_or_name_list:
912 try:
913 result_objects.append(cls.smart_get(id_or_name))
914 except cls.DoesNotExist:
915 invalid_inputs.append(id_or_name)
916 if invalid_inputs:
mbligh7a3ebe32008-12-01 17:10:33 +0000917 raise cls.DoesNotExist('The following %ss do not exist: %s'
918 % (cls.__name__.lower(),
919 ', '.join(invalid_inputs)))
showardbe3ec042008-11-12 18:16:07 +0000920 return result_objects
921
922
showard8bfb5cb2009-10-07 20:49:15 +0000923 def get_object_dict(self, extra_fields=None):
jadmanski0afbb632008-06-06 21:10:57 +0000924 """\
showard8bfb5cb2009-10-07 20:49:15 +0000925 Return a dictionary mapping fields to this object's values. @param
926 extra_fields: list of extra attribute names to include, in addition to
927 the fields defined on this object.
jadmanski0afbb632008-06-06 21:10:57 +0000928 """
showard8bfb5cb2009-10-07 20:49:15 +0000929 fields = self.get_field_dict().keys()
930 if extra_fields:
931 fields += extra_fields
jadmanski0afbb632008-06-06 21:10:57 +0000932 object_dict = dict((field_name, getattr(self, field_name))
showarde732ee72008-09-23 19:15:43 +0000933 for field_name in fields)
jadmanski0afbb632008-06-06 21:10:57 +0000934 self.clean_object_dicts([object_dict])
showardd3dc1992009-04-22 21:01:40 +0000935 self._postprocess_object_dict(object_dict)
jadmanski0afbb632008-06-06 21:10:57 +0000936 return object_dict
showard7c785282008-05-29 19:45:12 +0000937
938
showardd3dc1992009-04-22 21:01:40 +0000939 def _postprocess_object_dict(self, object_dict):
940 """For subclasses to override."""
941 pass
942
943
jadmanski0afbb632008-06-06 21:10:57 +0000944 @classmethod
945 def get_valid_manager(cls):
946 return cls.objects
showard7c785282008-05-29 19:45:12 +0000947
948
showard2bab8f42008-11-12 18:15:22 +0000949 def _record_attributes(self, attributes):
950 """
951 See on_attribute_changed.
952 """
953 assert not isinstance(attributes, basestring)
954 self._recorded_attributes = dict((attribute, getattr(self, attribute))
955 for attribute in attributes)
956
957
958 def _check_for_updated_attributes(self):
959 """
960 See on_attribute_changed.
961 """
962 for attribute, original_value in self._recorded_attributes.iteritems():
963 new_value = getattr(self, attribute)
964 if original_value != new_value:
965 self.on_attribute_changed(attribute, original_value)
966 self._record_attributes(self._recorded_attributes.keys())
967
968
969 def on_attribute_changed(self, attribute, old_value):
970 """
971 Called whenever an attribute is updated. To be overridden.
972
973 To use this method, you must:
974 * call _record_attributes() from __init__() (after making the super
975 call) with a list of attributes for which you want to be notified upon
976 change.
977 * call _check_for_updated_attributes() from save().
978 """
979 pass
980
981
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700982 def serialize(self, include_dependencies=True):
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700983 """Serializes the object with dependencies.
984
985 The variable SERIALIZATION_LINKS_TO_FOLLOW defines which dependencies
986 this function will serialize with the object.
987
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700988 @param include_dependencies: Whether or not to follow relations to
989 objects this object depends on.
990 This parameter is used when uploading
991 jobs from a shard to the master, as the
992 master already has all the dependent
993 objects.
994
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700995 @returns: Dictionary representation of the object.
996 """
997 serialized = {}
998 for field in self._meta.concrete_model._meta.local_fields:
999 if field.rel is None:
1000 serialized[field.name] = field._get_val_from_obj(self)
1001
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001002 if include_dependencies:
1003 for link in self.SERIALIZATION_LINKS_TO_FOLLOW:
1004 serialized[link] = self._serialize_relation(link)
Jakob Juelich3bb7c802014-09-02 16:31:11 -07001005
1006 return serialized
1007
1008
1009 def _serialize_relation(self, link):
1010 """Serializes dependent objects given the name of the relation.
1011
1012 @param link: Name of the relation to take objects from.
1013
1014 @returns For To-Many relationships a list of the serialized related
1015 objects, for To-One relationships the serialized related object.
1016 """
1017 try:
1018 attr = getattr(self, link)
1019 except AttributeError:
1020 # One-To-One relationships that point to None may raise this
1021 return None
1022
1023 if attr is None:
1024 return None
1025 if hasattr(attr, 'all'):
1026 return [obj.serialize() for obj in attr.all()]
1027 return attr.serialize()
1028
1029
Jakob Juelichf88fa932014-09-03 17:58:04 -07001030 @classmethod
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001031 def _split_local_from_foreign_values(cls, data):
1032 """This splits local from foreign values in a serialized object.
1033
1034 @param data: The serialized object.
1035
1036 @returns A tuple of two lists, both containing tuples in the form
1037 (link_name, link_value). The first list contains all links
1038 for local fields, the second one contains those for foreign
1039 fields/objects.
1040 """
1041 links_to_local_values, links_to_related_values = [], []
1042 for link, value in data.iteritems():
1043 if link in cls.SERIALIZATION_LINKS_TO_FOLLOW:
1044 # It's a foreign key
1045 links_to_related_values.append((link, value))
1046 else:
1047 # It's a local attribute
1048 links_to_local_values.append((link, value))
1049 return links_to_local_values, links_to_related_values
1050
1051
1052 def _deserialize_local(self, data):
1053 """Set local attributes from a list of tuples.
1054
1055 @param data: List of tuples like returned by
1056 _split_local_from_foreign_values.
1057 """
1058 for link, value in data:
1059 setattr(self, link, value)
1060 # Overwridden save() methods are prone to errors, so don't execute them.
1061 # This is because:
1062 # - the overwritten methods depend on ACL groups that don't yet exist
1063 # and don't handle errors
1064 # - the overwritten methods think this object already exists in the db
1065 # because the id is already set
1066 super(type(self), self).save()
1067
1068
1069 def _deserialize_relations(self, data):
1070 """Set foreign attributes from a list of tuples.
1071
1072 This deserialized the related objects using their own deserialize()
1073 function and then sets the relation.
1074
1075 @param data: List of tuples like returned by
1076 _split_local_from_foreign_values.
1077 """
1078 for link, value in data:
1079 self._deserialize_relation(link, value)
1080 # See comment in _deserialize_local
1081 super(type(self), self).save()
1082
1083
1084 @classmethod
Jakob Juelichf88fa932014-09-03 17:58:04 -07001085 def deserialize(cls, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001086 """Recursively deserializes and saves an object with it's dependencies.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001087
1088 This takes the result of the serialize method and creates objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001089 in the database that are just like the original.
1090
1091 If an object of the same type with the same id already exists, it's
1092 local values will be left untouched.
1093 Deserialize will still recursively propagate to all related objects
1094 present in data though.
1095 I.e. this is necessary to add users to an already existing acl-group.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001096
1097 @param data: Representation of an object and its dependencies, as
1098 returned by serialize.
1099
1100 @returns: The object represented by data if it didn't exist before,
1101 otherwise the object that existed before and has the same type
1102 and id as the one described by data.
1103 """
1104 if data is None:
1105 return None
1106
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001107 local, related = cls._split_local_from_foreign_values(data)
1108
Jakob Juelichf88fa932014-09-03 17:58:04 -07001109 try:
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001110 instance = cls.objects.get(id=data['id'])
Jakob Juelichf88fa932014-09-03 17:58:04 -07001111 except cls.DoesNotExist:
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001112 instance = cls()
1113 instance._deserialize_local(local)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001114
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001115 instance._deserialize_relations(related)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001116
1117 return instance
1118
1119
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001120 def update_from_serialized(self, serialized):
1121 """Updates local fields of an existing object from a serialized form.
1122
1123 This is different than the normal deserialize() in the way that it
1124 does update local values, which deserialize doesn't, but doesn't
1125 recursively propagate to related objects, which deserialize() does.
1126
1127 The use case of this function is to update job records on the master
1128 after the jobs have been executed on a slave, as the master is not
1129 interested in updates for users, labels, specialtasks, etc.
1130
1131 @param serialized: Representation of an object and its dependencies, as
1132 returned by serialize.
1133
1134 @raises ValueError: if serialized contains related objects, i.e. not
1135 only local fields.
1136 """
1137 local, related = (
1138 self._split_local_from_foreign_values(serialized))
1139 if related:
1140 raise ValueError('Serialized must not contain foreign '
1141 'objects: %s' % related)
1142
1143 self._deserialize_local(local)
1144
1145
Jakob Juelichf88fa932014-09-03 17:58:04 -07001146 def custom_deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001147 """Allows overriding the deserialization behaviour by subclasses."""
Jakob Juelichf88fa932014-09-03 17:58:04 -07001148 raise NotImplementedError(
1149 'custom_deserialize_relation must be implemented by subclass %s '
1150 'for relation %s' % (type(self), link))
1151
1152
1153 def _deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001154 """Deserializes related objects and sets references on this object.
1155
1156 Relations that point to a list of objects are handled automatically.
1157 For many-to-one or one-to-one relations custom_deserialize_relation
1158 must be overridden by the subclass.
1159
1160 Related objects are deserialized using their deserialize() method.
1161 Thereby they and their dependencies are created if they don't exist
1162 and saved to the database.
1163
1164 @param link: Name of the relation.
1165 @param data: Serialized representation of the related object(s).
1166 This means a list of dictionaries for to-many relations,
1167 just a dictionary for to-one relations.
1168 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001169 field = getattr(self, link)
1170
1171 if field and hasattr(field, 'all'):
1172 self._deserialize_2m_relation(link, data, field.model)
1173 else:
1174 self.custom_deserialize_relation(link, data)
1175
1176
1177 def _deserialize_2m_relation(self, link, data, related_class):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001178 """Deserialize related objects for one to-many relationship.
1179
1180 @param link: Name of the relation.
1181 @param data: Serialized representation of the related objects.
1182 This is a list with of dictionaries.
1183 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001184 relation_set = getattr(self, link)
1185 for serialized in data:
1186 relation_set.add(related_class.deserialize(serialized))
1187
1188
showard7c785282008-05-29 19:45:12 +00001189class ModelWithInvalid(ModelExtensions):
jadmanski0afbb632008-06-06 21:10:57 +00001190 """
1191 Overrides model methods save() and delete() to support invalidation in
1192 place of actual deletion. Subclasses must have a boolean "invalid"
1193 field.
1194 """
showard7c785282008-05-29 19:45:12 +00001195
showarda5288b42009-07-28 20:06:08 +00001196 def save(self, *args, **kwargs):
showardddb90992009-02-11 23:39:32 +00001197 first_time = (self.id is None)
1198 if first_time:
1199 # see if this object was previously added and invalidated
1200 my_name = getattr(self, self.name_field)
1201 filters = {self.name_field : my_name, 'invalid' : True}
1202 try:
1203 old_object = self.__class__.objects.get(**filters)
showardafd97de2009-10-01 18:45:09 +00001204 self.resurrect_object(old_object)
showardddb90992009-02-11 23:39:32 +00001205 except self.DoesNotExist:
1206 # no existing object
1207 pass
showard7c785282008-05-29 19:45:12 +00001208
showarda5288b42009-07-28 20:06:08 +00001209 super(ModelWithInvalid, self).save(*args, **kwargs)
showard7c785282008-05-29 19:45:12 +00001210
1211
showardafd97de2009-10-01 18:45:09 +00001212 def resurrect_object(self, old_object):
1213 """
1214 Called when self is about to be saved for the first time and is actually
1215 "undeleting" a previously deleted object. Can be overridden by
1216 subclasses to copy data as desired from the deleted entry (but this
1217 superclass implementation must normally be called).
1218 """
1219 self.id = old_object.id
1220
1221
jadmanski0afbb632008-06-06 21:10:57 +00001222 def clean_object(self):
1223 """
1224 This method is called when an object is marked invalid.
1225 Subclasses should override this to clean up relationships that
showardafd97de2009-10-01 18:45:09 +00001226 should no longer exist if the object were deleted.
1227 """
jadmanski0afbb632008-06-06 21:10:57 +00001228 pass
showard7c785282008-05-29 19:45:12 +00001229
1230
jadmanski0afbb632008-06-06 21:10:57 +00001231 def delete(self):
Dale Curtis74a314b2011-06-23 14:55:46 -07001232 self.invalid = self.invalid
jadmanski0afbb632008-06-06 21:10:57 +00001233 assert not self.invalid
1234 self.invalid = True
1235 self.save()
1236 self.clean_object()
showard7c785282008-05-29 19:45:12 +00001237
1238
jadmanski0afbb632008-06-06 21:10:57 +00001239 @classmethod
1240 def get_valid_manager(cls):
1241 return cls.valid_objects
showard7c785282008-05-29 19:45:12 +00001242
1243
jadmanski0afbb632008-06-06 21:10:57 +00001244 class Manipulator(object):
1245 """
1246 Force default manipulators to look only at valid objects -
1247 otherwise they will match against invalid objects when checking
1248 uniqueness.
1249 """
1250 @classmethod
1251 def _prepare(cls, model):
1252 super(ModelWithInvalid.Manipulator, cls)._prepare(model)
1253 cls.manager = model.valid_objects
showardf8b19042009-05-12 17:22:49 +00001254
1255
1256class ModelWithAttributes(object):
1257 """
1258 Mixin class for models that have an attribute model associated with them.
1259 The attribute model is assumed to have its value field named "value".
1260 """
1261
1262 def _get_attribute_model_and_args(self, attribute):
1263 """
1264 Subclasses should override this to return a tuple (attribute_model,
1265 keyword_args), where attribute_model is a model class and keyword_args
1266 is a dict of args to pass to attribute_model.objects.get() to get an
1267 instance of the given attribute on this object.
1268 """
Dale Curtis74a314b2011-06-23 14:55:46 -07001269 raise NotImplementedError
showardf8b19042009-05-12 17:22:49 +00001270
1271
1272 def set_attribute(self, attribute, value):
1273 attribute_model, get_args = self._get_attribute_model_and_args(
1274 attribute)
1275 attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
1276 attribute_object.value = value
1277 attribute_object.save()
1278
1279
1280 def delete_attribute(self, attribute):
1281 attribute_model, get_args = self._get_attribute_model_and_args(
1282 attribute)
1283 try:
1284 attribute_model.objects.get(**get_args).delete()
showard16245422009-09-08 16:28:15 +00001285 except attribute_model.DoesNotExist:
showardf8b19042009-05-12 17:22:49 +00001286 pass
1287
1288
1289 def set_or_delete_attribute(self, attribute, value):
1290 if value is None:
1291 self.delete_attribute(attribute)
1292 else:
1293 self.set_attribute(attribute, value)
showard26b7ec72009-12-21 22:43:57 +00001294
1295
1296class ModelWithHashManager(dbmodels.Manager):
1297 """Manager for use with the ModelWithHash abstract model class"""
1298
1299 def create(self, **kwargs):
1300 raise Exception('ModelWithHash manager should use get_or_create() '
1301 'instead of create()')
1302
1303
1304 def get_or_create(self, **kwargs):
1305 kwargs['the_hash'] = self.model._compute_hash(**kwargs)
1306 return super(ModelWithHashManager, self).get_or_create(**kwargs)
1307
1308
1309class ModelWithHash(dbmodels.Model):
1310 """Superclass with methods for dealing with a hash column"""
1311
1312 the_hash = dbmodels.CharField(max_length=40, unique=True)
1313
1314 objects = ModelWithHashManager()
1315
1316 class Meta:
1317 abstract = True
1318
1319
1320 @classmethod
1321 def _compute_hash(cls, **kwargs):
1322 raise NotImplementedError('Subclasses must override _compute_hash()')
1323
1324
1325 def save(self, force_insert=False, **kwargs):
1326 """Prevents saving the model in most cases
1327
1328 We want these models to be immutable, so the generic save() operation
1329 will not work. These models should be instantiated through their the
1330 model.objects.get_or_create() method instead.
1331
1332 The exception is that save(force_insert=True) will be allowed, since
1333 that creates a new row. However, the preferred way to make instances of
1334 these models is through the get_or_create() method.
1335 """
1336 if not force_insert:
1337 # Allow a forced insert to happen; if it's a duplicate, the unique
1338 # constraint will catch it later anyways
1339 raise Exception('ModelWithHash is immutable')
1340 super(ModelWithHash, self).save(force_insert=force_insert, **kwargs)