blob: ed1b2e67e70a29766303239892807b8bebdeb51a [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 = {
Dan Shie4cb9e22014-08-29 11:56:04 -0700716 state: value,
Michael Liang8864e862014-07-22 08:36:05 -0700717 'hostname': self.hostname,
718 }
719 es_utils.ESMetadata().post(type_str=type_str, metadata=metadata)
720
showard7c785282008-05-29 19:45:12 +0000721
jadmanski0afbb632008-06-06 21:10:57 +0000722 def update_object(self, data={}, **kwargs):
723 """\
724 Updates the object with the given data (a dictionary mapping
725 field names to values). Merges any extra keyword args into
726 data.
727 """
Prashanth B489b91d2014-03-15 12:17:16 -0700728 data = dict(data)
729 data.update(kwargs)
730 data = self.prepare_data_args(data)
731 self.convert_human_readable_values(data)
jadmanski0afbb632008-06-06 21:10:57 +0000732 for field_name, value in data.iteritems():
showardb0a73032009-03-27 18:35:41 +0000733 setattr(self, field_name, value)
Michael Liang8864e862014-07-22 08:36:05 -0700734 # Other fields such as label (when updated) are not sent over
735 # the es because it doesn't contribute to putting together host
736 # host history. Locks are important in host history because if
737 # a device is locked then we don't really care what state it is in.
738 if field_name == 'locked':
739 self.record_state('lock_history', 'locked', value)
jadmanski0afbb632008-06-06 21:10:57 +0000740 self.do_validate()
741 self.save()
showard7c785282008-05-29 19:45:12 +0000742
743
showard8bfb5cb2009-10-07 20:49:15 +0000744 # see query_objects()
745 _SPECIAL_FILTER_KEYS = ('query_start', 'query_limit', 'sort_by',
746 'extra_args', 'extra_where', 'no_distinct')
747
748
jadmanski0afbb632008-06-06 21:10:57 +0000749 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000750 def _extract_special_params(cls, filter_data):
751 """
752 @returns a tuple of dicts (special_params, regular_filters), where
753 special_params contains the parameters we handle specially and
754 regular_filters is the remaining data to be handled by Django.
755 """
756 regular_filters = dict(filter_data)
757 special_params = {}
758 for key in cls._SPECIAL_FILTER_KEYS:
759 if key in regular_filters:
760 special_params[key] = regular_filters.pop(key)
761 return special_params, regular_filters
762
763
764 @classmethod
765 def apply_presentation(cls, query, filter_data):
766 """
767 Apply presentation parameters -- sorting and paging -- to the given
768 query.
769 @returns new query with presentation applied
770 """
771 special_params, _ = cls._extract_special_params(filter_data)
772 sort_by = special_params.get('sort_by', None)
773 if sort_by:
774 assert isinstance(sort_by, list) or isinstance(sort_by, tuple)
showard8b0ea222009-12-23 19:23:03 +0000775 query = query.extra(order_by=sort_by)
showard8bfb5cb2009-10-07 20:49:15 +0000776
777 query_start = special_params.get('query_start', None)
778 query_limit = special_params.get('query_limit', None)
779 if query_start is not None:
780 if query_limit is None:
781 raise ValueError('Cannot pass query_start without query_limit')
782 # query_limit is passed as a page size
showard7074b742009-10-12 20:30:04 +0000783 query_limit += query_start
784 return query[query_start:query_limit]
showard8bfb5cb2009-10-07 20:49:15 +0000785
786
787 @classmethod
788 def query_objects(cls, filter_data, valid_only=True, initial_query=None,
789 apply_presentation=True):
jadmanski0afbb632008-06-06 21:10:57 +0000790 """\
791 Returns a QuerySet object for querying the given model_class
792 with the given filter_data. Optional special arguments in
793 filter_data include:
794 -query_start: index of first return to return
795 -query_limit: maximum number of results to return
796 -sort_by: list of fields to sort on. prefixing a '-' onto a
797 field name changes the sort to descending order.
798 -extra_args: keyword args to pass to query.extra() (see Django
799 DB layer documentation)
showarda5288b42009-07-28 20:06:08 +0000800 -extra_where: extra WHERE clause to append
showard8bfb5cb2009-10-07 20:49:15 +0000801 -no_distinct: if True, a DISTINCT will not be added to the SELECT
jadmanski0afbb632008-06-06 21:10:57 +0000802 """
showard8bfb5cb2009-10-07 20:49:15 +0000803 special_params, regular_filters = cls._extract_special_params(
804 filter_data)
showard7c785282008-05-29 19:45:12 +0000805
showard7ac7b7a2008-07-21 20:24:29 +0000806 if initial_query is None:
807 if valid_only:
808 initial_query = cls.get_valid_manager()
809 else:
810 initial_query = cls.objects
showard8bfb5cb2009-10-07 20:49:15 +0000811
812 query = initial_query.filter(**regular_filters)
813
814 use_distinct = not special_params.get('no_distinct', False)
showard7ac7b7a2008-07-21 20:24:29 +0000815 if use_distinct:
816 query = query.distinct()
showard7c785282008-05-29 19:45:12 +0000817
showard8bfb5cb2009-10-07 20:49:15 +0000818 extra_args = special_params.get('extra_args', {})
819 extra_where = special_params.get('extra_where', None)
820 if extra_where:
821 # escape %'s
822 extra_where = cls.objects.escape_user_sql(extra_where)
823 extra_args.setdefault('where', []).append(extra_where)
jadmanski0afbb632008-06-06 21:10:57 +0000824 if extra_args:
825 query = query.extra(**extra_args)
showard09096d82008-07-07 23:20:49 +0000826 query = query._clone(klass=ReadonlyQuerySet)
showard7c785282008-05-29 19:45:12 +0000827
showard8bfb5cb2009-10-07 20:49:15 +0000828 if apply_presentation:
829 query = cls.apply_presentation(query, filter_data)
830
831 return query
showard7c785282008-05-29 19:45:12 +0000832
833
jadmanski0afbb632008-06-06 21:10:57 +0000834 @classmethod
showard585c2ab2008-07-23 19:29:49 +0000835 def query_count(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000836 """\
837 Like query_objects, but retreive only the count of results.
838 """
839 filter_data.pop('query_start', None)
840 filter_data.pop('query_limit', None)
showard585c2ab2008-07-23 19:29:49 +0000841 query = cls.query_objects(filter_data, initial_query=initial_query)
842 return query.count()
showard7c785282008-05-29 19:45:12 +0000843
844
jadmanski0afbb632008-06-06 21:10:57 +0000845 @classmethod
846 def clean_object_dicts(cls, field_dicts):
847 """\
848 Take a list of dicts corresponding to object (as returned by
849 query.values()) and clean the data to be more suitable for
850 returning to the user.
851 """
showarde732ee72008-09-23 19:15:43 +0000852 for field_dict in field_dicts:
853 cls.clean_foreign_keys(field_dict)
showard21baa452008-10-21 00:08:39 +0000854 cls._convert_booleans(field_dict)
showarde732ee72008-09-23 19:15:43 +0000855 cls.convert_human_readable_values(field_dict,
856 to_human_readable=True)
showard7c785282008-05-29 19:45:12 +0000857
858
jadmanski0afbb632008-06-06 21:10:57 +0000859 @classmethod
showard8bfb5cb2009-10-07 20:49:15 +0000860 def list_objects(cls, filter_data, initial_query=None):
jadmanski0afbb632008-06-06 21:10:57 +0000861 """\
862 Like query_objects, but return a list of dictionaries.
863 """
showard7ac7b7a2008-07-21 20:24:29 +0000864 query = cls.query_objects(filter_data, initial_query=initial_query)
showard8bfb5cb2009-10-07 20:49:15 +0000865 extra_fields = query.query.extra_select.keys()
866 field_dicts = [model_object.get_object_dict(extra_fields=extra_fields)
showarde732ee72008-09-23 19:15:43 +0000867 for model_object in query]
jadmanski0afbb632008-06-06 21:10:57 +0000868 return field_dicts
showard7c785282008-05-29 19:45:12 +0000869
870
jadmanski0afbb632008-06-06 21:10:57 +0000871 @classmethod
showarda4ea5742009-02-17 20:56:23 +0000872 def smart_get(cls, id_or_name, valid_only=True):
jadmanski0afbb632008-06-06 21:10:57 +0000873 """\
874 smart_get(integer) -> get object by ID
875 smart_get(string) -> get object by name_field
jadmanski0afbb632008-06-06 21:10:57 +0000876 """
showarda4ea5742009-02-17 20:56:23 +0000877 if valid_only:
878 manager = cls.get_valid_manager()
879 else:
880 manager = cls.objects
881
882 if isinstance(id_or_name, (int, long)):
883 return manager.get(pk=id_or_name)
jamesren3e9f6092010-03-11 21:32:10 +0000884 if isinstance(id_or_name, basestring) and hasattr(cls, 'name_field'):
showarda4ea5742009-02-17 20:56:23 +0000885 return manager.get(**{cls.name_field : id_or_name})
886 raise ValueError(
887 'Invalid positional argument: %s (%s)' % (id_or_name,
888 type(id_or_name)))
showard7c785282008-05-29 19:45:12 +0000889
890
showardbe3ec042008-11-12 18:16:07 +0000891 @classmethod
892 def smart_get_bulk(cls, id_or_name_list):
893 invalid_inputs = []
894 result_objects = []
895 for id_or_name in id_or_name_list:
896 try:
897 result_objects.append(cls.smart_get(id_or_name))
898 except cls.DoesNotExist:
899 invalid_inputs.append(id_or_name)
900 if invalid_inputs:
mbligh7a3ebe32008-12-01 17:10:33 +0000901 raise cls.DoesNotExist('The following %ss do not exist: %s'
902 % (cls.__name__.lower(),
903 ', '.join(invalid_inputs)))
showardbe3ec042008-11-12 18:16:07 +0000904 return result_objects
905
906
showard8bfb5cb2009-10-07 20:49:15 +0000907 def get_object_dict(self, extra_fields=None):
jadmanski0afbb632008-06-06 21:10:57 +0000908 """\
showard8bfb5cb2009-10-07 20:49:15 +0000909 Return a dictionary mapping fields to this object's values. @param
910 extra_fields: list of extra attribute names to include, in addition to
911 the fields defined on this object.
jadmanski0afbb632008-06-06 21:10:57 +0000912 """
showard8bfb5cb2009-10-07 20:49:15 +0000913 fields = self.get_field_dict().keys()
914 if extra_fields:
915 fields += extra_fields
jadmanski0afbb632008-06-06 21:10:57 +0000916 object_dict = dict((field_name, getattr(self, field_name))
showarde732ee72008-09-23 19:15:43 +0000917 for field_name in fields)
jadmanski0afbb632008-06-06 21:10:57 +0000918 self.clean_object_dicts([object_dict])
showardd3dc1992009-04-22 21:01:40 +0000919 self._postprocess_object_dict(object_dict)
jadmanski0afbb632008-06-06 21:10:57 +0000920 return object_dict
showard7c785282008-05-29 19:45:12 +0000921
922
showardd3dc1992009-04-22 21:01:40 +0000923 def _postprocess_object_dict(self, object_dict):
924 """For subclasses to override."""
925 pass
926
927
jadmanski0afbb632008-06-06 21:10:57 +0000928 @classmethod
929 def get_valid_manager(cls):
930 return cls.objects
showard7c785282008-05-29 19:45:12 +0000931
932
showard2bab8f42008-11-12 18:15:22 +0000933 def _record_attributes(self, attributes):
934 """
935 See on_attribute_changed.
936 """
937 assert not isinstance(attributes, basestring)
938 self._recorded_attributes = dict((attribute, getattr(self, attribute))
939 for attribute in attributes)
940
941
942 def _check_for_updated_attributes(self):
943 """
944 See on_attribute_changed.
945 """
946 for attribute, original_value in self._recorded_attributes.iteritems():
947 new_value = getattr(self, attribute)
948 if original_value != new_value:
949 self.on_attribute_changed(attribute, original_value)
950 self._record_attributes(self._recorded_attributes.keys())
951
952
953 def on_attribute_changed(self, attribute, old_value):
954 """
955 Called whenever an attribute is updated. To be overridden.
956
957 To use this method, you must:
958 * call _record_attributes() from __init__() (after making the super
959 call) with a list of attributes for which you want to be notified upon
960 change.
961 * call _check_for_updated_attributes() from save().
962 """
963 pass
964
965
showard7c785282008-05-29 19:45:12 +0000966class ModelWithInvalid(ModelExtensions):
jadmanski0afbb632008-06-06 21:10:57 +0000967 """
968 Overrides model methods save() and delete() to support invalidation in
969 place of actual deletion. Subclasses must have a boolean "invalid"
970 field.
971 """
showard7c785282008-05-29 19:45:12 +0000972
showarda5288b42009-07-28 20:06:08 +0000973 def save(self, *args, **kwargs):
showardddb90992009-02-11 23:39:32 +0000974 first_time = (self.id is None)
975 if first_time:
976 # see if this object was previously added and invalidated
977 my_name = getattr(self, self.name_field)
978 filters = {self.name_field : my_name, 'invalid' : True}
979 try:
980 old_object = self.__class__.objects.get(**filters)
showardafd97de2009-10-01 18:45:09 +0000981 self.resurrect_object(old_object)
showardddb90992009-02-11 23:39:32 +0000982 except self.DoesNotExist:
983 # no existing object
984 pass
showard7c785282008-05-29 19:45:12 +0000985
showarda5288b42009-07-28 20:06:08 +0000986 super(ModelWithInvalid, self).save(*args, **kwargs)
showard7c785282008-05-29 19:45:12 +0000987
988
showardafd97de2009-10-01 18:45:09 +0000989 def resurrect_object(self, old_object):
990 """
991 Called when self is about to be saved for the first time and is actually
992 "undeleting" a previously deleted object. Can be overridden by
993 subclasses to copy data as desired from the deleted entry (but this
994 superclass implementation must normally be called).
995 """
996 self.id = old_object.id
997
998
jadmanski0afbb632008-06-06 21:10:57 +0000999 def clean_object(self):
1000 """
1001 This method is called when an object is marked invalid.
1002 Subclasses should override this to clean up relationships that
showardafd97de2009-10-01 18:45:09 +00001003 should no longer exist if the object were deleted.
1004 """
jadmanski0afbb632008-06-06 21:10:57 +00001005 pass
showard7c785282008-05-29 19:45:12 +00001006
1007
jadmanski0afbb632008-06-06 21:10:57 +00001008 def delete(self):
Dale Curtis74a314b2011-06-23 14:55:46 -07001009 self.invalid = self.invalid
jadmanski0afbb632008-06-06 21:10:57 +00001010 assert not self.invalid
1011 self.invalid = True
1012 self.save()
1013 self.clean_object()
showard7c785282008-05-29 19:45:12 +00001014
1015
jadmanski0afbb632008-06-06 21:10:57 +00001016 @classmethod
1017 def get_valid_manager(cls):
1018 return cls.valid_objects
showard7c785282008-05-29 19:45:12 +00001019
1020
jadmanski0afbb632008-06-06 21:10:57 +00001021 class Manipulator(object):
1022 """
1023 Force default manipulators to look only at valid objects -
1024 otherwise they will match against invalid objects when checking
1025 uniqueness.
1026 """
1027 @classmethod
1028 def _prepare(cls, model):
1029 super(ModelWithInvalid.Manipulator, cls)._prepare(model)
1030 cls.manager = model.valid_objects
showardf8b19042009-05-12 17:22:49 +00001031
1032
1033class ModelWithAttributes(object):
1034 """
1035 Mixin class for models that have an attribute model associated with them.
1036 The attribute model is assumed to have its value field named "value".
1037 """
1038
1039 def _get_attribute_model_and_args(self, attribute):
1040 """
1041 Subclasses should override this to return a tuple (attribute_model,
1042 keyword_args), where attribute_model is a model class and keyword_args
1043 is a dict of args to pass to attribute_model.objects.get() to get an
1044 instance of the given attribute on this object.
1045 """
Dale Curtis74a314b2011-06-23 14:55:46 -07001046 raise NotImplementedError
showardf8b19042009-05-12 17:22:49 +00001047
1048
1049 def set_attribute(self, attribute, value):
1050 attribute_model, get_args = self._get_attribute_model_and_args(
1051 attribute)
1052 attribute_object, _ = attribute_model.objects.get_or_create(**get_args)
1053 attribute_object.value = value
1054 attribute_object.save()
1055
1056
1057 def delete_attribute(self, attribute):
1058 attribute_model, get_args = self._get_attribute_model_and_args(
1059 attribute)
1060 try:
1061 attribute_model.objects.get(**get_args).delete()
showard16245422009-09-08 16:28:15 +00001062 except attribute_model.DoesNotExist:
showardf8b19042009-05-12 17:22:49 +00001063 pass
1064
1065
1066 def set_or_delete_attribute(self, attribute, value):
1067 if value is None:
1068 self.delete_attribute(attribute)
1069 else:
1070 self.set_attribute(attribute, value)
showard26b7ec72009-12-21 22:43:57 +00001071
1072
1073class ModelWithHashManager(dbmodels.Manager):
1074 """Manager for use with the ModelWithHash abstract model class"""
1075
1076 def create(self, **kwargs):
1077 raise Exception('ModelWithHash manager should use get_or_create() '
1078 'instead of create()')
1079
1080
1081 def get_or_create(self, **kwargs):
1082 kwargs['the_hash'] = self.model._compute_hash(**kwargs)
1083 return super(ModelWithHashManager, self).get_or_create(**kwargs)
1084
1085
1086class ModelWithHash(dbmodels.Model):
1087 """Superclass with methods for dealing with a hash column"""
1088
1089 the_hash = dbmodels.CharField(max_length=40, unique=True)
1090
1091 objects = ModelWithHashManager()
1092
1093 class Meta:
1094 abstract = True
1095
1096
1097 @classmethod
1098 def _compute_hash(cls, **kwargs):
1099 raise NotImplementedError('Subclasses must override _compute_hash()')
1100
1101
1102 def save(self, force_insert=False, **kwargs):
1103 """Prevents saving the model in most cases
1104
1105 We want these models to be immutable, so the generic save() operation
1106 will not work. These models should be instantiated through their the
1107 model.objects.get_or_create() method instead.
1108
1109 The exception is that save(force_insert=True) will be allowed, since
1110 that creates a new row. However, the preferred way to make instances of
1111 these models is through the get_or_create() method.
1112 """
1113 if not force_insert:
1114 # Allow a forced insert to happen; if it's a duplicate, the unique
1115 # constraint will catch it later anyways
1116 raise Exception('ModelWithHash is immutable')
1117 super(ModelWithHash, self).save(force_insert=force_insert, **kwargs)