schema.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388
  1. from django.db.backends.base.schema import BaseDatabaseSchemaEditor
  2. from django.db.backends.ddl_references import IndexColumns
  3. from django.db.backends.postgresql.psycopg_any import sql
  4. from django.db.backends.utils import strip_quotes
  5. class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
  6. # Setting all constraints to IMMEDIATE to allow changing data in the same
  7. # transaction.
  8. sql_update_with_default = (
  9. "UPDATE %(table)s SET %(column)s = %(default)s WHERE %(column)s IS NULL"
  10. "; SET CONSTRAINTS ALL IMMEDIATE"
  11. )
  12. sql_alter_sequence_type = "ALTER SEQUENCE IF EXISTS %(sequence)s AS %(type)s"
  13. sql_delete_sequence = "DROP SEQUENCE IF EXISTS %(sequence)s CASCADE"
  14. sql_create_index = (
  15. "CREATE INDEX %(name)s ON %(table)s%(using)s "
  16. "(%(columns)s)%(include)s%(extra)s%(condition)s"
  17. )
  18. sql_create_index_concurrently = (
  19. "CREATE INDEX CONCURRENTLY %(name)s ON %(table)s%(using)s "
  20. "(%(columns)s)%(include)s%(extra)s%(condition)s"
  21. )
  22. sql_delete_index = "DROP INDEX IF EXISTS %(name)s"
  23. sql_delete_index_concurrently = "DROP INDEX CONCURRENTLY IF EXISTS %(name)s"
  24. # Setting the constraint to IMMEDIATE to allow changing data in the same
  25. # transaction.
  26. sql_create_column_inline_fk = (
  27. "CONSTRAINT %(name)s REFERENCES %(to_table)s(%(to_column)s)%(deferrable)s"
  28. "; SET CONSTRAINTS %(namespace)s%(name)s IMMEDIATE"
  29. )
  30. # Setting the constraint to IMMEDIATE runs any deferred checks to allow
  31. # dropping it in the same transaction.
  32. sql_delete_fk = (
  33. "SET CONSTRAINTS %(name)s IMMEDIATE; "
  34. "ALTER TABLE %(table)s DROP CONSTRAINT %(name)s"
  35. )
  36. sql_delete_procedure = "DROP FUNCTION %(procedure)s(%(param_types)s)"
  37. def execute(self, sql, params=()):
  38. # Merge the query client-side, as PostgreSQL won't do it server-side.
  39. if params is None:
  40. return super().execute(sql, params)
  41. sql = self.connection.ops.compose_sql(str(sql), params)
  42. # Don't let the superclass touch anything.
  43. return super().execute(sql, None)
  44. sql_add_identity = (
  45. "ALTER TABLE %(table)s ALTER COLUMN %(column)s ADD "
  46. "GENERATED BY DEFAULT AS IDENTITY"
  47. )
  48. sql_drop_indentity = (
  49. "ALTER TABLE %(table)s ALTER COLUMN %(column)s DROP IDENTITY IF EXISTS"
  50. )
  51. def quote_value(self, value):
  52. return sql.quote(value, self.connection.connection)
  53. def _field_indexes_sql(self, model, field):
  54. output = super()._field_indexes_sql(model, field)
  55. like_index_statement = self._create_like_index_sql(model, field)
  56. if like_index_statement is not None:
  57. output.append(like_index_statement)
  58. return output
  59. def _field_data_type(self, field):
  60. if field.is_relation:
  61. return field.rel_db_type(self.connection)
  62. return self.connection.data_types.get(
  63. field.get_internal_type(),
  64. field.db_type(self.connection),
  65. )
  66. def _field_base_data_types(self, field):
  67. # Yield base data types for array fields.
  68. if field.base_field.get_internal_type() == "ArrayField":
  69. yield from self._field_base_data_types(field.base_field)
  70. else:
  71. yield self._field_data_type(field.base_field)
  72. def _create_like_index_sql(self, model, field):
  73. """
  74. Return the statement to create an index with varchar operator pattern
  75. when the column type is 'varchar' or 'text', otherwise return None.
  76. """
  77. db_type = field.db_type(connection=self.connection)
  78. if db_type is not None and (field.db_index or field.unique):
  79. # Fields with database column types of `varchar` and `text` need
  80. # a second index that specifies their operator class, which is
  81. # needed when performing correct LIKE queries outside the
  82. # C locale. See #12234.
  83. #
  84. # The same doesn't apply to array fields such as varchar[size]
  85. # and text[size], so skip them.
  86. if "[" in db_type:
  87. return None
  88. # Non-deterministic collations on Postgresql don't support indexes
  89. # for operator classes varchar_pattern_ops/text_pattern_ops.
  90. collation_name = getattr(field, "db_collation", None)
  91. if not collation_name and field.is_relation:
  92. collation_name = getattr(field.target_field, "db_collation", None)
  93. if collation_name and not self._is_collation_deterministic(collation_name):
  94. return None
  95. if db_type.startswith("varchar"):
  96. return self._create_index_sql(
  97. model,
  98. fields=[field],
  99. suffix="_like",
  100. opclasses=["varchar_pattern_ops"],
  101. )
  102. elif db_type.startswith("text"):
  103. return self._create_index_sql(
  104. model,
  105. fields=[field],
  106. suffix="_like",
  107. opclasses=["text_pattern_ops"],
  108. )
  109. return None
  110. def _using_sql(self, new_field, old_field):
  111. using_sql = " USING %(column)s::%(type)s"
  112. new_internal_type = new_field.get_internal_type()
  113. old_internal_type = old_field.get_internal_type()
  114. if new_internal_type == "ArrayField" and new_internal_type == old_internal_type:
  115. # Compare base data types for array fields.
  116. if list(self._field_base_data_types(old_field)) != list(
  117. self._field_base_data_types(new_field)
  118. ):
  119. return using_sql
  120. elif self._field_data_type(old_field) != self._field_data_type(new_field):
  121. return using_sql
  122. return ""
  123. def _get_sequence_name(self, table, column):
  124. with self.connection.cursor() as cursor:
  125. for sequence in self.connection.introspection.get_sequences(cursor, table):
  126. if sequence["column"] == column:
  127. return sequence["name"]
  128. return None
  129. def _alter_column_type_sql(
  130. self, model, old_field, new_field, new_type, old_collation, new_collation
  131. ):
  132. # Drop indexes on varchar/text/citext columns that are changing to a
  133. # different type.
  134. old_db_params = old_field.db_parameters(connection=self.connection)
  135. old_type = old_db_params["type"]
  136. if (old_field.db_index or old_field.unique) and (
  137. (old_type.startswith("varchar") and not new_type.startswith("varchar"))
  138. or (old_type.startswith("text") and not new_type.startswith("text"))
  139. or (old_type.startswith("citext") and not new_type.startswith("citext"))
  140. ):
  141. index_name = self._create_index_name(
  142. model._meta.db_table, [old_field.column], suffix="_like"
  143. )
  144. self.execute(self._delete_index_sql(model, index_name))
  145. self.sql_alter_column_type = (
  146. "ALTER COLUMN %(column)s TYPE %(type)s%(collation)s"
  147. )
  148. # Cast when data type changed.
  149. if using_sql := self._using_sql(new_field, old_field):
  150. self.sql_alter_column_type += using_sql
  151. new_internal_type = new_field.get_internal_type()
  152. old_internal_type = old_field.get_internal_type()
  153. # Make ALTER TYPE with IDENTITY make sense.
  154. table = strip_quotes(model._meta.db_table)
  155. auto_field_types = {
  156. "AutoField",
  157. "BigAutoField",
  158. "SmallAutoField",
  159. }
  160. old_is_auto = old_internal_type in auto_field_types
  161. new_is_auto = new_internal_type in auto_field_types
  162. if new_is_auto and not old_is_auto:
  163. column = strip_quotes(new_field.column)
  164. return (
  165. (
  166. self.sql_alter_column_type
  167. % {
  168. "column": self.quote_name(column),
  169. "type": new_type,
  170. "collation": "",
  171. },
  172. [],
  173. ),
  174. [
  175. (
  176. self.sql_add_identity
  177. % {
  178. "table": self.quote_name(table),
  179. "column": self.quote_name(column),
  180. },
  181. [],
  182. ),
  183. ],
  184. )
  185. elif old_is_auto and not new_is_auto:
  186. # Drop IDENTITY if exists (pre-Django 4.1 serial columns don't have
  187. # it).
  188. self.execute(
  189. self.sql_drop_indentity
  190. % {
  191. "table": self.quote_name(table),
  192. "column": self.quote_name(strip_quotes(new_field.column)),
  193. }
  194. )
  195. column = strip_quotes(new_field.column)
  196. fragment, _ = super()._alter_column_type_sql(
  197. model, old_field, new_field, new_type, old_collation, new_collation
  198. )
  199. # Drop the sequence if exists (Django 4.1+ identity columns don't
  200. # have it).
  201. other_actions = []
  202. if sequence_name := self._get_sequence_name(table, column):
  203. other_actions = [
  204. (
  205. self.sql_delete_sequence
  206. % {
  207. "sequence": self.quote_name(sequence_name),
  208. },
  209. [],
  210. )
  211. ]
  212. return fragment, other_actions
  213. elif new_is_auto and old_is_auto and old_internal_type != new_internal_type:
  214. fragment, _ = super()._alter_column_type_sql(
  215. model, old_field, new_field, new_type, old_collation, new_collation
  216. )
  217. column = strip_quotes(new_field.column)
  218. db_types = {
  219. "AutoField": "integer",
  220. "BigAutoField": "bigint",
  221. "SmallAutoField": "smallint",
  222. }
  223. # Alter the sequence type if exists (Django 4.1+ identity columns
  224. # don't have it).
  225. other_actions = []
  226. if sequence_name := self._get_sequence_name(table, column):
  227. other_actions = [
  228. (
  229. self.sql_alter_sequence_type
  230. % {
  231. "sequence": self.quote_name(sequence_name),
  232. "type": db_types[new_internal_type],
  233. },
  234. [],
  235. ),
  236. ]
  237. return fragment, other_actions
  238. else:
  239. return super()._alter_column_type_sql(
  240. model, old_field, new_field, new_type, old_collation, new_collation
  241. )
  242. def _alter_column_collation_sql(
  243. self, model, new_field, new_type, new_collation, old_field
  244. ):
  245. sql = self.sql_alter_column_collate
  246. # Cast when data type changed.
  247. if using_sql := self._using_sql(new_field, old_field):
  248. sql += using_sql
  249. return (
  250. sql
  251. % {
  252. "column": self.quote_name(new_field.column),
  253. "type": new_type,
  254. "collation": " " + self._collate_sql(new_collation)
  255. if new_collation
  256. else "",
  257. },
  258. [],
  259. )
  260. def _alter_field(
  261. self,
  262. model,
  263. old_field,
  264. new_field,
  265. old_type,
  266. new_type,
  267. old_db_params,
  268. new_db_params,
  269. strict=False,
  270. ):
  271. super()._alter_field(
  272. model,
  273. old_field,
  274. new_field,
  275. old_type,
  276. new_type,
  277. old_db_params,
  278. new_db_params,
  279. strict,
  280. )
  281. # Added an index? Create any PostgreSQL-specific indexes.
  282. if (not (old_field.db_index or old_field.unique) and new_field.db_index) or (
  283. not old_field.unique and new_field.unique
  284. ):
  285. like_index_statement = self._create_like_index_sql(model, new_field)
  286. if like_index_statement is not None:
  287. self.execute(like_index_statement)
  288. # Removed an index? Drop any PostgreSQL-specific indexes.
  289. if old_field.unique and not (new_field.db_index or new_field.unique):
  290. index_to_remove = self._create_index_name(
  291. model._meta.db_table, [old_field.column], suffix="_like"
  292. )
  293. self.execute(self._delete_index_sql(model, index_to_remove))
  294. def _index_columns(self, table, columns, col_suffixes, opclasses):
  295. if opclasses:
  296. return IndexColumns(
  297. table,
  298. columns,
  299. self.quote_name,
  300. col_suffixes=col_suffixes,
  301. opclasses=opclasses,
  302. )
  303. return super()._index_columns(table, columns, col_suffixes, opclasses)
  304. def add_index(self, model, index, concurrently=False):
  305. self.execute(
  306. index.create_sql(model, self, concurrently=concurrently), params=None
  307. )
  308. def remove_index(self, model, index, concurrently=False):
  309. self.execute(index.remove_sql(model, self, concurrently=concurrently))
  310. def _delete_index_sql(self, model, name, sql=None, concurrently=False):
  311. sql = (
  312. self.sql_delete_index_concurrently
  313. if concurrently
  314. else self.sql_delete_index
  315. )
  316. return super()._delete_index_sql(model, name, sql)
  317. def _create_index_sql(
  318. self,
  319. model,
  320. *,
  321. fields=None,
  322. name=None,
  323. suffix="",
  324. using="",
  325. db_tablespace=None,
  326. col_suffixes=(),
  327. sql=None,
  328. opclasses=(),
  329. condition=None,
  330. concurrently=False,
  331. include=None,
  332. expressions=None,
  333. ):
  334. sql = sql or (
  335. self.sql_create_index
  336. if not concurrently
  337. else self.sql_create_index_concurrently
  338. )
  339. return super()._create_index_sql(
  340. model,
  341. fields=fields,
  342. name=name,
  343. suffix=suffix,
  344. using=using,
  345. db_tablespace=db_tablespace,
  346. col_suffixes=col_suffixes,
  347. sql=sql,
  348. opclasses=opclasses,
  349. condition=condition,
  350. include=include,
  351. expressions=expressions,
  352. )
  353. def _is_collation_deterministic(self, collation_name):
  354. with self.connection.cursor() as cursor:
  355. cursor.execute(
  356. """
  357. SELECT collisdeterministic
  358. FROM pg_collation
  359. WHERE collname = %s
  360. """,
  361. [collation_name],
  362. )
  363. row = cursor.fetchone()
  364. return row[0] if row else None