schema.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238
  1. import psycopg2
  2. from django.db.backends.base.schema import BaseDatabaseSchemaEditor
  3. from django.db.backends.ddl_references import IndexColumns
  4. from django.db.backends.utils import strip_quotes
  5. class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
  6. sql_create_sequence = "CREATE SEQUENCE %(sequence)s"
  7. sql_delete_sequence = "DROP SEQUENCE IF EXISTS %(sequence)s CASCADE"
  8. sql_set_sequence_max = "SELECT setval('%(sequence)s', MAX(%(column)s)) FROM %(table)s"
  9. sql_set_sequence_owner = 'ALTER SEQUENCE %(sequence)s OWNED BY %(table)s.%(column)s'
  10. sql_create_index = (
  11. 'CREATE INDEX %(name)s ON %(table)s%(using)s '
  12. '(%(columns)s)%(include)s%(extra)s%(condition)s'
  13. )
  14. sql_create_index_concurrently = (
  15. 'CREATE INDEX CONCURRENTLY %(name)s ON %(table)s%(using)s '
  16. '(%(columns)s)%(include)s%(extra)s%(condition)s'
  17. )
  18. sql_delete_index = "DROP INDEX IF EXISTS %(name)s"
  19. sql_delete_index_concurrently = "DROP INDEX CONCURRENTLY IF EXISTS %(name)s"
  20. # Setting the constraint to IMMEDIATE to allow changing data in the same
  21. # transaction.
  22. sql_create_column_inline_fk = (
  23. 'CONSTRAINT %(name)s REFERENCES %(to_table)s(%(to_column)s)%(deferrable)s'
  24. '; SET CONSTRAINTS %(namespace)s%(name)s IMMEDIATE'
  25. )
  26. # Setting the constraint to IMMEDIATE runs any deferred checks to allow
  27. # dropping it in the same transaction.
  28. sql_delete_fk = "SET CONSTRAINTS %(name)s IMMEDIATE; ALTER TABLE %(table)s DROP CONSTRAINT %(name)s"
  29. sql_delete_procedure = 'DROP FUNCTION %(procedure)s(%(param_types)s)'
  30. def quote_value(self, value):
  31. if isinstance(value, str):
  32. value = value.replace('%', '%%')
  33. adapted = psycopg2.extensions.adapt(value)
  34. if hasattr(adapted, 'encoding'):
  35. adapted.encoding = 'utf8'
  36. # getquoted() returns a quoted bytestring of the adapted value.
  37. return adapted.getquoted().decode()
  38. def _field_indexes_sql(self, model, field):
  39. output = super()._field_indexes_sql(model, field)
  40. like_index_statement = self._create_like_index_sql(model, field)
  41. if like_index_statement is not None:
  42. output.append(like_index_statement)
  43. return output
  44. def _field_data_type(self, field):
  45. if field.is_relation:
  46. return field.rel_db_type(self.connection)
  47. return self.connection.data_types.get(
  48. field.get_internal_type(),
  49. field.db_type(self.connection),
  50. )
  51. def _field_base_data_types(self, field):
  52. # Yield base data types for array fields.
  53. if field.base_field.get_internal_type() == 'ArrayField':
  54. yield from self._field_base_data_types(field.base_field)
  55. else:
  56. yield self._field_data_type(field.base_field)
  57. def _create_like_index_sql(self, model, field):
  58. """
  59. Return the statement to create an index with varchar operator pattern
  60. when the column type is 'varchar' or 'text', otherwise return None.
  61. """
  62. db_type = field.db_type(connection=self.connection)
  63. if db_type is not None and (field.db_index or field.unique):
  64. # Fields with database column types of `varchar` and `text` need
  65. # a second index that specifies their operator class, which is
  66. # needed when performing correct LIKE queries outside the
  67. # C locale. See #12234.
  68. #
  69. # The same doesn't apply to array fields such as varchar[size]
  70. # and text[size], so skip them.
  71. if '[' in db_type:
  72. return None
  73. if db_type.startswith('varchar'):
  74. return self._create_index_sql(
  75. model,
  76. fields=[field],
  77. suffix='_like',
  78. opclasses=['varchar_pattern_ops'],
  79. )
  80. elif db_type.startswith('text'):
  81. return self._create_index_sql(
  82. model,
  83. fields=[field],
  84. suffix='_like',
  85. opclasses=['text_pattern_ops'],
  86. )
  87. return None
  88. def _alter_column_type_sql(self, model, old_field, new_field, new_type):
  89. self.sql_alter_column_type = 'ALTER COLUMN %(column)s TYPE %(type)s'
  90. # Cast when data type changed.
  91. using_sql = ' USING %(column)s::%(type)s'
  92. new_internal_type = new_field.get_internal_type()
  93. old_internal_type = old_field.get_internal_type()
  94. if new_internal_type == 'ArrayField' and new_internal_type == old_internal_type:
  95. # Compare base data types for array fields.
  96. if list(self._field_base_data_types(old_field)) != list(self._field_base_data_types(new_field)):
  97. self.sql_alter_column_type += using_sql
  98. elif self._field_data_type(old_field) != self._field_data_type(new_field):
  99. self.sql_alter_column_type += using_sql
  100. # Make ALTER TYPE with SERIAL make sense.
  101. table = strip_quotes(model._meta.db_table)
  102. serial_fields_map = {'bigserial': 'bigint', 'serial': 'integer', 'smallserial': 'smallint'}
  103. if new_type.lower() in serial_fields_map:
  104. column = strip_quotes(new_field.column)
  105. sequence_name = "%s_%s_seq" % (table, column)
  106. return (
  107. (
  108. self.sql_alter_column_type % {
  109. "column": self.quote_name(column),
  110. "type": serial_fields_map[new_type.lower()],
  111. },
  112. [],
  113. ),
  114. [
  115. (
  116. self.sql_delete_sequence % {
  117. "sequence": self.quote_name(sequence_name),
  118. },
  119. [],
  120. ),
  121. (
  122. self.sql_create_sequence % {
  123. "sequence": self.quote_name(sequence_name),
  124. },
  125. [],
  126. ),
  127. (
  128. self.sql_alter_column % {
  129. "table": self.quote_name(table),
  130. "changes": self.sql_alter_column_default % {
  131. "column": self.quote_name(column),
  132. "default": "nextval('%s')" % self.quote_name(sequence_name),
  133. }
  134. },
  135. [],
  136. ),
  137. (
  138. self.sql_set_sequence_max % {
  139. "table": self.quote_name(table),
  140. "column": self.quote_name(column),
  141. "sequence": self.quote_name(sequence_name),
  142. },
  143. [],
  144. ),
  145. (
  146. self.sql_set_sequence_owner % {
  147. 'table': self.quote_name(table),
  148. 'column': self.quote_name(column),
  149. 'sequence': self.quote_name(sequence_name),
  150. },
  151. [],
  152. ),
  153. ],
  154. )
  155. elif old_field.db_parameters(connection=self.connection)['type'] in serial_fields_map:
  156. # Drop the sequence if migrating away from AutoField.
  157. column = strip_quotes(new_field.column)
  158. sequence_name = '%s_%s_seq' % (table, column)
  159. fragment, _ = super()._alter_column_type_sql(model, old_field, new_field, new_type)
  160. return fragment, [
  161. (
  162. self.sql_delete_sequence % {
  163. 'sequence': self.quote_name(sequence_name),
  164. },
  165. [],
  166. ),
  167. ]
  168. else:
  169. return super()._alter_column_type_sql(model, old_field, new_field, new_type)
  170. def _alter_field(self, model, old_field, new_field, old_type, new_type,
  171. old_db_params, new_db_params, strict=False):
  172. # Drop indexes on varchar/text/citext columns that are changing to a
  173. # different type.
  174. if (old_field.db_index or old_field.unique) and (
  175. (old_type.startswith('varchar') and not new_type.startswith('varchar')) or
  176. (old_type.startswith('text') and not new_type.startswith('text')) or
  177. (old_type.startswith('citext') and not new_type.startswith('citext'))
  178. ):
  179. index_name = self._create_index_name(model._meta.db_table, [old_field.column], suffix='_like')
  180. self.execute(self._delete_index_sql(model, index_name))
  181. super()._alter_field(
  182. model, old_field, new_field, old_type, new_type, old_db_params,
  183. new_db_params, strict,
  184. )
  185. # Added an index? Create any PostgreSQL-specific indexes.
  186. if ((not (old_field.db_index or old_field.unique) and new_field.db_index) or
  187. (not old_field.unique and new_field.unique)):
  188. like_index_statement = self._create_like_index_sql(model, new_field)
  189. if like_index_statement is not None:
  190. self.execute(like_index_statement)
  191. # Removed an index? Drop any PostgreSQL-specific indexes.
  192. if old_field.unique and not (new_field.db_index or new_field.unique):
  193. index_to_remove = self._create_index_name(model._meta.db_table, [old_field.column], suffix='_like')
  194. self.execute(self._delete_index_sql(model, index_to_remove))
  195. def _index_columns(self, table, columns, col_suffixes, opclasses):
  196. if opclasses:
  197. return IndexColumns(table, columns, self.quote_name, col_suffixes=col_suffixes, opclasses=opclasses)
  198. return super()._index_columns(table, columns, col_suffixes, opclasses)
  199. def add_index(self, model, index, concurrently=False):
  200. self.execute(index.create_sql(model, self, concurrently=concurrently), params=None)
  201. def remove_index(self, model, index, concurrently=False):
  202. self.execute(index.remove_sql(model, self, concurrently=concurrently))
  203. def _delete_index_sql(self, model, name, sql=None, concurrently=False):
  204. sql = self.sql_delete_index_concurrently if concurrently else self.sql_delete_index
  205. return super()._delete_index_sql(model, name, sql)
  206. def _create_index_sql(
  207. self, model, *, fields=None, name=None, suffix='', using='',
  208. db_tablespace=None, col_suffixes=(), sql=None, opclasses=(),
  209. condition=None, concurrently=False, include=None, expressions=None,
  210. ):
  211. sql = self.sql_create_index if not concurrently else self.sql_create_index_concurrently
  212. return super()._create_index_sql(
  213. model, fields=fields, name=name, suffix=suffix, using=using,
  214. db_tablespace=db_tablespace, col_suffixes=col_suffixes, sql=sql,
  215. opclasses=opclasses, condition=condition, include=include,
  216. expressions=expressions,
  217. )