schema.py 37 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098
  1. from typing import TYPE_CHECKING, Any, List, Optional, Type, cast
  2. from unittest import mock
  3. import django
  4. from django.core.exceptions import (
  5. FieldDoesNotExist,
  6. ImproperlyConfigured,
  7. SuspiciousOperation,
  8. )
  9. from django.db import transaction
  10. from django.db.backends.ddl_references import Statement
  11. from django.db.backends.postgresql.schema import ( # type: ignore[import]
  12. DatabaseSchemaEditor,
  13. )
  14. from django.db.models import Field, Model
  15. from psqlextra.settings import (
  16. postgres_prepend_local_search_path,
  17. postgres_reset_local_search_path,
  18. )
  19. from psqlextra.type_assertions import is_sql_with_params
  20. from psqlextra.types import PostgresPartitioningMethod
  21. from . import base_impl
  22. from .introspection import PostgresIntrospection
  23. from .side_effects import (
  24. HStoreRequiredSchemaEditorSideEffect,
  25. HStoreUniqueSchemaEditorSideEffect,
  26. )
  27. if TYPE_CHECKING:
  28. class SchemaEditor(DatabaseSchemaEditor):
  29. pass
  30. else:
  31. SchemaEditor = base_impl.schema_editor()
  32. class PostgresSchemaEditor(SchemaEditor):
  33. """Schema editor that adds extra methods for PostgreSQL specific features
  34. and hooks into existing implementations to add side effects specific to
  35. PostgreSQL."""
  36. sql_add_pk = "ALTER TABLE %s ADD PRIMARY KEY (%s)"
  37. sql_create_fk_not_valid = f"{SchemaEditor.sql_create_fk} NOT VALID"
  38. sql_validate_fk = "ALTER TABLE %s VALIDATE CONSTRAINT %s"
  39. sql_create_sequence_with_owner = "CREATE SEQUENCE %s OWNED BY %s.%s"
  40. sql_alter_table_storage_setting = "ALTER TABLE %s SET (%s = %s)"
  41. sql_reset_table_storage_setting = "ALTER TABLE %s RESET (%s)"
  42. sql_alter_table_schema = "ALTER TABLE %s SET SCHEMA %s"
  43. sql_create_schema = "CREATE SCHEMA %s"
  44. sql_delete_schema = "DROP SCHEMA %s"
  45. sql_delete_schema_cascade = "DROP SCHEMA %s CASCADE"
  46. sql_create_view = "CREATE VIEW %s AS (%s)"
  47. sql_replace_view = "CREATE OR REPLACE VIEW %s AS (%s)"
  48. sql_drop_view = "DROP VIEW IF EXISTS %s"
  49. sql_create_materialized_view = (
  50. "CREATE MATERIALIZED VIEW %s AS (%s) WITH DATA"
  51. )
  52. sql_drop_materialized_view = "DROP MATERIALIZED VIEW %s"
  53. sql_refresh_materialized_view = "REFRESH MATERIALIZED VIEW %s"
  54. sql_refresh_materialized_view_concurrently = (
  55. "REFRESH MATERIALIZED VIEW CONCURRENTLY %s"
  56. )
  57. sql_partition_by = " PARTITION BY %s (%s)"
  58. sql_add_default_partition = "CREATE TABLE %s PARTITION OF %s DEFAULT"
  59. sql_add_hash_partition = "CREATE TABLE %s PARTITION OF %s FOR VALUES WITH (MODULUS %s, REMAINDER %s)"
  60. sql_add_range_partition = (
  61. "CREATE TABLE %s PARTITION OF %s FOR VALUES FROM (%s) TO (%s)"
  62. )
  63. sql_add_list_partition = (
  64. "CREATE TABLE %s PARTITION OF %s FOR VALUES IN (%s)"
  65. )
  66. sql_delete_partition = "DROP TABLE %s"
  67. sql_table_comment = "COMMENT ON TABLE %s IS %s"
  68. side_effects: List[DatabaseSchemaEditor] = [
  69. cast(DatabaseSchemaEditor, HStoreUniqueSchemaEditorSideEffect()),
  70. cast(DatabaseSchemaEditor, HStoreRequiredSchemaEditorSideEffect()),
  71. ]
  72. def __init__(self, connection, collect_sql=False, atomic=True):
  73. super().__init__(connection, collect_sql, atomic)
  74. for side_effect in self.side_effects:
  75. side_effect.execute = self.execute
  76. side_effect.quote_name = self.quote_name
  77. self.deferred_sql = []
  78. self.introspection = PostgresIntrospection(self.connection)
  79. def create_schema(self, name: str) -> None:
  80. """Creates a Postgres schema."""
  81. self.execute(self.sql_create_schema % self.quote_name(name))
  82. def delete_schema(self, name: str, cascade: bool) -> None:
  83. """Drops a Postgres schema."""
  84. sql = (
  85. self.sql_delete_schema
  86. if not cascade
  87. else self.sql_delete_schema_cascade
  88. )
  89. self.execute(sql % self.quote_name(name))
  90. def create_model(self, model: Type[Model]) -> None:
  91. """Creates a new model."""
  92. super().create_model(model)
  93. for side_effect in self.side_effects:
  94. side_effect.create_model(model)
  95. def delete_model(self, model: Type[Model]) -> None:
  96. """Drops/deletes an existing model."""
  97. for side_effect in self.side_effects:
  98. side_effect.delete_model(model)
  99. super().delete_model(model)
  100. def clone_model_structure_to_schema(
  101. self, model: Type[Model], *, schema_name: str
  102. ) -> None:
  103. """Creates a clone of the columns for the specified model in a separate
  104. schema.
  105. The table will have exactly the same name as the model table
  106. in the default schema. It will have none of the constraints,
  107. foreign keys and indexes.
  108. Use this to create a temporary clone of a model table to
  109. replace the original model table later on. The lack of
  110. indices and constraints allows for greater write speeds.
  111. The original model table will be unaffected.
  112. Arguments:
  113. model:
  114. Model to clone the table of into the
  115. specified schema.
  116. schema_name:
  117. Name of the schema to create the cloned
  118. table in.
  119. """
  120. table_name = model._meta.db_table
  121. quoted_table_name = self.quote_name(model._meta.db_table)
  122. quoted_schema_name = self.quote_name(schema_name)
  123. quoted_table_fqn = f"{quoted_schema_name}.{quoted_table_name}"
  124. self.execute(
  125. self.sql_create_table
  126. % {
  127. "table": quoted_table_fqn,
  128. "definition": f"LIKE {quoted_table_name} INCLUDING ALL EXCLUDING CONSTRAINTS EXCLUDING INDEXES",
  129. }
  130. )
  131. # Copy sequences
  132. #
  133. # Django 4.0 and older do not use IDENTITY so Postgres does
  134. # not copy the sequences into the new table. We do it manually.
  135. if django.VERSION < (4, 1):
  136. with self.connection.cursor() as cursor:
  137. sequences = self.introspection.get_sequences(cursor, table_name)
  138. for sequence in sequences:
  139. if sequence["table"] != table_name:
  140. continue
  141. quoted_sequence_name = self.quote_name(sequence["name"])
  142. quoted_sequence_fqn = (
  143. f"{quoted_schema_name}.{quoted_sequence_name}"
  144. )
  145. quoted_column_name = self.quote_name(sequence["column"])
  146. self.execute(
  147. self.sql_create_sequence_with_owner
  148. % (
  149. quoted_sequence_fqn,
  150. quoted_table_fqn,
  151. quoted_column_name,
  152. )
  153. )
  154. self.execute(
  155. self.sql_alter_column
  156. % {
  157. "table": quoted_table_fqn,
  158. "changes": self.sql_alter_column_default
  159. % {
  160. "column": quoted_column_name,
  161. "default": "nextval('%s')" % quoted_sequence_fqn,
  162. },
  163. }
  164. )
  165. # Copy storage settings
  166. #
  167. # Postgres only copies column-level storage options, not
  168. # the table-level storage options.
  169. with self.connection.cursor() as cursor:
  170. storage_settings = self.introspection.get_storage_settings(
  171. cursor, model._meta.db_table
  172. )
  173. for setting_name, setting_value in storage_settings.items():
  174. self.alter_table_storage_setting(
  175. quoted_table_fqn, setting_name, setting_value
  176. )
  177. def clone_model_constraints_and_indexes_to_schema(
  178. self, model: Type[Model], *, schema_name: str
  179. ) -> None:
  180. """Adds the constraints, foreign keys and indexes to a model table that
  181. was cloned into a separate table without them by
  182. `clone_model_structure_to_schema`.
  183. Arguments:
  184. model:
  185. Model for which the cloned table was created.
  186. schema_name:
  187. Name of the schema in which the cloned table
  188. resides.
  189. """
  190. with postgres_prepend_local_search_path(
  191. [schema_name], using=self.connection.alias
  192. ):
  193. for constraint in model._meta.constraints:
  194. self.add_constraint(model, constraint) # type: ignore[attr-defined]
  195. for index in model._meta.indexes:
  196. self.add_index(model, index)
  197. if model._meta.unique_together:
  198. self.alter_unique_together(
  199. model, tuple(), model._meta.unique_together
  200. )
  201. if model._meta.index_together:
  202. self.alter_index_together(
  203. model, tuple(), model._meta.index_together
  204. )
  205. for field in model._meta.local_concrete_fields: # type: ignore[attr-defined]
  206. # Django creates primary keys later added to the model with
  207. # a custom name. We want the name as it was created originally.
  208. if field.primary_key:
  209. with postgres_reset_local_search_path(
  210. using=self.connection.alias
  211. ):
  212. [primary_key_name] = self._constraint_names( # type: ignore[attr-defined]
  213. model, primary_key=True
  214. )
  215. self.execute(
  216. self.sql_create_pk
  217. % {
  218. "table": self.quote_name(model._meta.db_table),
  219. "name": self.quote_name(primary_key_name),
  220. "columns": self.quote_name(
  221. field.db_column or field.attname
  222. ),
  223. }
  224. )
  225. continue
  226. # Django creates foreign keys in a single statement which acquires
  227. # a AccessExclusiveLock on the referenced table. We want to avoid
  228. # that and created the FK as NOT VALID. We can run VALIDATE in
  229. # a separate transaction later to validate the entries without
  230. # acquiring a AccessExclusiveLock.
  231. if field.remote_field:
  232. with postgres_reset_local_search_path(
  233. using=self.connection.alias
  234. ):
  235. [fk_name] = self._constraint_names( # type: ignore[attr-defined]
  236. model, [field.column], foreign_key=True
  237. )
  238. sql = Statement(
  239. self.sql_create_fk_not_valid,
  240. table=self.quote_name(model._meta.db_table),
  241. name=self.quote_name(fk_name),
  242. column=self.quote_name(field.column),
  243. to_table=self.quote_name(
  244. field.target_field.model._meta.db_table
  245. ),
  246. to_column=self.quote_name(field.target_field.column),
  247. deferrable=self.connection.ops.deferrable_sql(),
  248. )
  249. self.execute(sql)
  250. # It's hard to alter a field's check because it is defined
  251. # by the field class, not the field instance. Handle this
  252. # manually.
  253. field_check = field.db_parameters(self.connection).get("check")
  254. if field_check:
  255. with postgres_reset_local_search_path(
  256. using=self.connection.alias
  257. ):
  258. [field_check_name] = self._constraint_names( # type: ignore[attr-defined]
  259. model,
  260. [field.column],
  261. check=True,
  262. exclude={
  263. constraint.name
  264. for constraint in model._meta.constraints
  265. },
  266. )
  267. self.execute(
  268. self._create_check_sql( # type: ignore[attr-defined]
  269. model, field_check_name, field_check
  270. )
  271. )
  272. # Clone the field and alter its state to math our current
  273. # table definition. This will cause Django see the missing
  274. # indices and create them.
  275. if field.remote_field:
  276. # We add the foreign key constraint ourselves with NOT VALID,
  277. # hence, we specify `db_constraint=False` on both old/new.
  278. # Django won't touch the foreign key constraint.
  279. old_field = self._clone_model_field(
  280. field, db_index=False, unique=False, db_constraint=False
  281. )
  282. new_field = self._clone_model_field(
  283. field, db_constraint=False
  284. )
  285. self.alter_field(model, old_field, new_field)
  286. else:
  287. old_field = self._clone_model_field(
  288. field, db_index=False, unique=False
  289. )
  290. new_field = self._clone_model_field(field)
  291. self.alter_field(model, old_field, new_field)
  292. def clone_model_foreign_keys_to_schema(
  293. self, model: Type[Model], schema_name: str
  294. ) -> None:
  295. """Validates the foreign keys in the cloned model table created by
  296. `clone_model_structure_to_schema` and
  297. `clone_model_constraints_and_indexes_to_schema`.
  298. Do NOT run this in the same transaction as the
  299. foreign keys were added to the table. It WILL
  300. acquire a long-lived AccessExclusiveLock.
  301. Arguments:
  302. model:
  303. Model for which the cloned table was created.
  304. schema_name:
  305. Name of the schema in which the cloned table
  306. resides.
  307. """
  308. constraint_names = self._constraint_names(model, foreign_key=True) # type: ignore[attr-defined]
  309. with postgres_prepend_local_search_path(
  310. [schema_name], using=self.connection.alias
  311. ):
  312. for fk_name in constraint_names:
  313. self.execute(
  314. self.sql_validate_fk
  315. % (
  316. self.quote_name(model._meta.db_table),
  317. self.quote_name(fk_name),
  318. )
  319. )
  320. def alter_table_storage_setting(
  321. self, table_name: str, name: str, value: str
  322. ) -> None:
  323. """Alters a storage setting for a table.
  324. See: https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-STORAGE-PARAMETERS
  325. Arguments:
  326. table_name:
  327. Name of the table to alter the setting for.
  328. name:
  329. Name of the setting to alter.
  330. value:
  331. Value to alter the setting to.
  332. Note that this is always a string, even if it looks
  333. like a number or a boolean. That's how Postgres
  334. stores storage settings internally.
  335. """
  336. self.execute(
  337. self.sql_alter_table_storage_setting
  338. % (self.quote_name(table_name), name, value)
  339. )
  340. def alter_model_storage_setting(
  341. self, model: Type[Model], name: str, value: str
  342. ) -> None:
  343. """Alters a storage setting for the model's table.
  344. See: https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-STORAGE-PARAMETERS
  345. Arguments:
  346. model:
  347. Model of which to alter the table
  348. setting.
  349. name:
  350. Name of the setting to alter.
  351. value:
  352. Value to alter the setting to.
  353. Note that this is always a string, even if it looks
  354. like a number or a boolean. That's how Postgres
  355. stores storage settings internally.
  356. """
  357. self.alter_table_storage_setting(model._meta.db_table, name, value)
  358. def reset_table_storage_setting(self, table_name: str, name: str) -> None:
  359. """Resets a table's storage setting to the database or server default.
  360. See: https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-STORAGE-PARAMETERS
  361. Arguments:
  362. table_name:
  363. Name of the table to reset the setting for.
  364. name:
  365. Name of the setting to reset.
  366. """
  367. self.execute(
  368. self.sql_reset_table_storage_setting
  369. % (self.quote_name(table_name), name)
  370. )
  371. def reset_model_storage_setting(
  372. self, model: Type[Model], name: str
  373. ) -> None:
  374. """Resets a model's table storage setting to the database or server
  375. default.
  376. See: https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-STORAGE-PARAMETERS
  377. Arguments:
  378. table_name:
  379. model:
  380. Model for which to reset the table setting for.
  381. name:
  382. Name of the setting to reset.
  383. """
  384. self.reset_table_storage_setting(model._meta.db_table, name)
  385. def alter_table_schema(self, table_name: str, schema_name: str) -> None:
  386. """Moves the specified table into the specified schema.
  387. WARNING: Moving models into a different schema than the default
  388. will break querying the model.
  389. Arguments:
  390. table_name:
  391. Name of the table to move into the specified schema.
  392. schema_name:
  393. Name of the schema to move the table to.
  394. """
  395. self.execute(
  396. self.sql_alter_table_schema
  397. % (self.quote_name(table_name), self.quote_name(schema_name))
  398. )
  399. def alter_model_schema(self, model: Type[Model], schema_name: str) -> None:
  400. """Moves the specified model's table into the specified schema.
  401. WARNING: Moving models into a different schema than the default
  402. will break querying the model.
  403. Arguments:
  404. model:
  405. Model of which to move the table.
  406. schema_name:
  407. Name of the schema to move the model's table to.
  408. """
  409. self.execute(
  410. self.sql_alter_table_schema
  411. % (
  412. self.quote_name(model._meta.db_table),
  413. self.quote_name(schema_name),
  414. )
  415. )
  416. def refresh_materialized_view_model(
  417. self, model: Type[Model], concurrently: bool = False
  418. ) -> None:
  419. """Refreshes a materialized view."""
  420. sql_template = (
  421. self.sql_refresh_materialized_view_concurrently
  422. if concurrently
  423. else self.sql_refresh_materialized_view
  424. )
  425. sql = sql_template % self.quote_name(model._meta.db_table)
  426. self.execute(sql)
  427. def create_view_model(self, model: Type[Model]) -> None:
  428. """Creates a new view model."""
  429. self._create_view_model(self.sql_create_view, model)
  430. def replace_view_model(self, model: Type[Model]) -> None:
  431. """Replaces a view model with a newer version.
  432. This is used to alter the backing query of a view.
  433. """
  434. self._create_view_model(self.sql_replace_view, model)
  435. def delete_view_model(self, model: Type[Model]) -> None:
  436. """Deletes a view model."""
  437. sql = self.sql_drop_view % self.quote_name(model._meta.db_table)
  438. self.execute(sql)
  439. def create_materialized_view_model(self, model: Type[Model]) -> None:
  440. """Creates a new materialized view model."""
  441. self._create_view_model(self.sql_create_materialized_view, model)
  442. def replace_materialized_view_model(self, model: Type[Model]) -> None:
  443. """Replaces a materialized view with a newer version.
  444. This is used to alter the backing query of a materialized view.
  445. Replacing a materialized view is a lot trickier than a normal view.
  446. For normal views we can use `CREATE OR REPLACE VIEW`, but for
  447. materialized views, we have to create the new view, copy all
  448. indexes and constraints and drop the old one.
  449. This operation is atomic as it runs in a transaction.
  450. """
  451. with self.connection.cursor() as cursor:
  452. constraints = self.introspection.get_constraints(
  453. cursor, model._meta.db_table
  454. )
  455. with transaction.atomic():
  456. self.delete_materialized_view_model(model)
  457. self.create_materialized_view_model(model)
  458. for constraint_name, constraint_options in constraints.items():
  459. if not constraint_options["definition"]:
  460. raise SuspiciousOperation(
  461. "Table %s has a constraint '%s' that no definition could be generated for",
  462. (model._meta.db_table, constraint_name),
  463. )
  464. self.execute(constraint_options["definition"])
  465. def delete_materialized_view_model(self, model: Type[Model]) -> None:
  466. """Deletes a materialized view model."""
  467. sql = self.sql_drop_materialized_view % self.quote_name(
  468. model._meta.db_table
  469. )
  470. self.execute(sql)
  471. def create_partitioned_model(self, model: Type[Model]) -> None:
  472. """Creates a new partitioned model."""
  473. meta = self._partitioning_properties_for_model(model)
  474. # get the sql statement that django creates for normal
  475. # table creations..
  476. sql, params = self._extract_sql(self.create_model, model)
  477. partitioning_key_sql = ", ".join(
  478. self.quote_name(field_name) for field_name in meta.key
  479. )
  480. # create a composite key that includes the partitioning key
  481. sql = sql.replace(" PRIMARY KEY", "")
  482. if model._meta.pk and model._meta.pk.name not in meta.key:
  483. sql = sql[:-1] + ", PRIMARY KEY (%s, %s))" % (
  484. self.quote_name(model._meta.pk.name),
  485. partitioning_key_sql,
  486. )
  487. else:
  488. sql = sql[:-1] + ", PRIMARY KEY (%s))" % (partitioning_key_sql,)
  489. # extend the standard CREATE TABLE statement with
  490. # 'PARTITION BY ...'
  491. sql += self.sql_partition_by % (
  492. meta.method.upper(),
  493. partitioning_key_sql,
  494. )
  495. self.execute(sql, params)
  496. def delete_partitioned_model(self, model: Type[Model]) -> None:
  497. """Drops the specified partitioned model."""
  498. return self.delete_model(model)
  499. def add_range_partition(
  500. self,
  501. model: Type[Model],
  502. name: str,
  503. from_values: Any,
  504. to_values: Any,
  505. comment: Optional[str] = None,
  506. ) -> None:
  507. """Creates a new range partition for the specified partitioned model.
  508. Arguments:
  509. model:
  510. Partitioned model to create a partition for.
  511. name:
  512. Name to give to the new partition.
  513. Final name will be "{table_name}_{partition_name}"
  514. from_values:
  515. Start of the partitioning key range of
  516. values that need to be stored in this
  517. partition.
  518. to_values:
  519. End of the partitioning key range of
  520. values that need to be stored in this
  521. partition.
  522. comment:
  523. Optionally, a comment to add on this
  524. partition table.
  525. """
  526. # asserts the model is a model set up for partitioning
  527. self._partitioning_properties_for_model(model)
  528. table_name = self.create_partition_table_name(model, name)
  529. sql = self.sql_add_range_partition % (
  530. self.quote_name(table_name),
  531. self.quote_name(model._meta.db_table),
  532. "%s",
  533. "%s",
  534. )
  535. with transaction.atomic():
  536. self.execute(sql, (from_values, to_values))
  537. if comment:
  538. self.set_comment_on_table(table_name, comment)
  539. def add_list_partition(
  540. self,
  541. model: Type[Model],
  542. name: str,
  543. values: List[Any],
  544. comment: Optional[str] = None,
  545. ) -> None:
  546. """Creates a new list partition for the specified partitioned model.
  547. Arguments:
  548. model:
  549. Partitioned model to create a partition for.
  550. name:
  551. Name to give to the new partition.
  552. Final name will be "{table_name}_{partition_name}"
  553. values:
  554. Partition key values that should be
  555. stored in this partition.
  556. comment:
  557. Optionally, a comment to add on this
  558. partition table.
  559. """
  560. # asserts the model is a model set up for partitioning
  561. self._partitioning_properties_for_model(model)
  562. table_name = self.create_partition_table_name(model, name)
  563. sql = self.sql_add_list_partition % (
  564. self.quote_name(table_name),
  565. self.quote_name(model._meta.db_table),
  566. ",".join(["%s" for _ in range(len(values))]),
  567. )
  568. with transaction.atomic():
  569. self.execute(sql, values)
  570. if comment:
  571. self.set_comment_on_table(table_name, comment)
  572. def add_hash_partition(
  573. self,
  574. model: Type[Model],
  575. name: str,
  576. modulus: int,
  577. remainder: int,
  578. comment: Optional[str] = None,
  579. ) -> None:
  580. """Creates a new hash partition for the specified partitioned model.
  581. Arguments:
  582. model:
  583. Partitioned model to create a partition for.
  584. name:
  585. Name to give to the new partition.
  586. Final name will be "{table_name}_{partition_name}"
  587. modulus:
  588. Integer value by which the key is divided.
  589. remainder:
  590. The remainder of the hash value when divided by modulus.
  591. comment:
  592. Optionally, a comment to add on this partition table.
  593. """
  594. # asserts the model is a model set up for partitioning
  595. self._partitioning_properties_for_model(model)
  596. table_name = self.create_partition_table_name(model, name)
  597. sql = self.sql_add_hash_partition % (
  598. self.quote_name(table_name),
  599. self.quote_name(model._meta.db_table),
  600. "%s",
  601. "%s",
  602. )
  603. with transaction.atomic():
  604. self.execute(sql, (modulus, remainder))
  605. if comment:
  606. self.set_comment_on_table(table_name, comment)
  607. def add_default_partition(
  608. self, model: Type[Model], name: str, comment: Optional[str] = None
  609. ) -> None:
  610. """Creates a new default partition for the specified partitioned model.
  611. A default partition is a partition where rows are routed to when
  612. no more specific partition is a match.
  613. Arguments:
  614. model:
  615. Partitioned model to create a partition for.
  616. name:
  617. Name to give to the new partition.
  618. Final name will be "{table_name}_{partition_name}"
  619. comment:
  620. Optionally, a comment to add on this
  621. partition table.
  622. """
  623. # asserts the model is a model set up for partitioning
  624. self._partitioning_properties_for_model(model)
  625. table_name = self.create_partition_table_name(model, name)
  626. sql = self.sql_add_default_partition % (
  627. self.quote_name(table_name),
  628. self.quote_name(model._meta.db_table),
  629. )
  630. with transaction.atomic():
  631. self.execute(sql)
  632. if comment:
  633. self.set_comment_on_table(table_name, comment)
  634. def delete_partition(self, model: Type[Model], name: str) -> None:
  635. """Deletes the partition with the specified name."""
  636. sql = self.sql_delete_partition % self.quote_name(
  637. self.create_partition_table_name(model, name)
  638. )
  639. self.execute(sql)
  640. def alter_db_table(
  641. self, model: Type[Model], old_db_table: str, new_db_table: str
  642. ) -> None:
  643. """Alters a table/model."""
  644. super().alter_db_table(model, old_db_table, new_db_table)
  645. for side_effect in self.side_effects:
  646. side_effect.alter_db_table(model, old_db_table, new_db_table)
  647. def add_field(self, model: Type[Model], field: Field) -> None:
  648. """Adds a new field to an exisiting model."""
  649. super().add_field(model, field)
  650. for side_effect in self.side_effects:
  651. side_effect.add_field(model, field)
  652. def remove_field(self, model: Type[Model], field: Field) -> None:
  653. """Removes a field from an existing model."""
  654. for side_effect in self.side_effects:
  655. side_effect.remove_field(model, field)
  656. super().remove_field(model, field)
  657. def alter_field(
  658. self,
  659. model: Type[Model],
  660. old_field: Field,
  661. new_field: Field,
  662. strict: bool = False,
  663. ) -> None:
  664. """Alters an existing field on an existing model."""
  665. super().alter_field(model, old_field, new_field, strict)
  666. for side_effect in self.side_effects:
  667. side_effect.alter_field(model, old_field, new_field, strict)
  668. def vacuum_table(
  669. self,
  670. table_name: str,
  671. columns: List[str] = [],
  672. *,
  673. full: bool = False,
  674. freeze: bool = False,
  675. verbose: bool = False,
  676. analyze: bool = False,
  677. disable_page_skipping: bool = False,
  678. skip_locked: bool = False,
  679. index_cleanup: bool = False,
  680. truncate: bool = False,
  681. parallel: Optional[int] = None,
  682. ) -> None:
  683. """Runs the VACUUM statement on the specified table with the specified
  684. options.
  685. Arguments:
  686. table_name:
  687. Name of the table to run VACUUM on.
  688. columns:
  689. Optionally, a list of columns to vacuum. If not
  690. specified, all columns are vacuumed.
  691. """
  692. if self.connection.in_atomic_block:
  693. raise SuspiciousOperation("Vacuum cannot be done in a transaction")
  694. options = []
  695. if full:
  696. options.append("FULL")
  697. if freeze:
  698. options.append("FREEZE")
  699. if verbose:
  700. options.append("VERBOSE")
  701. if analyze:
  702. options.append("ANALYZE")
  703. if disable_page_skipping:
  704. options.append("DISABLE_PAGE_SKIPPING")
  705. if skip_locked:
  706. options.append("SKIP_LOCKED")
  707. if index_cleanup:
  708. options.append("INDEX_CLEANUP")
  709. if truncate:
  710. options.append("TRUNCATE")
  711. if parallel is not None:
  712. options.append(f"PARALLEL {parallel}")
  713. sql = "VACUUM"
  714. if options:
  715. options_sql = ", ".join(options)
  716. sql += f" ({options_sql})"
  717. sql += f" {self.quote_name(table_name)}"
  718. if columns:
  719. columns_sql = ", ".join(
  720. [self.quote_name(column) for column in columns]
  721. )
  722. sql += f" ({columns_sql})"
  723. self.execute(sql)
  724. def vacuum_model(
  725. self, model: Type[Model], fields: List[Field] = [], **kwargs
  726. ) -> None:
  727. """Runs the VACUUM statement on the table of the specified model with
  728. the specified options.
  729. Arguments:
  730. table_name:
  731. model:
  732. Model of which to run VACUUM the table.
  733. fields:
  734. Optionally, a list of fields to vacuum. If not
  735. specified, all fields are vacuumed.
  736. """
  737. columns = [
  738. field.column
  739. for field in fields
  740. if getattr(field, "concrete", False) and field.column
  741. ]
  742. self.vacuum_table(model._meta.db_table, columns, **kwargs)
  743. def set_comment_on_table(self, table_name: str, comment: str) -> None:
  744. """Sets the comment on the specified table."""
  745. sql = self.sql_table_comment % (self.quote_name(table_name), "%s")
  746. self.execute(sql, (comment,))
  747. def _create_view_model(self, sql: str, model: Type[Model]) -> None:
  748. """Creates a new view model using the specified SQL query."""
  749. meta = self._view_properties_for_model(model)
  750. with self.connection.cursor() as cursor:
  751. view_sql = cursor.mogrify(*meta.query)
  752. if isinstance(view_sql, bytes):
  753. view_sql = view_sql.decode("utf-8")
  754. self.execute(sql % (self.quote_name(model._meta.db_table), view_sql))
  755. def _extract_sql(self, method, *args):
  756. """Calls the specified method with the specified arguments and
  757. intercepts the SQL statement it WOULD execute.
  758. We use this to figure out the exact SQL statement Django would
  759. execute. We can then make a small modification and execute it
  760. ourselves.
  761. """
  762. with mock.patch.object(self, "execute") as execute:
  763. method(*args)
  764. return tuple(execute.mock_calls[0])[1]
  765. @staticmethod
  766. def _view_properties_for_model(model: Type[Model]):
  767. """Gets the view options for the specified model.
  768. Raises:
  769. ImproperlyConfigured:
  770. When the specified model is not set up
  771. as a view.
  772. """
  773. meta = getattr(model, "_view_meta", None)
  774. if not meta:
  775. raise ImproperlyConfigured(
  776. (
  777. "Model '%s' is not properly configured to be a view."
  778. " Create the `ViewMeta` class as a child of '%s'."
  779. )
  780. % (model.__name__, model.__name__)
  781. )
  782. if not is_sql_with_params(meta.query):
  783. raise ImproperlyConfigured(
  784. (
  785. "Model '%s' is not properly configured to be a view."
  786. " Set the `query` and `key` attribute on the"
  787. " `ViewMeta` class as a child of '%s'"
  788. )
  789. % (model.__name__, model.__name__)
  790. )
  791. return meta
  792. @staticmethod
  793. def _partitioning_properties_for_model(model: Type[Model]):
  794. """Gets the partitioning options for the specified model.
  795. Raises:
  796. ImproperlyConfigured:
  797. When the specified model is not set up
  798. for partitioning.
  799. """
  800. meta = getattr(model, "_partitioning_meta", None)
  801. if not meta:
  802. raise ImproperlyConfigured(
  803. (
  804. "Model '%s' is not properly configured to be partitioned."
  805. " Create the `PartitioningMeta` class as a child of '%s'."
  806. )
  807. % (model.__name__, model.__name__)
  808. )
  809. if not meta.method or not meta.key:
  810. raise ImproperlyConfigured(
  811. (
  812. "Model '%s' is not properly configured to be partitioned."
  813. " Set the `method` and `key` attributes on the"
  814. " `PartitioningMeta` class as a child of '%s'"
  815. )
  816. % (model.__name__, model.__name__)
  817. )
  818. if meta.method not in PostgresPartitioningMethod:
  819. raise ImproperlyConfigured(
  820. (
  821. "Model '%s' is not properly configured to be partitioned."
  822. " '%s' is not a member of the PostgresPartitioningMethod enum."
  823. )
  824. % (model.__name__, meta.method)
  825. )
  826. if not isinstance(meta.key, list):
  827. raise ImproperlyConfigured(
  828. (
  829. "Model '%s' is not properly configured to be partitioned."
  830. " Partitioning key should be a list (of field names or values,"
  831. " depending on the partitioning method)."
  832. )
  833. % model.__name__
  834. )
  835. try:
  836. for field_name in meta.key:
  837. model._meta.get_field(field_name)
  838. except FieldDoesNotExist:
  839. raise ImproperlyConfigured(
  840. (
  841. "Model '%s' is not properly configured to be partitioned."
  842. " Field '%s' in partitioning key %s is not a valid field on"
  843. " '%s'."
  844. )
  845. % (model.__name__, field_name, meta.key, model.__name__)
  846. )
  847. return meta
  848. def create_partition_table_name(self, model: Type[Model], name: str) -> str:
  849. return "%s_%s" % (model._meta.db_table.lower(), name.lower())
  850. def _clone_model_field(self, field: Field, **overrides) -> Field:
  851. """Clones the specified model field and overrides its kwargs with the
  852. specified overrides.
  853. The cloned field will not be contributed to the model.
  854. """
  855. _, _, field_args, field_kwargs = field.deconstruct()
  856. cloned_field_args = field_args[:]
  857. cloned_field_kwargs = {**field_kwargs, **overrides}
  858. cloned_field = field.__class__(
  859. *cloned_field_args, **cloned_field_kwargs
  860. )
  861. cloned_field.model = field.model
  862. cloned_field.set_attributes_from_name(field.name)
  863. if cloned_field.remote_field and field.remote_field:
  864. cloned_field.remote_field.model = field.remote_field.model
  865. cloned_field.set_attributes_from_rel() # type: ignore[attr-defined]
  866. return cloned_field