blob: 645312811ae7329a7a77255ce7ca7b896f4b4de6 [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
jadmanski0afbb632008-06-06 21:10:57 +0000567 @classmethod
568 def convert_human_readable_values(cls, data, to_human_readable=False):
569 """\
570 Performs conversions on user-supplied field data, to make it
571 easier for users to pass human-readable data.
showard7c785282008-05-29 19:45:12 +0000572
jadmanski0afbb632008-06-06 21:10:57 +0000573 For all fields that have choice sets, convert their values
574 from human-readable strings to enum values, if necessary. This
575 allows users to pass strings instead of the corresponding
576 integer values.
showard7c785282008-05-29 19:45:12 +0000577
jadmanski0afbb632008-06-06 21:10:57 +0000578 For all foreign key fields, call smart_get with the supplied
579 data. This allows the user to pass either an ID value or
580 the name of the object as a string.
showard7c785282008-05-29 19:45:12 +0000581
jadmanski0afbb632008-06-06 21:10:57 +0000582 If to_human_readable=True, perform the inverse - i.e. convert
583 numeric values to human readable values.
showard7c785282008-05-29 19:45:12 +0000584
jadmanski0afbb632008-06-06 21:10:57 +0000585 This method modifies data in-place.
586 """
587 field_dict = cls.get_field_dict()
588 for field_name in data:
showarde732ee72008-09-23 19:15:43 +0000589 if field_name not in field_dict or data[field_name] is None:
jadmanski0afbb632008-06-06 21:10:57 +0000590 continue
591 field_obj = field_dict[field_name]
592 # convert enum values
593 if field_obj.choices:
594 for choice_data in field_obj.choices:
595 # choice_data is (value, name)
596 if to_human_readable:
597 from_val, to_val = choice_data
598 else:
599 to_val, from_val = choice_data
600 if from_val == data[field_name]:
601 data[field_name] = to_val
602 break
603 # convert foreign key values
604 elif field_obj.rel:
showarda4ea5742009-02-17 20:56:23 +0000605 dest_obj = field_obj.rel.to.smart_get(data[field_name],
606 valid_only=False)
showardf8b19042009-05-12 17:22:49 +0000607 if to_human_readable:
Paul Pendlebury5a8c6ad2011-02-01 07:20:17 -0800608 # parameterized_jobs do not have a name_field
609 if (field_name != 'parameterized_job' and
610 dest_obj.name_field is not None):
showardf8b19042009-05-12 17:22:49 +0000611 data[field_name] = getattr(dest_obj,
612 dest_obj.name_field)
jadmanski0afbb632008-06-06 21:10:57 +0000613 else:
showardb0a73032009-03-27 18:35:41 +0000614 data[field_name] = dest_obj
showard7c785282008-05-29 19:45:12 +0000615
616
showard7c785282008-05-29 19:45:12 +0000617
618
Dale Curtis74a314b2011-06-23 14:55:46 -0700619 def _validate_unique(self):
jadmanski0afbb632008-06-06 21:10:57 +0000620 """\
621 Validate that unique fields are unique. Django manipulators do
622 this too, but they're a huge pain to use manually. Trust me.
623 """
624 errors = {}
625 cls = type(self)
626 field_dict = self.get_field_dict()
627 manager = cls.get_valid_manager()
628 for field_name, field_obj in field_dict.iteritems():
629 if not field_obj.unique:
630 continue
showard7c785282008-05-29 19:45:12 +0000631
jadmanski0afbb632008-06-06 21:10:57 +0000632 value = getattr(self, field_name)
showardbd18ab72009-09-18 21:20:27 +0000633 if value is None and field_obj.auto_created:
634 # don't bother checking autoincrement fields about to be
635 # generated
636 continue
637
jadmanski0afbb632008-06-06 21:10:57 +0000638 existing_objs = manager.filter(**{field_name : value})
639 num_existing = existing_objs.count()
showard7c785282008-05-29 19:45:12 +0000640
jadmanski0afbb632008-06-06 21:10:57 +0000641 if num_existing == 0:
642 continue
643 if num_existing == 1 and existing_objs[0].id == self.id:
644 continue
645 errors[field_name] = (
646 'This value must be unique (%s)' % (value))
647 return errors
showard7c785282008-05-29 19:45:12 +0000648
649
showarda5288b42009-07-28 20:06:08 +0000650 def _validate(self):
651 """
652 First coerces all fields on this instance to their proper Python types.
653 Then runs validation on every field. Returns a dictionary of
654 field_name -> error_list.
655
656 Based on validate() from django.db.models.Model in Django 0.96, which
657 was removed in Django 1.0. It should reappear in a later version. See:
658 http://code.djangoproject.com/ticket/6845
659 """
660 error_dict = {}
661 for f in self._meta.fields:
662 try:
663 python_value = f.to_python(
664 getattr(self, f.attname, f.get_default()))
665 except django.core.exceptions.ValidationError, e:
jamesren1e0a4ce2010-04-21 17:45:11 +0000666 error_dict[f.name] = str(e)
showarda5288b42009-07-28 20:06:08 +0000667 continue
668
669 if not f.blank and not python_value:
670 error_dict[f.name] = 'This field is required.'
671 continue
672
673 setattr(self, f.attname, python_value)
674
675 return error_dict
676
677
jadmanski0afbb632008-06-06 21:10:57 +0000678 def do_validate(self):
showarda5288b42009-07-28 20:06:08 +0000679 errors = self._validate()
Dale Curtis74a314b2011-06-23 14:55:46 -0700680 unique_errors = self._validate_unique()
jadmanski0afbb632008-06-06 21:10:57 +0000681 for field_name, error in unique_errors.iteritems():
682 errors.setdefault(field_name, error)
683 if errors:
684 raise ValidationError(errors)
showard7c785282008-05-29 19:45:12 +0000685
686
jadmanski0afbb632008-06-06 21:10:57 +0000687 # actually (externally) useful methods follow
showard7c785282008-05-29 19:45:12 +0000688
jadmanski0afbb632008-06-06 21:10:57 +0000689 @classmethod
690 def add_object(cls, data={}, **kwargs):
691 """\
692 Returns a new object created with the given data (a dictionary
693 mapping field names to values). Merges any extra keyword args
694 into data.
695 """
Prashanth B489b91d2014-03-15 12:17:16 -0700696 data = dict(data)
697 data.update(kwargs)
698 data = cls.prepare_data_args(data)
699 cls.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000700 data = cls.provide_default_values(data)
Prashanth B489b91d2014-03-15 12:17:16 -0700701
jadmanski0afbb632008-06-06 21:10:57 +0000702 obj = cls(**data)
703 obj.do_validate()
704 obj.save()
705 return obj
showard7c785282008-05-29 19:45:12 +0000706
Michael Liang8864e862014-07-22 08:36:05 -0700707 def record_state(self, type_str, state, value):
708 """Record metadata in elasticsearch.
709
710 @param type_str: sets the _type field in elasticsearch db.
711 @param state: string representing what state we are recording,
712 e.g. 'locked'
713 @param value: value of the state, e.g. True
714 """
715 metadata = {
716 'time_recorded': time.time(),
717 state: value,
718 'hostname': self.hostname,
719 }
720 es_utils.ESMetadata().post(type_str=type_str, metadata=metadata)
721
showard7c785282008-05-29 19:45:12 +0000722
jadmanski0afbb632008-06-06 21:10:57 +0000723 def update_object(self, data={}, **kwargs):
724 """\
725 Updates the object with the given data (a dictionary mapping
726 field names to values). Merges any extra keyword args into
727 data.
728 """
Prashanth B489b91d2014-03-15 12:17:16 -0700729 data = dict(data)
730 data.update(kwargs)
731 data = self.prepare_data_args(data)
732 self.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000733 for field_name, value in data.iteritems():
showardb0a73032009-03-27 18:35:41 +0000734 setattr(self, field_name, value)
Michael Liang8864e862014-07-22 08:36:05 -0700735 # Other fields such as label (when updated) are not sent over
736 # the es because it doesn't contribute to putting together host
737 # host history. Locks are important in host history because if
738 # a device is locked then we don't really care what state it is in.
739 if field_name == 'locked':
740 self.record_state('lock_history', 'locked', value)
jadmanski0afbb632008-06-06 21:10:57 +0000741 self.do_validate()
742 self.save()
showard7c785282008-05-29 19:45:12 +0000743
744
showard8bfb5cb2009-10-07 20:49:15 +0000745 # see query_objects()
746 _SPECIAL_FILTER_KEYS = ('query_start', 'query_limit', 'sort_by',
747 'extra_args', 'extra_where', 'no_distinct')
748
749
jadmanski0afbb632008-06-06 21:10:57 +0000750 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000751 def _extract_special_params(cls, filter_data):
752 """
753 @returns a tuple of dicts (special_params, regular_filters), where
754 special_params contains the parameters we handle specially and
755 regular_filters is the remaining data to be handled by Django.
756 """
757 regular_filters = dict(filter_data)
758 special_params = {}
759 for key in cls._SPECIAL_FILTER_KEYS:
760 if key in regular_filters:
761 special_params[key] = regular_filters.pop(key)
762 return special_params, regular_filters
763
764
765 @classmethod
766 def apply_presentation(cls, query, filter_data):
767 """
768 Apply presentation parameters -- sorting and paging -- to the given
769 query.
770 @returns new query with presentation applied
771 """
772 special_params, _ = cls._extract_special_params(filter_data)
773 sort_by = special_params.get('sort_by', None)
774 if sort_by:
775 assert isinstance(sort_by, list) or isinstance(sort_by, tuple)
showard8b0ea222009-12-23 19:23:03 +0000776 query = query.extra(order_by=sort_by)
showard8bfb5cb2009-10-07 20:49:15 +0000777
778 query_start = special_params.get('query_start', None)
779 query_limit = special_params.get('query_limit', None)
780 if query_start is not None:
781 if query_limit is None:
782 raise ValueError('Cannot pass query_start without query_limit')
783 # query_limit is passed as a page size
showard7074b742009-10-12 20:30:04 +0000784 query_limit += query_start
785 return query[query_start:query_limit]
showard8bfb5cb2009-10-07 20:49:15 +0000786
787
788 @classmethod
789 def query_objects(cls, filter_data, valid_only=True, initial_query=None,
790 apply_presentation=True):
jadmanski0afbb632008-06-06 21:10:57 +0000791 """\
792 Returns a QuerySet object for querying the given model_class
793 with the given filter_data. Optional special arguments in
794 filter_data include:
795 -query_start: index of first return to return
796 -query_limit: maximum number of results to return
797 -sort_by: list of fields to sort on. prefixing a '-' onto a
798 field name changes the sort to descending order.
799 -extra_args: keyword args to pass to query.extra() (see Django
800 DB layer documentation)
showarda5288b42009-07-28 20:06:08 +0000801 -extra_where: extra WHERE clause to append
showard8bfb5cb2009-10-07 20:49:15 +0000802 -no_distinct: if True, a DISTINCT will not be added to the SELECT
jadmanski0afbb632008-06-06 21:10:57 +0000803 """
showard8bfb5cb2009-10-07 20:49:15 +0000804 special_params, regular_filters = cls._extract_special_params(
805 filter_data)
showard7c785282008-05-29 19:45:12 +0000806
showard7ac7b7a2008-07-21 20:24:29 +0000807 if initial_query is None:
808 if valid_only:
809 initial_query = cls.get_valid_manager()
810 else:
811 initial_query = cls.objects
showard8bfb5cb2009-10-07 20:49:15 +0000812
813 query = initial_query.filter(**regular_filters)
814
815 use_distinct = not special_params.get('no_distinct', False)
showard7ac7b7a2008-07-21 20:24:29 +0000816 if use_distinct:
817 query = query.distinct()
showard7c785282008-05-29 19:45:12 +0000818
showard8bfb5cb2009-10-07 20:49:15 +0000819 extra_args = special_params.get('extra_args', {})
820 extra_where = special_params.get('extra_where', None)
821 if extra_where:
822 # escape %'s
823 extra_where = cls.objects.escape_user_sql(extra_where)
824 extra_args.setdefault('where', []).append(extra_where)
jadmanski0afbb632008-06-06 21:10:57 +0000825 if extra_args:
826 query = query.extra(**extra_args)
showard09096d82008-07-07 23:20:49 +0000827 query = query._clone(klass=ReadonlyQuerySet)
showard7c785282008-05-29 19:45:12 +0000828
showard8bfb5cb2009-10-07 20:49:15 +0000829 if apply_presentation:
830 query = cls.apply_presentation(query, filter_data)
831
832 return query
showard7c785282008-05-29 19:45:12 +0000833
834
jadmanski0afbb632008-06-06 21:10:57 +0000835 @classmethod
showard585c2ab2008-07-23 19:29:49 +0000836 def query_count(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000837 """\
838 Like query_objects, but retreive only the count of results.
839 """
840 filter_data.pop('query_start', None)
841 filter_data.pop('query_limit', None)
showard585c2ab2008-07-23 19:29:49 +0000842 query = cls.query_objects(filter_data, initial_query=initial_query)
843 return query.count()
showard7c785282008-05-29 19:45:12 +0000844
845
jadmanski0afbb632008-06-06 21:10:57 +0000846 @classmethod
847 def clean_object_dicts(cls, field_dicts):
848 """\
849 Take a list of dicts corresponding to object (as returned by
850 query.values()) and clean the data to be more suitable for
851 returning to the user.
852 """
showarde732ee72008-09-23 19:15:43 +0000853 for field_dict in field_dicts:
854 cls.clean_foreign_keys(field_dict)
showard21baa452008-10-21 00:08:39 +0000855 cls._convert_booleans(field_dict)
showarde732ee72008-09-23 19:15:43 +0000856 cls.convert_human_readable_values(field_dict,
857 to_human_readable=True)
showard7c785282008-05-29 19:45:12 +0000858
859
jadmanski0afbb632008-06-06 21:10:57 +0000860 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000861 def list_objects(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000862 """\
863 Like query_objects, but return a list of dictionaries.
864 """
showard7ac7b7a2008-07-21 20:24:29 +0000865 query = cls.query_objects(filter_data, initial_query=initial_query)
showard8bfb5cb2009-10-07 20:49:15 +0000866 extra_fields = query.query.extra_select.keys()
867 field_dicts = [model_object.get_object_dict(extra_fields=extra_fields)
showarde732ee72008-09-23 19:15:43 +0000868 for model_object in query]
jadmanski0afbb632008-06-06 21:10:57 +0000869 return field_dicts
showard7c785282008-05-29 19:45:12 +0000870
871
jadmanski0afbb632008-06-06 21:10:57 +0000872 @classmethod
showarda4ea5742009-02-17 20:56:23 +0000873 def smart_get(cls, id_or_name, valid_only=True):
jadmanski0afbb632008-06-06 21:10:57 +0000874 """\
875 smart_get(integer) -> get object by ID
876 smart_get(string) -> get object by name_field
jadmanski0afbb632008-06-06 21:10:57 +0000877 """
showarda4ea5742009-02-17 20:56:23 +0000878 if valid_only:
879 manager = cls.get_valid_manager()
880 else:
881 manager = cls.objects
882
883 if isinstance(id_or_name, (int, long)):
884 return manager.get(pk=id_or_name)
jamesren3e9f6092010-03-11 21:32:10 +0000885 if isinstance(id_or_name, basestring) and hasattr(cls, 'name_field'):
showarda4ea5742009-02-17 20:56:23 +0000886 return manager.get(**{cls.name_field : id_or_name})
887 raise ValueError(
888 'Invalid positional argument: %s (%s)' % (id_or_name,
889 type(id_or_name)))
showard7c785282008-05-29 19:45:12 +0000890
891
showardbe3ec042008-11-12 18:16:07 +0000892 @classmethod
893 def smart_get_bulk(cls, id_or_name_list):
894 invalid_inputs = []
895 result_objects = []
896 for id_or_name in id_or_name_list:
897 try:
898 result_objects.append(cls.smart_get(id_or_name))
899 except cls.DoesNotExist:
900 invalid_inputs.append(id_or_name)
901 if invalid_inputs:
mbligh7a3ebe32008-12-01 17:10:33 +0000902 raise cls.DoesNotExist('The following %ss do not exist: %s'
903 % (cls.__name__.lower(),
904 ', '.join(invalid_inputs)))
showardbe3ec042008-11-12 18:16:07 +0000905 return result_objects
906
907
showard8bfb5cb2009-10-07 20:49:15 +0000908 def get_object_dict(self, extra_fields=None):
jadmanski0afbb632008-06-06 21:10:57 +0000909 """\
showard8bfb5cb2009-10-07 20:49:15 +0000910 Return a dictionary mapping fields to this object's values. @param
911 extra_fields: list of extra attribute names to include, in addition to
912 the fields defined on this object.
jadmanski0afbb632008-06-06 21:10:57 +0000913 """
showard8bfb5cb2009-10-07 20:49:15 +0000914 fields = self.get_field_dict().keys()
915 if extra_fields:
916 fields += extra_fields
jadmanski0afbb632008-06-06 21:10:57 +0000917 object_dict = dict((field_name, getattr(self, field_name))
showarde732ee72008-09-23 19:15:43 +0000918 for field_name in fields)
jadmanski0afbb632008-06-06 21:10:57 +0000919 self.clean_object_dicts([object_dict])
showardd3dc1992009-04-22 21:01:40 +0000920 self._postprocess_object_dict(object_dict)
jadmanski0afbb632008-06-06 21:10:57 +0000921 return object_dict
showard7c785282008-05-29 19:45:12 +0000922
923
showardd3dc1992009-04-22 21:01:40 +0000924 def _postprocess_object_dict(self, object_dict):
925 """For subclasses to override."""
926 pass
927
928
jadmanski0afbb632008-06-06 21:10:57 +0000929 @classmethod
930 def get_valid_manager(cls):
931 return cls.objects
showard7c785282008-05-29 19:45:12 +0000932
933
showard2bab8f42008-11-12 18:15:22 +0000934 def _record_attributes(self, attributes):
935 """
936 See on_attribute_changed.
937 """
938 assert not isinstance(attributes, basestring)
939 self._recorded_attributes = dict((attribute, getattr(self, attribute))
940 for attribute in attributes)
941
942
943 def _check_for_updated_attributes(self):
944 """
945 See on_attribute_changed.
946 """
947 for attribute, original_value in self._recorded_attributes.iteritems():
948 new_value = getattr(self, attribute)
949 if original_value != new_value:
950 self.on_attribute_changed(attribute, original_value)
951 self._record_attributes(self._recorded_attributes.keys())
952
953
954 def on_attribute_changed(self, attribute, old_value):
955 """
956 Called whenever an attribute is updated. To be overridden.
957
958 To use this method, you must:
959 * call _record_attributes() from __init__() (after making the super
960 call) with a list of attributes for which you want to be notified upon
961 change.
962 * call _check_for_updated_attributes() from save().
963 """
964 pass
965
966
showard7c785282008-05-29 19:45:12 +0000967class ModelWithInvalid(ModelExtensions):
jadmanski0afbb632008-06-06 21:10:57 +0000968 """
969 Overrides model methods save() and delete() to support invalidation in
970 place of actual deletion. Subclasses must have a boolean "invalid"
971 field.
972 """
showard7c785282008-05-29 19:45:12 +0000973
showarda5288b42009-07-28 20:06:08 +0000974 def save(self, *args, **kwargs):
showardddb90992009-02-11 23:39:32 +0000975 first_time = (self.id is None)
976 if first_time:
977 # see if this object was previously added and invalidated
978 my_name = getattr(self, self.name_field)
979 filters = {self.name_field : my_name, 'invalid' : True}
980 try:
981 old_object = self.__class__.objects.get(**filters)
showardafd97de2009-10-01 18:45:09 +0000982 self.resurrect_object(old_object)
showardddb90992009-02-11 23:39:32 +0000983 except self.DoesNotExist:
984 # no existing object
985 pass
showard7c785282008-05-29 19:45:12 +0000986
showarda5288b42009-07-28 20:06:08 +0000987 super(ModelWithInvalid, self).save(*args, **kwargs)
showard7c785282008-05-29 19:45:12 +0000988
989
showardafd97de2009-10-01 18:45:09 +0000990 def resurrect_object(self, old_object):
991 """
992 Called when self is about to be saved for the first time and is actually
993 "undeleting" a previously deleted object. Can be overridden by
994 subclasses to copy data as desired from the deleted entry (but this
995 superclass implementation must normally be called).
996 """
997 self.id = old_object.id
998
999
jadmanski0afbb632008-06-06 21:10:57 +00001000 def clean_object(self):
1001 """
1002 This method is called when an object is marked invalid.
1003 Subclasses should override this to clean up relationships that
showardafd97de2009-10-01 18:45:09 +00001004 should no longer exist if the object were deleted.
1005 """
jadmanski0afbb632008-06-06 21:10:57 +00001006 pass
showard7c785282008-05-29 19:45:12 +00001007
1008
jadmanski0afbb632008-06-06 21:10:57 +00001009 def delete(self):
Dale Curtis74a314b2011-06-23 14:55:46 -07001010 self.invalid = self.invalid
jadmanski0afbb632008-06-06 21:10:57 +00001011 assert not self.invalid
1012 self.invalid = True
1013 self.save()
1014 self.clean_object()
showard7c785282008-05-29 19:45:12 +00001015
1016
jadmanski0afbb632008-06-06 21:10:57 +00001017 @classmethod
1018 def get_valid_manager(cls):
1019 return cls.valid_objects
showard7c785282008-05-29 19:45:12 +00001020
1021
jadmanski0afbb632008-06-06 21:10:57 +00001022 class Manipulator(object):
1023 """
1024 Force default manipulators to look only at valid objects -
1025 otherwise they will match against invalid objects when checking
1026 uniqueness.
1027 """
1028 @classmethod
1029 def _prepare(cls, model):
1030 super(ModelWithInvalid.Manipulator, cls)._prepare(model)
1031 cls.manager = model.valid_objects
showardf8b19042009-05-12 17:22:49 +00001032
1033
1034class ModelWithAttributes(object):
1035 """
1036 Mixin class for models that have an attribute model associated with them.
1037 The attribute model is assumed to have its value field named "value".
1038 """
1039
1040 def _get_attribute_model_and_args(self, attribute):
1041 """
1042 Subclasses should override this to return a tuple (attribute_model,
1043 keyword_args), where attribute_model is a model class and keyword_args
1044 is a dict of args to pass to attribute_model.objects.get() to get an
1045 instance of the given attribute on this object.
1046 """
Dale Curtis74a314b2011-06-23 14:55:46 -07001047 raise NotImplementedError
showardf8b19042009-05-12 17:22:49 +00001048
1049
1050 def set_attribute(self, attribute, value):
1051 attribute_model, get_args = self._get_attribute_model_and_args(
1052 attribute)
1053 attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
1054 attribute_object.value = value
1055 attribute_object.save()
1056
1057
1058 def delete_attribute(self, attribute):
1059 attribute_model, get_args = self._get_attribute_model_and_args(
1060 attribute)
1061 try:
1062 attribute_model.objects.get(**get_args).delete()
showard16245422009-09-08 16:28:15 +00001063 except attribute_model.DoesNotExist:
showardf8b19042009-05-12 17:22:49 +00001064 pass
1065
1066
1067 def set_or_delete_attribute(self, attribute, value):
1068 if value is None:
1069 self.delete_attribute(attribute)
1070 else:
1071 self.set_attribute(attribute, value)
showard26b7ec72009-12-21 22:43:57 +00001072
1073
1074class ModelWithHashManager(dbmodels.Manager):
1075 """Manager for use with the ModelWithHash abstract model class"""
1076
1077 def create(self, **kwargs):
1078 raise Exception('ModelWithHash manager should use get_or_create() '
1079 'instead of create()')
1080
1081
1082 def get_or_create(self, **kwargs):
1083 kwargs['the_hash'] = self.model._compute_hash(**kwargs)
1084 return super(ModelWithHashManager, self).get_or_create(**kwargs)
1085
1086
1087class ModelWithHash(dbmodels.Model):
1088 """Superclass with methods for dealing with a hash column"""
1089
1090 the_hash = dbmodels.CharField(max_length=40, unique=True)
1091
1092 objects = ModelWithHashManager()
1093
1094 class Meta:
1095 abstract = True
1096
1097
1098 @classmethod
1099 def _compute_hash(cls, **kwargs):
1100 raise NotImplementedError('Subclasses must override _compute_hash()')
1101
1102
1103 def save(self, force_insert=False, **kwargs):
1104 """Prevents saving the model in most cases
1105
1106 We want these models to be immutable, so the generic save() operation
1107 will not work. These models should be instantiated through their the
1108 model.objects.get_or_create() method instead.
1109
1110 The exception is that save(force_insert=True) will be allowed, since
1111 that creates a new row. However, the preferred way to make instances of
1112 these models is through the get_or_create() method.
1113 """
1114 if not force_insert:
1115 # Allow a forced insert to happen; if it's a duplicate, the unique
1116 # constraint will catch it later anyways
1117 raise Exception('ModelWithHash is immutable')
1118 super(ModelWithHash, self).save(force_insert=force_insert, **kwargs)