blob: 5781646c0d557ade5bf880b8326088ac09d89c01 [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
Prashanth Balasubramanian75be1d32014-11-25 18:03:09 -08008from django.db import backend
9from django.db import connection
10from django.db import connections
11from django.db import models as dbmodels
12from django.db import transaction
showarda5288b42009-07-28 20:06:08 +000013from django.db.models.sql import query
showard7e67b432010-01-20 01:13:04 +000014import django.db.models.sql.where
showard7c785282008-05-29 19:45:12 +000015from django.utils import datastructures
Prashanth B489b91d2014-03-15 12:17:16 -070016from autotest_lib.frontend.afe import rdb_model_extensions
showard7c785282008-05-29 19:45:12 +000017
Prashanth B489b91d2014-03-15 12:17:16 -070018
19class ValidationError(django.core.exceptions.ValidationError):
jadmanski0afbb632008-06-06 21:10:57 +000020 """\
showarda5288b42009-07-28 20:06:08 +000021 Data validation error in adding or updating an object. The associated
jadmanski0afbb632008-06-06 21:10:57 +000022 value is a dictionary mapping field names to error strings.
23 """
showard7c785282008-05-29 19:45:12 +000024
showarda5288b42009-07-28 20:06:08 +000025def _quote_name(name):
26 """Shorthand for connection.ops.quote_name()."""
27 return connection.ops.quote_name(name)
28
29
beepscc9fc702013-12-02 12:45:38 -080030class LeasedHostManager(dbmodels.Manager):
31 """Query manager for unleased, unlocked hosts.
32 """
33 def get_query_set(self):
34 return (super(LeasedHostManager, self).get_query_set().filter(
35 leased=0, locked=0))
36
37
showard7c785282008-05-29 19:45:12 +000038class ExtendedManager(dbmodels.Manager):
jadmanski0afbb632008-06-06 21:10:57 +000039 """\
40 Extended manager supporting subquery filtering.
41 """
showard7c785282008-05-29 19:45:12 +000042
showardf828c772010-01-25 21:49:42 +000043 class CustomQuery(query.Query):
showard7e67b432010-01-20 01:13:04 +000044 def __init__(self, *args, **kwargs):
showardf828c772010-01-25 21:49:42 +000045 super(ExtendedManager.CustomQuery, self).__init__(*args, **kwargs)
showard7e67b432010-01-20 01:13:04 +000046 self._custom_joins = []
47
48
showarda5288b42009-07-28 20:06:08 +000049 def clone(self, klass=None, **kwargs):
showardf828c772010-01-25 21:49:42 +000050 obj = super(ExtendedManager.CustomQuery, self).clone(klass)
showard7e67b432010-01-20 01:13:04 +000051 obj._custom_joins = list(self._custom_joins)
showarda5288b42009-07-28 20:06:08 +000052 return obj
showard08f981b2008-06-24 21:59:03 +000053
showard7e67b432010-01-20 01:13:04 +000054
55 def combine(self, rhs, connector):
showardf828c772010-01-25 21:49:42 +000056 super(ExtendedManager.CustomQuery, self).combine(rhs, connector)
showard7e67b432010-01-20 01:13:04 +000057 if hasattr(rhs, '_custom_joins'):
58 self._custom_joins.extend(rhs._custom_joins)
59
60
61 def add_custom_join(self, table, condition, join_type,
62 condition_values=(), alias=None):
63 if alias is None:
64 alias = table
65 join_dict = dict(table=table,
66 condition=condition,
67 condition_values=condition_values,
68 join_type=join_type,
69 alias=alias)
70 self._custom_joins.append(join_dict)
71
72
showard7e67b432010-01-20 01:13:04 +000073 @classmethod
74 def convert_query(self, query_set):
75 """
showardf828c772010-01-25 21:49:42 +000076 Convert the query set's "query" attribute to a CustomQuery.
showard7e67b432010-01-20 01:13:04 +000077 """
78 # Make a copy of the query set
79 query_set = query_set.all()
80 query_set.query = query_set.query.clone(
showardf828c772010-01-25 21:49:42 +000081 klass=ExtendedManager.CustomQuery,
showard7e67b432010-01-20 01:13:04 +000082 _custom_joins=[])
83 return query_set
showard43a3d262008-11-12 18:17:05 +000084
85
showard7e67b432010-01-20 01:13:04 +000086 class _WhereClause(object):
87 """Object allowing us to inject arbitrary SQL into Django queries.
showard43a3d262008-11-12 18:17:05 +000088
showard7e67b432010-01-20 01:13:04 +000089 By using this instead of extra(where=...), we can still freely combine
90 queries with & and |.
showarda5288b42009-07-28 20:06:08 +000091 """
showard7e67b432010-01-20 01:13:04 +000092 def __init__(self, clause, values=()):
93 self._clause = clause
94 self._values = values
showarda5288b42009-07-28 20:06:08 +000095
showard7e67b432010-01-20 01:13:04 +000096
Dale Curtis74a314b2011-06-23 14:55:46 -070097 def as_sql(self, qn=None, connection=None):
showard7e67b432010-01-20 01:13:04 +000098 return self._clause, self._values
99
100
101 def relabel_aliases(self, change_map):
102 return
showard43a3d262008-11-12 18:17:05 +0000103
104
showard8b0ea222009-12-23 19:23:03 +0000105 def add_join(self, query_set, join_table, join_key, join_condition='',
showard7e67b432010-01-20 01:13:04 +0000106 join_condition_values=(), join_from_key=None, alias=None,
107 suffix='', exclude=False, force_left_join=False):
108 """Add a join to query_set.
109
110 Join looks like this:
111 (INNER|LEFT) JOIN <join_table> AS <alias>
112 ON (<this table>.<join_from_key> = <join_table>.<join_key>
113 and <join_condition>)
114
showard0957a842009-05-11 19:25:08 +0000115 @param join_table table to join to
116 @param join_key field referencing back to this model to use for the join
117 @param join_condition extra condition for the ON clause of the join
showard7e67b432010-01-20 01:13:04 +0000118 @param join_condition_values values to substitute into join_condition
119 @param join_from_key column on this model to join from.
showard8b0ea222009-12-23 19:23:03 +0000120 @param alias alias to use for for join
121 @param suffix suffix to add to join_table for the join alias, if no
122 alias is provided
showard0957a842009-05-11 19:25:08 +0000123 @param exclude if true, exclude rows that match this join (will use a
showarda5288b42009-07-28 20:06:08 +0000124 LEFT OUTER JOIN and an appropriate WHERE condition)
showardc4780402009-08-31 18:31:34 +0000125 @param force_left_join - if true, a LEFT OUTER JOIN will be used
126 instead of an INNER JOIN regardless of other options
showard0957a842009-05-11 19:25:08 +0000127 """
showard7e67b432010-01-20 01:13:04 +0000128 join_from_table = query_set.model._meta.db_table
129 if join_from_key is None:
130 join_from_key = self.model._meta.pk.name
131 if alias is None:
132 alias = join_table + suffix
133 full_join_key = _quote_name(alias) + '.' + _quote_name(join_key)
134 full_join_condition = '%s = %s.%s' % (full_join_key,
135 _quote_name(join_from_table),
136 _quote_name(join_from_key))
showard43a3d262008-11-12 18:17:05 +0000137 if join_condition:
138 full_join_condition += ' AND (' + join_condition + ')'
139 if exclude or force_left_join:
showarda5288b42009-07-28 20:06:08 +0000140 join_type = query_set.query.LOUTER
showard43a3d262008-11-12 18:17:05 +0000141 else:
showarda5288b42009-07-28 20:06:08 +0000142 join_type = query_set.query.INNER
showard43a3d262008-11-12 18:17:05 +0000143
showardf828c772010-01-25 21:49:42 +0000144 query_set = self.CustomQuery.convert_query(query_set)
showard7e67b432010-01-20 01:13:04 +0000145 query_set.query.add_custom_join(join_table,
146 full_join_condition,
147 join_type,
148 condition_values=join_condition_values,
149 alias=alias)
showard43a3d262008-11-12 18:17:05 +0000150
showard7e67b432010-01-20 01:13:04 +0000151 if exclude:
152 query_set = query_set.extra(where=[full_join_key + ' IS NULL'])
153
154 return query_set
155
156
157 def _info_for_many_to_one_join(self, field, join_to_query, alias):
158 """
159 @param field: the ForeignKey field on the related model
160 @param join_to_query: the query over the related model that we're
161 joining to
162 @param alias: alias of joined table
163 """
164 info = {}
165 rhs_table = join_to_query.model._meta.db_table
166 info['rhs_table'] = rhs_table
167 info['rhs_column'] = field.column
168 info['lhs_column'] = field.rel.get_related_field().column
169 rhs_where = join_to_query.query.where
170 rhs_where.relabel_aliases({rhs_table: alias})
Dale Curtis74a314b2011-06-23 14:55:46 -0700171 compiler = join_to_query.query.get_compiler(using=join_to_query.db)
172 initial_clause, values = compiler.as_sql()
173 all_clauses = (initial_clause,)
174 if hasattr(join_to_query.query, 'extra_where'):
175 all_clauses += join_to_query.query.extra_where
176 info['where_clause'] = (
177 ' AND '.join('(%s)' % clause for clause in all_clauses))
showard7e67b432010-01-20 01:13:04 +0000178 info['values'] = values
179 return info
180
181
182 def _info_for_many_to_many_join(self, m2m_field, join_to_query, alias,
183 m2m_is_on_this_model):
184 """
185 @param m2m_field: a Django field representing the M2M relationship.
186 It uses a pivot table with the following structure:
187 this model table <---> M2M pivot table <---> joined model table
188 @param join_to_query: the query over the related model that we're
189 joining to.
190 @param alias: alias of joined table
191 """
192 if m2m_is_on_this_model:
193 # referenced field on this model
194 lhs_id_field = self.model._meta.pk
195 # foreign key on the pivot table referencing lhs_id_field
196 m2m_lhs_column = m2m_field.m2m_column_name()
197 # foreign key on the pivot table referencing rhd_id_field
198 m2m_rhs_column = m2m_field.m2m_reverse_name()
199 # referenced field on related model
200 rhs_id_field = m2m_field.rel.get_related_field()
201 else:
202 lhs_id_field = m2m_field.rel.get_related_field()
203 m2m_lhs_column = m2m_field.m2m_reverse_name()
204 m2m_rhs_column = m2m_field.m2m_column_name()
205 rhs_id_field = join_to_query.model._meta.pk
206
207 info = {}
208 info['rhs_table'] = m2m_field.m2m_db_table()
209 info['rhs_column'] = m2m_lhs_column
210 info['lhs_column'] = lhs_id_field.column
211
212 # select the ID of related models relevant to this join. we can only do
213 # a single join, so we need to gather this information up front and
214 # include it in the join condition.
215 rhs_ids = join_to_query.values_list(rhs_id_field.attname, flat=True)
216 assert len(rhs_ids) == 1, ('Many-to-many custom field joins can only '
217 'match a single related object.')
218 rhs_id = rhs_ids[0]
219
220 info['where_clause'] = '%s.%s = %s' % (_quote_name(alias),
221 _quote_name(m2m_rhs_column),
222 rhs_id)
223 info['values'] = ()
224 return info
225
226
227 def join_custom_field(self, query_set, join_to_query, alias,
228 left_join=True):
229 """Join to a related model to create a custom field in the given query.
230
231 This method is used to construct a custom field on the given query based
232 on a many-valued relationsip. join_to_query should be a simple query
233 (no joins) on the related model which returns at most one related row
234 per instance of this model.
235
236 For many-to-one relationships, the joined table contains the matching
237 row from the related model it one is related, NULL otherwise.
238
239 For many-to-many relationships, the joined table contains the matching
240 row if it's related, NULL otherwise.
241 """
242 relationship_type, field = self.determine_relationship(
243 join_to_query.model)
244
245 if relationship_type == self.MANY_TO_ONE:
246 info = self._info_for_many_to_one_join(field, join_to_query, alias)
247 elif relationship_type == self.M2M_ON_RELATED_MODEL:
248 info = self._info_for_many_to_many_join(
249 m2m_field=field, join_to_query=join_to_query, alias=alias,
250 m2m_is_on_this_model=False)
251 elif relationship_type ==self.M2M_ON_THIS_MODEL:
252 info = self._info_for_many_to_many_join(
253 m2m_field=field, join_to_query=join_to_query, alias=alias,
254 m2m_is_on_this_model=True)
255
256 return self.add_join(query_set, info['rhs_table'], info['rhs_column'],
257 join_from_key=info['lhs_column'],
258 join_condition=info['where_clause'],
259 join_condition_values=info['values'],
260 alias=alias,
261 force_left_join=left_join)
262
263
showardf828c772010-01-25 21:49:42 +0000264 def key_on_joined_table(self, join_to_query):
265 """Get a non-null column on the table joined for the given query.
266
267 This analyzes the join that would be produced if join_to_query were
268 passed to join_custom_field.
269 """
270 relationship_type, field = self.determine_relationship(
271 join_to_query.model)
272 if relationship_type == self.MANY_TO_ONE:
273 return join_to_query.model._meta.pk.column
274 return field.m2m_column_name() # any column on the M2M table will do
275
276
showard7e67b432010-01-20 01:13:04 +0000277 def add_where(self, query_set, where, values=()):
278 query_set = query_set.all()
279 query_set.query.where.add(self._WhereClause(where, values),
280 django.db.models.sql.where.AND)
showardc4780402009-08-31 18:31:34 +0000281 return query_set
showard7c785282008-05-29 19:45:12 +0000282
283
showardeaccf8f2009-04-16 03:11:33 +0000284 def _get_quoted_field(self, table, field):
showarda5288b42009-07-28 20:06:08 +0000285 return _quote_name(table) + '.' + _quote_name(field)
showard5ef36e92008-07-02 16:37:09 +0000286
287
showard7c199df2008-10-03 10:17:15 +0000288 def get_key_on_this_table(self, key_field=None):
showard5ef36e92008-07-02 16:37:09 +0000289 if key_field is None:
290 # default to primary key
291 key_field = self.model._meta.pk.column
292 return self._get_quoted_field(self.model._meta.db_table, key_field)
293
294
showardeaccf8f2009-04-16 03:11:33 +0000295 def escape_user_sql(self, sql):
296 return sql.replace('%', '%%')
297
showard5ef36e92008-07-02 16:37:09 +0000298
showard0957a842009-05-11 19:25:08 +0000299 def _custom_select_query(self, query_set, selects):
Jakob Juelich7bef8412014-10-14 19:11:54 -0700300 """Execute a custom select query.
301
302 @param query_set: query set as returned by query_objects.
303 @param selects: Tables/Columns to select, e.g. tko_test_labels_list.id.
304
305 @returns: Result of the query as returned by cursor.fetchall().
306 """
Dale Curtis74a314b2011-06-23 14:55:46 -0700307 compiler = query_set.query.get_compiler(using=query_set.db)
308 sql, params = compiler.as_sql()
showarda5288b42009-07-28 20:06:08 +0000309 from_ = sql[sql.find(' FROM'):]
310
311 if query_set.query.distinct:
showard0957a842009-05-11 19:25:08 +0000312 distinct = 'DISTINCT '
313 else:
314 distinct = ''
showarda5288b42009-07-28 20:06:08 +0000315
316 sql_query = ('SELECT ' + distinct + ','.join(selects) + from_)
Jakob Juelich7bef8412014-10-14 19:11:54 -0700317 # Chose the connection that's responsible for this type of object
318 cursor = connections[query_set.db].cursor()
showard0957a842009-05-11 19:25:08 +0000319 cursor.execute(sql_query, params)
320 return cursor.fetchall()
321
322
showard68693f72009-05-20 00:31:53 +0000323 def _is_relation_to(self, field, model_class):
324 return field.rel and field.rel.to is model_class
showard0957a842009-05-11 19:25:08 +0000325
326
showard7e67b432010-01-20 01:13:04 +0000327 MANY_TO_ONE = object()
328 M2M_ON_RELATED_MODEL = object()
329 M2M_ON_THIS_MODEL = object()
330
331 def determine_relationship(self, related_model):
332 """
333 Determine the relationship between this model and related_model.
334
335 related_model must have some sort of many-valued relationship to this
336 manager's model.
337 @returns (relationship_type, field), where relationship_type is one of
338 MANY_TO_ONE, M2M_ON_RELATED_MODEL, M2M_ON_THIS_MODEL, and field
339 is the Django field object for the relationship.
340 """
341 # look for a foreign key field on related_model relating to this model
342 for field in related_model._meta.fields:
343 if self._is_relation_to(field, self.model):
344 return self.MANY_TO_ONE, field
345
346 # look for an M2M field on related_model relating to this model
347 for field in related_model._meta.many_to_many:
348 if self._is_relation_to(field, self.model):
349 return self.M2M_ON_RELATED_MODEL, field
350
351 # maybe this model has the many-to-many field
352 for field in self.model._meta.many_to_many:
353 if self._is_relation_to(field, related_model):
354 return self.M2M_ON_THIS_MODEL, field
355
356 raise ValueError('%s has no relation to %s' %
357 (related_model, self.model))
358
359
showard68693f72009-05-20 00:31:53 +0000360 def _get_pivot_iterator(self, base_objects_by_id, related_model):
showard0957a842009-05-11 19:25:08 +0000361 """
showard68693f72009-05-20 00:31:53 +0000362 Determine the relationship between this model and related_model, and
363 return a pivot iterator.
364 @param base_objects_by_id: dict of instances of this model indexed by
365 their IDs
366 @returns a pivot iterator, which yields a tuple (base_object,
367 related_object) for each relationship between a base object and a
368 related object. all base_object instances come from base_objects_by_id.
showard7e67b432010-01-20 01:13:04 +0000369 Note -- this depends on Django model internals.
showard0957a842009-05-11 19:25:08 +0000370 """
showard7e67b432010-01-20 01:13:04 +0000371 relationship_type, field = self.determine_relationship(related_model)
372 if relationship_type == self.MANY_TO_ONE:
373 return self._many_to_one_pivot(base_objects_by_id,
374 related_model, field)
375 elif relationship_type == self.M2M_ON_RELATED_MODEL:
376 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000377 base_objects_by_id, related_model, field.m2m_db_table(),
378 field.m2m_reverse_name(), field.m2m_column_name())
showard7e67b432010-01-20 01:13:04 +0000379 else:
380 assert relationship_type == self.M2M_ON_THIS_MODEL
381 return self._many_to_many_pivot(
showard68693f72009-05-20 00:31:53 +0000382 base_objects_by_id, related_model, field.m2m_db_table(),
383 field.m2m_column_name(), field.m2m_reverse_name())
showard0957a842009-05-11 19:25:08 +0000384
showard0957a842009-05-11 19:25:08 +0000385
showard68693f72009-05-20 00:31:53 +0000386 def _many_to_one_pivot(self, base_objects_by_id, related_model,
387 foreign_key_field):
388 """
389 @returns a pivot iterator - see _get_pivot_iterator()
390 """
391 filter_data = {foreign_key_field.name + '__pk__in':
392 base_objects_by_id.keys()}
393 for related_object in related_model.objects.filter(**filter_data):
showarda5a72c92009-08-20 23:35:21 +0000394 # lookup base object in the dict, rather than grabbing it from the
395 # related object. we need to return instances from the dict, not
396 # fresh instances of the same models (and grabbing model instances
397 # from the related models incurs a DB query each time).
398 base_object_id = getattr(related_object, foreign_key_field.attname)
399 base_object = base_objects_by_id[base_object_id]
showard68693f72009-05-20 00:31:53 +0000400 yield base_object, related_object
401
402
403 def _query_pivot_table(self, base_objects_by_id, pivot_table,
Jakob Juelich7bef8412014-10-14 19:11:54 -0700404 pivot_from_field, pivot_to_field, related_model):
showard0957a842009-05-11 19:25:08 +0000405 """
406 @param id_list list of IDs of self.model objects to include
407 @param pivot_table the name of the pivot table
408 @param pivot_from_field a field name on pivot_table referencing
409 self.model
410 @param pivot_to_field a field name on pivot_table referencing the
411 related model.
Jakob Juelich7bef8412014-10-14 19:11:54 -0700412 @param related_model the related model
413
showard68693f72009-05-20 00:31:53 +0000414 @returns pivot list of IDs (base_id, related_id)
showard0957a842009-05-11 19:25:08 +0000415 """
416 query = """
417 SELECT %(from_field)s, %(to_field)s
418 FROM %(table)s
419 WHERE %(from_field)s IN (%(id_list)s)
420 """ % dict(from_field=pivot_from_field,
421 to_field=pivot_to_field,
422 table=pivot_table,
showard68693f72009-05-20 00:31:53 +0000423 id_list=','.join(str(id_) for id_
424 in base_objects_by_id.iterkeys()))
Jakob Juelich7bef8412014-10-14 19:11:54 -0700425
426 # Chose the connection that's responsible for this type of object
427 # The databases for related_model and the current model will always
428 # be the same, related_model is just easier to obtain here because
429 # self is only a ExtendedManager, not the object.
430 cursor = connections[related_model.objects.db].cursor()
showard0957a842009-05-11 19:25:08 +0000431 cursor.execute(query)
showard68693f72009-05-20 00:31:53 +0000432 return cursor.fetchall()
showard0957a842009-05-11 19:25:08 +0000433
434
showard68693f72009-05-20 00:31:53 +0000435 def _many_to_many_pivot(self, base_objects_by_id, related_model,
436 pivot_table, pivot_from_field, pivot_to_field):
437 """
438 @param pivot_table: see _query_pivot_table
439 @param pivot_from_field: see _query_pivot_table
440 @param pivot_to_field: see _query_pivot_table
441 @returns a pivot iterator - see _get_pivot_iterator()
442 """
443 id_pivot = self._query_pivot_table(base_objects_by_id, pivot_table,
Jakob Juelich7bef8412014-10-14 19:11:54 -0700444 pivot_from_field, pivot_to_field,
445 related_model)
showard68693f72009-05-20 00:31:53 +0000446
447 all_related_ids = list(set(related_id for base_id, related_id
448 in id_pivot))
449 related_objects_by_id = related_model.objects.in_bulk(all_related_ids)
450
451 for base_id, related_id in id_pivot:
452 yield base_objects_by_id[base_id], related_objects_by_id[related_id]
453
454
455 def populate_relationships(self, base_objects, related_model,
showard0957a842009-05-11 19:25:08 +0000456 related_list_name):
457 """
showard68693f72009-05-20 00:31:53 +0000458 For each instance of this model in base_objects, add a field named
459 related_list_name listing all the related objects of type related_model.
460 related_model must be in a many-to-one or many-to-many relationship with
461 this model.
462 @param base_objects - list of instances of this model
463 @param related_model - model class related to this model
464 @param related_list_name - attribute name in which to store the related
465 object list.
showard0957a842009-05-11 19:25:08 +0000466 """
showard68693f72009-05-20 00:31:53 +0000467 if not base_objects:
showard0957a842009-05-11 19:25:08 +0000468 # if we don't bail early, we'll get a SQL error later
469 return
showard0957a842009-05-11 19:25:08 +0000470
showard68693f72009-05-20 00:31:53 +0000471 base_objects_by_id = dict((base_object._get_pk_val(), base_object)
472 for base_object in base_objects)
473 pivot_iterator = self._get_pivot_iterator(base_objects_by_id,
474 related_model)
showard0957a842009-05-11 19:25:08 +0000475
showard68693f72009-05-20 00:31:53 +0000476 for base_object in base_objects:
477 setattr(base_object, related_list_name, [])
478
479 for base_object, related_object in pivot_iterator:
480 getattr(base_object, related_list_name).append(related_object)
showard0957a842009-05-11 19:25:08 +0000481
482
jamesrene3656232010-03-02 00:00:30 +0000483class ModelWithInvalidQuerySet(dbmodels.query.QuerySet):
484 """
485 QuerySet that handles delete() properly for models with an "invalid" bit
486 """
487 def delete(self):
488 for model in self:
489 model.delete()
490
491
492class ModelWithInvalidManager(ExtendedManager):
493 """
494 Manager for objects with an "invalid" bit
495 """
496 def get_query_set(self):
497 return ModelWithInvalidQuerySet(self.model)
498
499
500class ValidObjectsManager(ModelWithInvalidManager):
jadmanski0afbb632008-06-06 21:10:57 +0000501 """
502 Manager returning only objects with invalid=False.
503 """
504 def get_query_set(self):
505 queryset = super(ValidObjectsManager, self).get_query_set()
506 return queryset.filter(invalid=False)
showard7c785282008-05-29 19:45:12 +0000507
508
Prashanth B489b91d2014-03-15 12:17:16 -0700509class ModelExtensions(rdb_model_extensions.ModelValidators):
jadmanski0afbb632008-06-06 21:10:57 +0000510 """\
Prashanth B489b91d2014-03-15 12:17:16 -0700511 Mixin with convenience functions for models, built on top of
512 the model validators in rdb_model_extensions.
jadmanski0afbb632008-06-06 21:10:57 +0000513 """
514 # TODO: at least some of these functions really belong in a custom
515 # Manager class
showard7c785282008-05-29 19:45:12 +0000516
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700517
518 SERIALIZATION_LINKS_TO_FOLLOW = set()
519 """
520 To be able to send jobs and hosts to shards, it's necessary to find their
521 dependencies.
522 The most generic approach for this would be to traverse all relationships
523 to other objects recursively. This would list all objects that are related
524 in any way.
525 But this approach finds too many objects: If a host should be transferred,
526 all it's relationships would be traversed. This would find an acl group.
527 If then the acl group's relationships are traversed, the relationship
528 would be followed backwards and many other hosts would be found.
529
530 This mapping tells that algorithm which relations to follow explicitly.
531 """
532
Jakob Juelichf865d332014-09-29 10:47:49 -0700533
534 SERIALIZATION_LOCAL_LINKS_TO_UPDATE = set()
535 """
536 On deserializion, if the object to persist already exists, local fields
537 will only be updated, if their name is in this set.
538 """
539
540
jadmanski0afbb632008-06-06 21:10:57 +0000541 @classmethod
542 def convert_human_readable_values(cls, data, to_human_readable=False):
543 """\
544 Performs conversions on user-supplied field data, to make it
545 easier for users to pass human-readable data.
showard7c785282008-05-29 19:45:12 +0000546
jadmanski0afbb632008-06-06 21:10:57 +0000547 For all fields that have choice sets, convert their values
548 from human-readable strings to enum values, if necessary. This
549 allows users to pass strings instead of the corresponding
550 integer values.
showard7c785282008-05-29 19:45:12 +0000551
jadmanski0afbb632008-06-06 21:10:57 +0000552 For all foreign key fields, call smart_get with the supplied
553 data. This allows the user to pass either an ID value or
554 the name of the object as a string.
showard7c785282008-05-29 19:45:12 +0000555
jadmanski0afbb632008-06-06 21:10:57 +0000556 If to_human_readable=True, perform the inverse - i.e. convert
557 numeric values to human readable values.
showard7c785282008-05-29 19:45:12 +0000558
jadmanski0afbb632008-06-06 21:10:57 +0000559 This method modifies data in-place.
560 """
561 field_dict = cls.get_field_dict()
562 for field_name in data:
showarde732ee72008-09-23 19:15:43 +0000563 if field_name not in field_dict or data[field_name] is None:
jadmanski0afbb632008-06-06 21:10:57 +0000564 continue
565 field_obj = field_dict[field_name]
566 # convert enum values
567 if field_obj.choices:
568 for choice_data in field_obj.choices:
569 # choice_data is (value, name)
570 if to_human_readable:
571 from_val, to_val = choice_data
572 else:
573 to_val, from_val = choice_data
574 if from_val == data[field_name]:
575 data[field_name] = to_val
576 break
577 # convert foreign key values
578 elif field_obj.rel:
showarda4ea5742009-02-17 20:56:23 +0000579 dest_obj = field_obj.rel.to.smart_get(data[field_name],
580 valid_only=False)
showardf8b19042009-05-12 17:22:49 +0000581 if to_human_readable:
Paul Pendlebury5a8c6ad2011-02-01 07:20:17 -0800582 # parameterized_jobs do not have a name_field
583 if (field_name != 'parameterized_job' and
584 dest_obj.name_field is not None):
showardf8b19042009-05-12 17:22:49 +0000585 data[field_name] = getattr(dest_obj,
586 dest_obj.name_field)
jadmanski0afbb632008-06-06 21:10:57 +0000587 else:
showardb0a73032009-03-27 18:35:41 +0000588 data[field_name] = dest_obj
showard7c785282008-05-29 19:45:12 +0000589
590
showard7c785282008-05-29 19:45:12 +0000591
592
Dale Curtis74a314b2011-06-23 14:55:46 -0700593 def _validate_unique(self):
jadmanski0afbb632008-06-06 21:10:57 +0000594 """\
595 Validate that unique fields are unique. Django manipulators do
596 this too, but they're a huge pain to use manually. Trust me.
597 """
598 errors = {}
599 cls = type(self)
600 field_dict = self.get_field_dict()
601 manager = cls.get_valid_manager()
602 for field_name, field_obj in field_dict.iteritems():
603 if not field_obj.unique:
604 continue
showard7c785282008-05-29 19:45:12 +0000605
jadmanski0afbb632008-06-06 21:10:57 +0000606 value = getattr(self, field_name)
showardbd18ab72009-09-18 21:20:27 +0000607 if value is None and field_obj.auto_created:
608 # don't bother checking autoincrement fields about to be
609 # generated
610 continue
611
jadmanski0afbb632008-06-06 21:10:57 +0000612 existing_objs = manager.filter(**{field_name : value})
613 num_existing = existing_objs.count()
showard7c785282008-05-29 19:45:12 +0000614
jadmanski0afbb632008-06-06 21:10:57 +0000615 if num_existing == 0:
616 continue
617 if num_existing == 1 and existing_objs[0].id == self.id:
618 continue
619 errors[field_name] = (
620 'This value must be unique (%s)' % (value))
621 return errors
showard7c785282008-05-29 19:45:12 +0000622
623
showarda5288b42009-07-28 20:06:08 +0000624 def _validate(self):
625 """
626 First coerces all fields on this instance to their proper Python types.
627 Then runs validation on every field. Returns a dictionary of
628 field_name -> error_list.
629
630 Based on validate() from django.db.models.Model in Django 0.96, which
631 was removed in Django 1.0. It should reappear in a later version. See:
632 http://code.djangoproject.com/ticket/6845
633 """
634 error_dict = {}
635 for f in self._meta.fields:
636 try:
637 python_value = f.to_python(
638 getattr(self, f.attname, f.get_default()))
639 except django.core.exceptions.ValidationError, e:
jamesren1e0a4ce2010-04-21 17:45:11 +0000640 error_dict[f.name] = str(e)
showarda5288b42009-07-28 20:06:08 +0000641 continue
642
643 if not f.blank and not python_value:
644 error_dict[f.name] = 'This field is required.'
645 continue
646
647 setattr(self, f.attname, python_value)
648
649 return error_dict
650
651
jadmanski0afbb632008-06-06 21:10:57 +0000652 def do_validate(self):
showarda5288b42009-07-28 20:06:08 +0000653 errors = self._validate()
Dale Curtis74a314b2011-06-23 14:55:46 -0700654 unique_errors = self._validate_unique()
jadmanski0afbb632008-06-06 21:10:57 +0000655 for field_name, error in unique_errors.iteritems():
656 errors.setdefault(field_name, error)
657 if errors:
658 raise ValidationError(errors)
showard7c785282008-05-29 19:45:12 +0000659
660
jadmanski0afbb632008-06-06 21:10:57 +0000661 # actually (externally) useful methods follow
showard7c785282008-05-29 19:45:12 +0000662
jadmanski0afbb632008-06-06 21:10:57 +0000663 @classmethod
664 def add_object(cls, data={}, **kwargs):
665 """\
666 Returns a new object created with the given data (a dictionary
667 mapping field names to values). Merges any extra keyword args
668 into data.
669 """
Prashanth B489b91d2014-03-15 12:17:16 -0700670 data = dict(data)
671 data.update(kwargs)
672 data = cls.prepare_data_args(data)
673 cls.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000674 data = cls.provide_default_values(data)
Prashanth B489b91d2014-03-15 12:17:16 -0700675
jadmanski0afbb632008-06-06 21:10:57 +0000676 obj = cls(**data)
677 obj.do_validate()
678 obj.save()
679 return obj
showard7c785282008-05-29 19:45:12 +0000680
681
jadmanski0afbb632008-06-06 21:10:57 +0000682 def update_object(self, data={}, **kwargs):
683 """\
684 Updates the object with the given data (a dictionary mapping
685 field names to values). Merges any extra keyword args into
686 data.
687 """
Prashanth B489b91d2014-03-15 12:17:16 -0700688 data = dict(data)
689 data.update(kwargs)
690 data = self.prepare_data_args(data)
691 self.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000692 for field_name, value in data.iteritems():
showardb0a73032009-03-27 18:35:41 +0000693 setattr(self, field_name, value)
jadmanski0afbb632008-06-06 21:10:57 +0000694 self.do_validate()
695 self.save()
showard7c785282008-05-29 19:45:12 +0000696
697
showard8bfb5cb2009-10-07 20:49:15 +0000698 # see query_objects()
699 _SPECIAL_FILTER_KEYS = ('query_start', 'query_limit', 'sort_by',
700 'extra_args', 'extra_where', 'no_distinct')
701
702
jadmanski0afbb632008-06-06 21:10:57 +0000703 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000704 def _extract_special_params(cls, filter_data):
705 """
706 @returns a tuple of dicts (special_params, regular_filters), where
707 special_params contains the parameters we handle specially and
708 regular_filters is the remaining data to be handled by Django.
709 """
710 regular_filters = dict(filter_data)
711 special_params = {}
712 for key in cls._SPECIAL_FILTER_KEYS:
713 if key in regular_filters:
714 special_params[key] = regular_filters.pop(key)
715 return special_params, regular_filters
716
717
718 @classmethod
719 def apply_presentation(cls, query, filter_data):
720 """
721 Apply presentation parameters -- sorting and paging -- to the given
722 query.
723 @returns new query with presentation applied
724 """
725 special_params, _ = cls._extract_special_params(filter_data)
726 sort_by = special_params.get('sort_by', None)
727 if sort_by:
728 assert isinstance(sort_by, list) or isinstance(sort_by, tuple)
showard8b0ea222009-12-23 19:23:03 +0000729 query = query.extra(order_by=sort_by)
showard8bfb5cb2009-10-07 20:49:15 +0000730
731 query_start = special_params.get('query_start', None)
732 query_limit = special_params.get('query_limit', None)
733 if query_start is not None:
734 if query_limit is None:
735 raise ValueError('Cannot pass query_start without query_limit')
736 # query_limit is passed as a page size
showard7074b742009-10-12 20:30:04 +0000737 query_limit += query_start
738 return query[query_start:query_limit]
showard8bfb5cb2009-10-07 20:49:15 +0000739
740
741 @classmethod
742 def query_objects(cls, filter_data, valid_only=True, initial_query=None,
743 apply_presentation=True):
jadmanski0afbb632008-06-06 21:10:57 +0000744 """\
745 Returns a QuerySet object for querying the given model_class
746 with the given filter_data. Optional special arguments in
747 filter_data include:
748 -query_start: index of first return to return
749 -query_limit: maximum number of results to return
750 -sort_by: list of fields to sort on. prefixing a '-' onto a
751 field name changes the sort to descending order.
752 -extra_args: keyword args to pass to query.extra() (see Django
753 DB layer documentation)
showarda5288b42009-07-28 20:06:08 +0000754 -extra_where: extra WHERE clause to append
showard8bfb5cb2009-10-07 20:49:15 +0000755 -no_distinct: if True, a DISTINCT will not be added to the SELECT
jadmanski0afbb632008-06-06 21:10:57 +0000756 """
showard8bfb5cb2009-10-07 20:49:15 +0000757 special_params, regular_filters = cls._extract_special_params(
758 filter_data)
showard7c785282008-05-29 19:45:12 +0000759
showard7ac7b7a2008-07-21 20:24:29 +0000760 if initial_query is None:
761 if valid_only:
762 initial_query = cls.get_valid_manager()
763 else:
764 initial_query = cls.objects
showard8bfb5cb2009-10-07 20:49:15 +0000765
766 query = initial_query.filter(**regular_filters)
767
768 use_distinct = not special_params.get('no_distinct', False)
showard7ac7b7a2008-07-21 20:24:29 +0000769 if use_distinct:
770 query = query.distinct()
showard7c785282008-05-29 19:45:12 +0000771
showard8bfb5cb2009-10-07 20:49:15 +0000772 extra_args = special_params.get('extra_args', {})
773 extra_where = special_params.get('extra_where', None)
774 if extra_where:
775 # escape %'s
776 extra_where = cls.objects.escape_user_sql(extra_where)
777 extra_args.setdefault('where', []).append(extra_where)
jadmanski0afbb632008-06-06 21:10:57 +0000778 if extra_args:
779 query = query.extra(**extra_args)
Jakob Juelich7bef8412014-10-14 19:11:54 -0700780 # TODO: Use readonly connection for these queries.
781 # This has been disabled, because it's not used anyway, as the
782 # configured readonly user is the same as the real user anyway.
showard7c785282008-05-29 19:45:12 +0000783
showard8bfb5cb2009-10-07 20:49:15 +0000784 if apply_presentation:
785 query = cls.apply_presentation(query, filter_data)
786
787 return query
showard7c785282008-05-29 19:45:12 +0000788
789
jadmanski0afbb632008-06-06 21:10:57 +0000790 @classmethod
showard585c2ab2008-07-23 19:29:49 +0000791 def query_count(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000792 """\
793 Like query_objects, but retreive only the count of results.
794 """
795 filter_data.pop('query_start', None)
796 filter_data.pop('query_limit', None)
showard585c2ab2008-07-23 19:29:49 +0000797 query = cls.query_objects(filter_data, initial_query=initial_query)
798 return query.count()
showard7c785282008-05-29 19:45:12 +0000799
800
jadmanski0afbb632008-06-06 21:10:57 +0000801 @classmethod
802 def clean_object_dicts(cls, field_dicts):
803 """\
804 Take a list of dicts corresponding to object (as returned by
805 query.values()) and clean the data to be more suitable for
806 returning to the user.
807 """
showarde732ee72008-09-23 19:15:43 +0000808 for field_dict in field_dicts:
809 cls.clean_foreign_keys(field_dict)
showard21baa452008-10-21 00:08:39 +0000810 cls._convert_booleans(field_dict)
showarde732ee72008-09-23 19:15:43 +0000811 cls.convert_human_readable_values(field_dict,
812 to_human_readable=True)
showard7c785282008-05-29 19:45:12 +0000813
814
jadmanski0afbb632008-06-06 21:10:57 +0000815 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000816 def list_objects(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000817 """\
818 Like query_objects, but return a list of dictionaries.
819 """
showard7ac7b7a2008-07-21 20:24:29 +0000820 query = cls.query_objects(filter_data, initial_query=initial_query)
showard8bfb5cb2009-10-07 20:49:15 +0000821 extra_fields = query.query.extra_select.keys()
822 field_dicts = [model_object.get_object_dict(extra_fields=extra_fields)
showarde732ee72008-09-23 19:15:43 +0000823 for model_object in query]
jadmanski0afbb632008-06-06 21:10:57 +0000824 return field_dicts
showard7c785282008-05-29 19:45:12 +0000825
826
jadmanski0afbb632008-06-06 21:10:57 +0000827 @classmethod
showarda4ea5742009-02-17 20:56:23 +0000828 def smart_get(cls, id_or_name, valid_only=True):
jadmanski0afbb632008-06-06 21:10:57 +0000829 """\
830 smart_get(integer) -> get object by ID
831 smart_get(string) -> get object by name_field
jadmanski0afbb632008-06-06 21:10:57 +0000832 """
showarda4ea5742009-02-17 20:56:23 +0000833 if valid_only:
834 manager = cls.get_valid_manager()
835 else:
836 manager = cls.objects
837
838 if isinstance(id_or_name, (int, long)):
839 return manager.get(pk=id_or_name)
jamesren3e9f6092010-03-11 21:32:10 +0000840 if isinstance(id_or_name, basestring) and hasattr(cls, 'name_field'):
showarda4ea5742009-02-17 20:56:23 +0000841 return manager.get(**{cls.name_field : id_or_name})
842 raise ValueError(
843 'Invalid positional argument: %s (%s)' % (id_or_name,
844 type(id_or_name)))
showard7c785282008-05-29 19:45:12 +0000845
846
showardbe3ec042008-11-12 18:16:07 +0000847 @classmethod
848 def smart_get_bulk(cls, id_or_name_list):
849 invalid_inputs = []
850 result_objects = []
851 for id_or_name in id_or_name_list:
852 try:
853 result_objects.append(cls.smart_get(id_or_name))
854 except cls.DoesNotExist:
855 invalid_inputs.append(id_or_name)
856 if invalid_inputs:
mbligh7a3ebe32008-12-01 17:10:33 +0000857 raise cls.DoesNotExist('The following %ss do not exist: %s'
858 % (cls.__name__.lower(),
859 ', '.join(invalid_inputs)))
showardbe3ec042008-11-12 18:16:07 +0000860 return result_objects
861
862
showard8bfb5cb2009-10-07 20:49:15 +0000863 def get_object_dict(self, extra_fields=None):
jadmanski0afbb632008-06-06 21:10:57 +0000864 """\
showard8bfb5cb2009-10-07 20:49:15 +0000865 Return a dictionary mapping fields to this object's values. @param
866 extra_fields: list of extra attribute names to include, in addition to
867 the fields defined on this object.
jadmanski0afbb632008-06-06 21:10:57 +0000868 """
showard8bfb5cb2009-10-07 20:49:15 +0000869 fields = self.get_field_dict().keys()
870 if extra_fields:
871 fields += extra_fields
jadmanski0afbb632008-06-06 21:10:57 +0000872 object_dict = dict((field_name, getattr(self, field_name))
showarde732ee72008-09-23 19:15:43 +0000873 for field_name in fields)
jadmanski0afbb632008-06-06 21:10:57 +0000874 self.clean_object_dicts([object_dict])
showardd3dc1992009-04-22 21:01:40 +0000875 self._postprocess_object_dict(object_dict)
jadmanski0afbb632008-06-06 21:10:57 +0000876 return object_dict
showard7c785282008-05-29 19:45:12 +0000877
878
showardd3dc1992009-04-22 21:01:40 +0000879 def _postprocess_object_dict(self, object_dict):
880 """For subclasses to override."""
881 pass
882
883
jadmanski0afbb632008-06-06 21:10:57 +0000884 @classmethod
885 def get_valid_manager(cls):
886 return cls.objects
showard7c785282008-05-29 19:45:12 +0000887
888
showard2bab8f42008-11-12 18:15:22 +0000889 def _record_attributes(self, attributes):
890 """
891 See on_attribute_changed.
892 """
893 assert not isinstance(attributes, basestring)
894 self._recorded_attributes = dict((attribute, getattr(self, attribute))
895 for attribute in attributes)
896
897
898 def _check_for_updated_attributes(self):
899 """
900 See on_attribute_changed.
901 """
902 for attribute, original_value in self._recorded_attributes.iteritems():
903 new_value = getattr(self, attribute)
904 if original_value != new_value:
905 self.on_attribute_changed(attribute, original_value)
906 self._record_attributes(self._recorded_attributes.keys())
907
908
909 def on_attribute_changed(self, attribute, old_value):
910 """
911 Called whenever an attribute is updated. To be overridden.
912
913 To use this method, you must:
914 * call _record_attributes() from __init__() (after making the super
915 call) with a list of attributes for which you want to be notified upon
916 change.
917 * call _check_for_updated_attributes() from save().
918 """
919 pass
920
921
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700922 def serialize(self, include_dependencies=True):
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700923 """Serializes the object with dependencies.
924
925 The variable SERIALIZATION_LINKS_TO_FOLLOW defines which dependencies
926 this function will serialize with the object.
927
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700928 @param include_dependencies: Whether or not to follow relations to
929 objects this object depends on.
930 This parameter is used when uploading
931 jobs from a shard to the master, as the
932 master already has all the dependent
933 objects.
934
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700935 @returns: Dictionary representation of the object.
936 """
937 serialized = {}
938 for field in self._meta.concrete_model._meta.local_fields:
939 if field.rel is None:
940 serialized[field.name] = field._get_val_from_obj(self)
941
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700942 if include_dependencies:
943 for link in self.SERIALIZATION_LINKS_TO_FOLLOW:
944 serialized[link] = self._serialize_relation(link)
Jakob Juelich3bb7c802014-09-02 16:31:11 -0700945
946 return serialized
947
948
949 def _serialize_relation(self, link):
950 """Serializes dependent objects given the name of the relation.
951
952 @param link: Name of the relation to take objects from.
953
954 @returns For To-Many relationships a list of the serialized related
955 objects, for To-One relationships the serialized related object.
956 """
957 try:
958 attr = getattr(self, link)
959 except AttributeError:
960 # One-To-One relationships that point to None may raise this
961 return None
962
963 if attr is None:
964 return None
965 if hasattr(attr, 'all'):
966 return [obj.serialize() for obj in attr.all()]
967 return attr.serialize()
968
969
Jakob Juelichf88fa932014-09-03 17:58:04 -0700970 @classmethod
Jakob Juelich116ff0f2014-09-17 18:25:16 -0700971 def _split_local_from_foreign_values(cls, data):
972 """This splits local from foreign values in a serialized object.
973
974 @param data: The serialized object.
975
976 @returns A tuple of two lists, both containing tuples in the form
977 (link_name, link_value). The first list contains all links
978 for local fields, the second one contains those for foreign
979 fields/objects.
980 """
981 links_to_local_values, links_to_related_values = [], []
982 for link, value in data.iteritems():
983 if link in cls.SERIALIZATION_LINKS_TO_FOLLOW:
984 # It's a foreign key
985 links_to_related_values.append((link, value))
986 else:
987 # It's a local attribute
988 links_to_local_values.append((link, value))
989 return links_to_local_values, links_to_related_values
990
991
Jakob Juelichf865d332014-09-29 10:47:49 -0700992 @classmethod
993 def _filter_update_allowed_fields(cls, data):
994 """Filters data and returns only files that updates are allowed on.
995
996 This is i.e. needed for syncing aborted bits from the master to shards.
997
998 Local links are only allowed to be updated, if they are in
999 SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1000 Overwriting existing values is allowed in order to be able to sync i.e.
1001 the aborted bit from the master to a shard.
1002
1003 The whitelisting mechanism is in place to prevent overwriting local
1004 status: If all fields were overwritten, jobs would be completely be
1005 set back to their original (unstarted) state.
1006
1007 @param data: List with tuples of the form (link_name, link_value), as
1008 returned by _split_local_from_foreign_values.
1009
1010 @returns List of the same format as data, but only containing data for
1011 fields that updates are allowed on.
1012 """
1013 return [pair for pair in data
1014 if pair[0] in cls.SERIALIZATION_LOCAL_LINKS_TO_UPDATE]
1015
1016
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001017 @classmethod
1018 def delete_matching_record(cls, **filter_args):
1019 """Delete records matching the filter.
1020
1021 @param filter_args: Arguments for the django filter
1022 used to locate the record to delete.
1023 """
1024 try:
1025 existing_record = cls.objects.get(**filter_args)
1026 except cls.DoesNotExist:
1027 return
1028 existing_record.delete()
1029
1030
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001031 def _deserialize_local(self, data):
1032 """Set local attributes from a list of tuples.
1033
1034 @param data: List of tuples like returned by
1035 _split_local_from_foreign_values.
1036 """
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001037 if not data:
1038 return
1039
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001040 for link, value in data:
1041 setattr(self, link, value)
1042 # Overwridden save() methods are prone to errors, so don't execute them.
1043 # This is because:
1044 # - the overwritten methods depend on ACL groups that don't yet exist
1045 # and don't handle errors
1046 # - the overwritten methods think this object already exists in the db
1047 # because the id is already set
1048 super(type(self), self).save()
1049
1050
1051 def _deserialize_relations(self, data):
1052 """Set foreign attributes from a list of tuples.
1053
1054 This deserialized the related objects using their own deserialize()
1055 function and then sets the relation.
1056
1057 @param data: List of tuples like returned by
1058 _split_local_from_foreign_values.
1059 """
1060 for link, value in data:
1061 self._deserialize_relation(link, value)
1062 # See comment in _deserialize_local
1063 super(type(self), self).save()
1064
1065
1066 @classmethod
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001067 def get_record(cls, data):
1068 """Retrieve a record with the data in the given input arg.
1069
1070 @param data: A dictionary containing the information to use in a query
1071 for data. If child models have different constraints of
1072 uniqueness they should override this model.
1073
1074 @return: An object with matching data.
1075
1076 @raises DoesNotExist: If a record with the given data doesn't exist.
1077 """
1078 return cls.objects.get(id=data['id'])
1079
1080
1081 @classmethod
Jakob Juelichf88fa932014-09-03 17:58:04 -07001082 def deserialize(cls, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001083 """Recursively deserializes and saves an object with it's dependencies.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001084
1085 This takes the result of the serialize method and creates objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001086 in the database that are just like the original.
1087
1088 If an object of the same type with the same id already exists, it's
Jakob Juelichf865d332014-09-29 10:47:49 -07001089 local values will be left untouched, unless they are explicitly
1090 whitelisted in SERIALIZATION_LOCAL_LINKS_TO_UPDATE.
1091
1092 Deserialize will always recursively propagate to all related objects
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001093 present in data though.
1094 I.e. this is necessary to add users to an already existing acl-group.
Jakob Juelichf88fa932014-09-03 17:58:04 -07001095
1096 @param data: Representation of an object and its dependencies, as
1097 returned by serialize.
1098
1099 @returns: The object represented by data if it didn't exist before,
1100 otherwise the object that existed before and has the same type
1101 and id as the one described by data.
1102 """
1103 if data is None:
1104 return None
1105
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001106 local, related = cls._split_local_from_foreign_values(data)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001107 try:
Prashanth Balasubramanianaf516642014-12-12 18:16:32 -08001108 instance = cls.get_record(data)
Jakob Juelichf865d332014-09-29 10:47:49 -07001109 local = cls._filter_update_allowed_fields(local)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001110 except cls.DoesNotExist:
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001111 instance = cls()
Jakob Juelichf88fa932014-09-03 17:58:04 -07001112
Jakob Juelichf865d332014-09-29 10:47:49 -07001113 instance._deserialize_local(local)
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001114 instance._deserialize_relations(related)
Jakob Juelichf88fa932014-09-03 17:58:04 -07001115
1116 return instance
1117
1118
Jakob Juelicha94efe62014-09-18 16:02:49 -07001119 def sanity_check_update_from_shard(self, shard, updated_serialized,
1120 *args, **kwargs):
1121 """Check if an update sent from a shard is legitimate.
1122
1123 @raises error.UnallowedRecordsSentToMaster if an update is not
1124 legitimate.
1125 """
1126 raise NotImplementedError(
1127 'sanity_check_update_from_shard must be implemented by subclass %s '
1128 'for type %s' % type(self))
1129
1130
Prashanth Balasubramanian75be1d32014-11-25 18:03:09 -08001131 @transaction.commit_on_success
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001132 def update_from_serialized(self, serialized):
1133 """Updates local fields of an existing object from a serialized form.
1134
1135 This is different than the normal deserialize() in the way that it
1136 does update local values, which deserialize doesn't, but doesn't
1137 recursively propagate to related objects, which deserialize() does.
1138
1139 The use case of this function is to update job records on the master
1140 after the jobs have been executed on a slave, as the master is not
1141 interested in updates for users, labels, specialtasks, etc.
1142
1143 @param serialized: Representation of an object and its dependencies, as
1144 returned by serialize.
1145
1146 @raises ValueError: if serialized contains related objects, i.e. not
1147 only local fields.
1148 """
1149 local, related = (
1150 self._split_local_from_foreign_values(serialized))
1151 if related:
1152 raise ValueError('Serialized must not contain foreign '
1153 'objects: %s' % related)
1154
1155 self._deserialize_local(local)
1156
1157
Jakob Juelichf88fa932014-09-03 17:58:04 -07001158 def custom_deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001159 """Allows overriding the deserialization behaviour by subclasses."""
Jakob Juelichf88fa932014-09-03 17:58:04 -07001160 raise NotImplementedError(
1161 'custom_deserialize_relation must be implemented by subclass %s '
1162 'for relation %s' % (type(self), link))
1163
1164
1165 def _deserialize_relation(self, link, data):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001166 """Deserializes related objects and sets references on this object.
1167
1168 Relations that point to a list of objects are handled automatically.
1169 For many-to-one or one-to-one relations custom_deserialize_relation
1170 must be overridden by the subclass.
1171
1172 Related objects are deserialized using their deserialize() method.
1173 Thereby they and their dependencies are created if they don't exist
1174 and saved to the database.
1175
1176 @param link: Name of the relation.
1177 @param data: Serialized representation of the related object(s).
1178 This means a list of dictionaries for to-many relations,
1179 just a dictionary for to-one relations.
1180 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001181 field = getattr(self, link)
1182
1183 if field and hasattr(field, 'all'):
1184 self._deserialize_2m_relation(link, data, field.model)
1185 else:
1186 self.custom_deserialize_relation(link, data)
1187
1188
1189 def _deserialize_2m_relation(self, link, data, related_class):
Jakob Juelich116ff0f2014-09-17 18:25:16 -07001190 """Deserialize related objects for one to-many relationship.
1191
1192 @param link: Name of the relation.
1193 @param data: Serialized representation of the related objects.
1194 This is a list with of dictionaries.
1195 """
Jakob Juelichf88fa932014-09-03 17:58:04 -07001196 relation_set = getattr(self, link)
1197 for serialized in data:
1198 relation_set.add(related_class.deserialize(serialized))
1199
1200
showard7c785282008-05-29 19:45:12 +00001201class ModelWithInvalid(ModelExtensions):
jadmanski0afbb632008-06-06 21:10:57 +00001202 """
1203 Overrides model methods save() and delete() to support invalidation in
1204 place of actual deletion. Subclasses must have a boolean "invalid"
1205 field.
1206 """
showard7c785282008-05-29 19:45:12 +00001207
showarda5288b42009-07-28 20:06:08 +00001208 def save(self, *args, **kwargs):
showardddb90992009-02-11 23:39:32 +00001209 first_time = (self.id is None)
1210 if first_time:
1211 # see if this object was previously added and invalidated
1212 my_name = getattr(self, self.name_field)
1213 filters = {self.name_field : my_name, 'invalid' : True}
1214 try:
1215 old_object = self.__class__.objects.get(**filters)
showardafd97de2009-10-01 18:45:09 +00001216 self.resurrect_object(old_object)
showardddb90992009-02-11 23:39:32 +00001217 except self.DoesNotExist:
1218 # no existing object
1219 pass
showard7c785282008-05-29 19:45:12 +00001220
showarda5288b42009-07-28 20:06:08 +00001221 super(ModelWithInvalid, self).save(*args, **kwargs)
showard7c785282008-05-29 19:45:12 +00001222
1223
showardafd97de2009-10-01 18:45:09 +00001224 def resurrect_object(self, old_object):
1225 """
1226 Called when self is about to be saved for the first time and is actually
1227 "undeleting" a previously deleted object. Can be overridden by
1228 subclasses to copy data as desired from the deleted entry (but this
1229 superclass implementation must normally be called).
1230 """
1231 self.id = old_object.id
1232
1233
jadmanski0afbb632008-06-06 21:10:57 +00001234 def clean_object(self):
1235 """
1236 This method is called when an object is marked invalid.
1237 Subclasses should override this to clean up relationships that
showardafd97de2009-10-01 18:45:09 +00001238 should no longer exist if the object were deleted.
1239 """
jadmanski0afbb632008-06-06 21:10:57 +00001240 pass
showard7c785282008-05-29 19:45:12 +00001241
1242
jadmanski0afbb632008-06-06 21:10:57 +00001243 def delete(self):
Dale Curtis74a314b2011-06-23 14:55:46 -07001244 self.invalid = self.invalid
jadmanski0afbb632008-06-06 21:10:57 +00001245 assert not self.invalid
1246 self.invalid = True
1247 self.save()
1248 self.clean_object()
showard7c785282008-05-29 19:45:12 +00001249
1250
jadmanski0afbb632008-06-06 21:10:57 +00001251 @classmethod
1252 def get_valid_manager(cls):
1253 return cls.valid_objects
showard7c785282008-05-29 19:45:12 +00001254
1255
jadmanski0afbb632008-06-06 21:10:57 +00001256 class Manipulator(object):
1257 """
1258 Force default manipulators to look only at valid objects -
1259 otherwise they will match against invalid objects when checking
1260 uniqueness.
1261 """
1262 @classmethod
1263 def _prepare(cls, model):
1264 super(ModelWithInvalid.Manipulator, cls)._prepare(model)
1265 cls.manager = model.valid_objects
showardf8b19042009-05-12 17:22:49 +00001266
1267
1268class ModelWithAttributes(object):
1269 """
1270 Mixin class for models that have an attribute model associated with them.
1271 The attribute model is assumed to have its value field named "value".
1272 """
1273
1274 def _get_attribute_model_and_args(self, attribute):
1275 """
1276 Subclasses should override this to return a tuple (attribute_model,
1277 keyword_args), where attribute_model is a model class and keyword_args
1278 is a dict of args to pass to attribute_model.objects.get() to get an
1279 instance of the given attribute on this object.
1280 """
Dale Curtis74a314b2011-06-23 14:55:46 -07001281 raise NotImplementedError
showardf8b19042009-05-12 17:22:49 +00001282
1283
1284 def set_attribute(self, attribute, value):
1285 attribute_model, get_args = self._get_attribute_model_and_args(
1286 attribute)
1287 attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
1288 attribute_object.value = value
1289 attribute_object.save()
1290
1291
1292 def delete_attribute(self, attribute):
1293 attribute_model, get_args = self._get_attribute_model_and_args(
1294 attribute)
1295 try:
1296 attribute_model.objects.get(**get_args).delete()
showard16245422009-09-08 16:28:15 +00001297 except attribute_model.DoesNotExist:
showardf8b19042009-05-12 17:22:49 +00001298 pass
1299
1300
1301 def set_or_delete_attribute(self, attribute, value):
1302 if value is None:
1303 self.delete_attribute(attribute)
1304 else:
1305 self.set_attribute(attribute, value)
showard26b7ec72009-12-21 22:43:57 +00001306
1307
1308class ModelWithHashManager(dbmodels.Manager):
1309 """Manager for use with the ModelWithHash abstract model class"""
1310
1311 def create(self, **kwargs):
1312 raise Exception('ModelWithHash manager should use get_or_create() '
1313 'instead of create()')
1314
1315
1316 def get_or_create(self, **kwargs):
1317 kwargs['the_hash'] = self.model._compute_hash(**kwargs)
1318 return super(ModelWithHashManager, self).get_or_create(**kwargs)
1319
1320
1321class ModelWithHash(dbmodels.Model):
1322 """Superclass with methods for dealing with a hash column"""
1323
1324 the_hash = dbmodels.CharField(max_length=40, unique=True)
1325
1326 objects = ModelWithHashManager()
1327
1328 class Meta:
1329 abstract = True
1330
1331
1332 @classmethod
1333 def _compute_hash(cls, **kwargs):
1334 raise NotImplementedError('Subclasses must override _compute_hash()')
1335
1336
1337 def save(self, force_insert=False, **kwargs):
1338 """Prevents saving the model in most cases
1339
1340 We want these models to be immutable, so the generic save() operation
1341 will not work. These models should be instantiated through their the
1342 model.objects.get_or_create() method instead.
1343
1344 The exception is that save(force_insert=True) will be allowed, since
1345 that creates a new row. However, the preferred way to make instances of
1346 these models is through the get_or_create() method.
1347 """
1348 if not force_insert:
1349 # Allow a forced insert to happen; if it's a duplicate, the unique
1350 # constraint will catch it later anyways
1351 raise Exception('ModelWithHash is immutable')
1352 super(ModelWithHash, self).save(force_insert=force_insert, **kwargs)