constraints.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195
  1. from enum import Enum
  2. from django.db.models.query_utils import Q
  3. from django.db.models.sql.query import Query
  4. __all__ = ['CheckConstraint', 'Deferrable', 'UniqueConstraint']
  5. class BaseConstraint:
  6. def __init__(self, name):
  7. self.name = name
  8. def constraint_sql(self, model, schema_editor):
  9. raise NotImplementedError('This method must be implemented by a subclass.')
  10. def create_sql(self, model, schema_editor):
  11. raise NotImplementedError('This method must be implemented by a subclass.')
  12. def remove_sql(self, model, schema_editor):
  13. raise NotImplementedError('This method must be implemented by a subclass.')
  14. def deconstruct(self):
  15. path = '%s.%s' % (self.__class__.__module__, self.__class__.__name__)
  16. path = path.replace('django.db.models.constraints', 'django.db.models')
  17. return (path, (), {'name': self.name})
  18. def clone(self):
  19. _, args, kwargs = self.deconstruct()
  20. return self.__class__(*args, **kwargs)
  21. class CheckConstraint(BaseConstraint):
  22. def __init__(self, *, check, name):
  23. self.check = check
  24. if not getattr(check, 'conditional', False):
  25. raise TypeError(
  26. 'CheckConstraint.check must be a Q instance or boolean '
  27. 'expression.'
  28. )
  29. super().__init__(name)
  30. def _get_check_sql(self, model, schema_editor):
  31. query = Query(model=model, alias_cols=False)
  32. where = query.build_where(self.check)
  33. compiler = query.get_compiler(connection=schema_editor.connection)
  34. sql, params = where.as_sql(compiler, schema_editor.connection)
  35. return sql % tuple(schema_editor.quote_value(p) for p in params)
  36. def constraint_sql(self, model, schema_editor):
  37. check = self._get_check_sql(model, schema_editor)
  38. return schema_editor._check_sql(self.name, check)
  39. def create_sql(self, model, schema_editor):
  40. check = self._get_check_sql(model, schema_editor)
  41. return schema_editor._create_check_sql(model, self.name, check)
  42. def remove_sql(self, model, schema_editor):
  43. return schema_editor._delete_check_sql(model, self.name)
  44. def __repr__(self):
  45. return "<%s: check='%s' name=%r>" % (self.__class__.__name__, self.check, self.name)
  46. def __eq__(self, other):
  47. if isinstance(other, CheckConstraint):
  48. return self.name == other.name and self.check == other.check
  49. return super().__eq__(other)
  50. def deconstruct(self):
  51. path, args, kwargs = super().deconstruct()
  52. kwargs['check'] = self.check
  53. return path, args, kwargs
  54. class Deferrable(Enum):
  55. DEFERRED = 'deferred'
  56. IMMEDIATE = 'immediate'
  57. class UniqueConstraint(BaseConstraint):
  58. def __init__(
  59. self,
  60. *,
  61. fields,
  62. name,
  63. condition=None,
  64. deferrable=None,
  65. include=None,
  66. opclasses=(),
  67. ):
  68. if not fields:
  69. raise ValueError('At least one field is required to define a unique constraint.')
  70. if not isinstance(condition, (type(None), Q)):
  71. raise ValueError('UniqueConstraint.condition must be a Q instance.')
  72. if condition and deferrable:
  73. raise ValueError(
  74. 'UniqueConstraint with conditions cannot be deferred.'
  75. )
  76. if include and deferrable:
  77. raise ValueError(
  78. 'UniqueConstraint with include fields cannot be deferred.'
  79. )
  80. if opclasses and deferrable:
  81. raise ValueError(
  82. 'UniqueConstraint with opclasses cannot be deferred.'
  83. )
  84. if not isinstance(deferrable, (type(None), Deferrable)):
  85. raise ValueError(
  86. 'UniqueConstraint.deferrable must be a Deferrable instance.'
  87. )
  88. if not isinstance(include, (type(None), list, tuple)):
  89. raise ValueError('UniqueConstraint.include must be a list or tuple.')
  90. if not isinstance(opclasses, (list, tuple)):
  91. raise ValueError('UniqueConstraint.opclasses must be a list or tuple.')
  92. if opclasses and len(fields) != len(opclasses):
  93. raise ValueError(
  94. 'UniqueConstraint.fields and UniqueConstraint.opclasses must '
  95. 'have the same number of elements.'
  96. )
  97. self.fields = tuple(fields)
  98. self.condition = condition
  99. self.deferrable = deferrable
  100. self.include = tuple(include) if include else ()
  101. self.opclasses = opclasses
  102. super().__init__(name)
  103. def _get_condition_sql(self, model, schema_editor):
  104. if self.condition is None:
  105. return None
  106. query = Query(model=model, alias_cols=False)
  107. where = query.build_where(self.condition)
  108. compiler = query.get_compiler(connection=schema_editor.connection)
  109. sql, params = where.as_sql(compiler, schema_editor.connection)
  110. return sql % tuple(schema_editor.quote_value(p) for p in params)
  111. def constraint_sql(self, model, schema_editor):
  112. fields = [model._meta.get_field(field_name).column for field_name in self.fields]
  113. include = [model._meta.get_field(field_name).column for field_name in self.include]
  114. condition = self._get_condition_sql(model, schema_editor)
  115. return schema_editor._unique_sql(
  116. model, fields, self.name, condition=condition,
  117. deferrable=self.deferrable, include=include,
  118. opclasses=self.opclasses,
  119. )
  120. def create_sql(self, model, schema_editor):
  121. fields = [model._meta.get_field(field_name).column for field_name in self.fields]
  122. include = [model._meta.get_field(field_name).column for field_name in self.include]
  123. condition = self._get_condition_sql(model, schema_editor)
  124. return schema_editor._create_unique_sql(
  125. model, fields, self.name, condition=condition,
  126. deferrable=self.deferrable, include=include,
  127. opclasses=self.opclasses,
  128. )
  129. def remove_sql(self, model, schema_editor):
  130. condition = self._get_condition_sql(model, schema_editor)
  131. include = [model._meta.get_field(field_name).column for field_name in self.include]
  132. return schema_editor._delete_unique_sql(
  133. model, self.name, condition=condition, deferrable=self.deferrable,
  134. include=include, opclasses=self.opclasses,
  135. )
  136. def __repr__(self):
  137. return '<%s: fields=%r name=%r%s%s%s%s>' % (
  138. self.__class__.__name__, self.fields, self.name,
  139. '' if self.condition is None else ' condition=%s' % self.condition,
  140. '' if self.deferrable is None else ' deferrable=%s' % self.deferrable,
  141. '' if not self.include else ' include=%s' % repr(self.include),
  142. '' if not self.opclasses else ' opclasses=%s' % repr(self.opclasses),
  143. )
  144. def __eq__(self, other):
  145. if isinstance(other, UniqueConstraint):
  146. return (
  147. self.name == other.name and
  148. self.fields == other.fields and
  149. self.condition == other.condition and
  150. self.deferrable == other.deferrable and
  151. self.include == other.include and
  152. self.opclasses == other.opclasses
  153. )
  154. return super().__eq__(other)
  155. def deconstruct(self):
  156. path, args, kwargs = super().deconstruct()
  157. kwargs['fields'] = self.fields
  158. if self.condition:
  159. kwargs['condition'] = self.condition
  160. if self.deferrable:
  161. kwargs['deferrable'] = self.deferrable
  162. if self.include:
  163. kwargs['include'] = self.include
  164. if self.opclasses:
  165. kwargs['opclasses'] = self.opclasses
  166. return path, args, kwargs