Skip to content

Commit 3f2486c

Browse files
PhotoniosPhotonios
committed
Add vacuum methods to schema editor
1 parent 498124a commit 3f2486c

File tree

2 files changed

+234
-0
lines changed

2 files changed

+234
-0
lines changed

‎psqlextra/backend/schema.py

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -777,6 +777,93 @@ def alter_field(
777777
for side_effect in self.side_effects:
778778
side_effect.alter_field(model, old_field, new_field, strict)
779779

780+
def vacuum_table(
781+
self,
782+
table_name: str,
783+
columns: List[str] = [],
784+
*,
785+
full: bool = False,
786+
freeze: bool = False,
787+
verbose: bool = False,
788+
analyze: bool = False,
789+
disable_page_skipping: bool = False,
790+
skip_locked: bool = False,
791+
index_cleanup: bool = False,
792+
truncate: bool = False,
793+
parallel: Optional[int] = None,
794+
) -> None:
795+
"""Runs the VACUUM statement on the specified table with the specified
796+
options.
797+
798+
Arguments:
799+
table_name:
800+
Name of the table to run VACUUM on.
801+
802+
columns:
803+
Optionally, a list of columns to vacuum. If not
804+
specified, all columns are vacuumed.
805+
"""
806+
807+
if self.connection.in_atomic_block:
808+
raise SuspiciousOperation("Vacuum cannot be done in a transaction")
809+
810+
options = []
811+
if full:
812+
options.append("FULL")
813+
if freeze:
814+
options.append("FREEZE")
815+
if verbose:
816+
options.append("VERBOSE")
817+
if analyze:
818+
options.append("ANALYZE")
819+
if disable_page_skipping:
820+
options.append("DISABLE_PAGE_SKIPPING")
821+
if skip_locked:
822+
options.append("SKIP_LOCKED")
823+
if index_cleanup:
824+
options.append("INDEX_CLEANUP")
825+
if truncate:
826+
options.append("TRUNCATE")
827+
if parallel is not None:
828+
options.append(f"PARALLEL {parallel}")
829+
830+
sql = "VACUUM"
831+
832+
if options:
833+
options_sql = ", ".join(options)
834+
sql += f" ({options_sql})"
835+
836+
sql += f" {self.quote_name(table_name)}"
837+
838+
if columns:
839+
columns_sql = ", ".join(
840+
[self.quote_name(column) for column in columns]
841+
)
842+
sql += f" ({columns_sql})"
843+
844+
self.execute(sql)
845+
846+
def vacuum_model(
847+
self, model: Type[Model], fields: List[Field] = [], **kwargs
848+
) -> None:
849+
"""Runs the VACUUM statement on the table of the specified model with
850+
the specified options.
851+
852+
Arguments:
853+
table_name:
854+
model:
855+
Model of which to run VACUUM the table.
856+
857+
fields:
858+
Optionally, a list of fields to vacuum. If not
859+
specified, all fields are vacuumed.
860+
"""
861+
862+
columns = [
863+
field.column for field in fields if field.concrete and field.column
864+
]
865+
self.vacuum_table(model._meta.db_table, columns, **kwargs)
866+
780867
def set_comment_on_table(self, table_name: str, comment: str) -> None:
781868
"""Sets the comment on the specified table."""
782869

‎tests/test_schema_editor_vacuum.py

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
import pytest
2+
3+
from django.core.exceptions import SuspiciousOperation
4+
from django.db import connection, models
5+
from django.test.utils import CaptureQueriesContext
6+
7+
from psqlextra.backend.schema import PostgresSchemaEditor
8+
9+
from .fake_model import delete_fake_model, get_fake_model
10+
11+
12+
@pytest.fixture
13+
def fake_model():
14+
model = get_fake_model(
15+
{
16+
"name": models.TextField(),
17+
}
18+
)
19+
20+
yield model
21+
22+
delete_fake_model(model)
23+
24+
25+
@pytest.fixture
26+
def fake_model_non_concrete_field(fake_model):
27+
model = get_fake_model(
28+
{
29+
"fk": models.ForeignKey(
30+
fake_model, on_delete=models.CASCADE, related_name="fakes"
31+
),
32+
}
33+
)
34+
35+
yield model
36+
37+
delete_fake_model(model)
38+
39+
40+
def test_schema_editor_vacuum_not_in_transaction(fake_model):
41+
schema_editor = PostgresSchemaEditor(connection)
42+
43+
with pytest.raises(SuspiciousOperation):
44+
schema_editor.vacuum_table(fake_model._meta.db_table)
45+
46+
47+
@pytest.mark.parametrize(
48+
"kwargs,query",
49+
[
50+
(dict(), "VACUUM %s"),
51+
(dict(full=True), "VACUUM (FULL) %s"),
52+
(dict(analyze=True), "VACUUM (ANALYZE) %s"),
53+
(dict(parallel=8), "VACUUM (PARALLEL 8) %s"),
54+
(dict(analyze=True, verbose=True), "VACUUM (VERBOSE, ANALYZE) %s"),
55+
(
56+
dict(analyze=True, parallel=8, verbose=True),
57+
"VACUUM (VERBOSE, ANALYZE, PARALLEL 8) %s",
58+
),
59+
(dict(freeze=True), "VACUUM (FREEZE) %s"),
60+
(dict(verbose=True), "VACUUM (VERBOSE) %s"),
61+
(dict(disable_page_skipping=True), "VACUUM (DISABLE_PAGE_SKIPPING) %s"),
62+
(dict(skip_locked=True), "VACUUM (SKIP_LOCKED) %s"),
63+
(dict(index_cleanup=True), "VACUUM (INDEX_CLEANUP) %s"),
64+
(dict(truncate=True), "VACUUM (TRUNCATE) %s"),
65+
],
66+
)
67+
@pytest.mark.django_db(transaction=True)
68+
def test_schema_editor_vacuum_table(fake_model, kwargs, query):
69+
schema_editor = PostgresSchemaEditor(connection)
70+
71+
with CaptureQueriesContext(connection) as ctx:
72+
schema_editor.vacuum_table(fake_model._meta.db_table, **kwargs)
73+
74+
queries = [query["sql"] for query in ctx.captured_queries]
75+
assert queries == [
76+
query % connection.ops.quote_name(fake_model._meta.db_table)
77+
]
78+
79+
80+
@pytest.mark.django_db(transaction=True)
81+
def test_schema_editor_vacuum_table_columns(fake_model):
82+
schema_editor = PostgresSchemaEditor(connection)
83+
84+
with CaptureQueriesContext(connection) as ctx:
85+
schema_editor.vacuum_table(
86+
fake_model._meta.db_table, ["id", "name"], analyze=True
87+
)
88+
89+
queries = [query["sql"] for query in ctx.captured_queries]
90+
assert queries == [
91+
'VACUUM (ANALYZE) %s ("id", "name")'
92+
% connection.ops.quote_name(fake_model._meta.db_table)
93+
]
94+
95+
96+
@pytest.mark.django_db(transaction=True)
97+
def test_schema_editor_vacuum_model(fake_model):
98+
schema_editor = PostgresSchemaEditor(connection)
99+
100+
with CaptureQueriesContext(connection) as ctx:
101+
schema_editor.vacuum_model(fake_model, analyze=True, parallel=8)
102+
103+
queries = [query["sql"] for query in ctx.captured_queries]
104+
assert queries == [
105+
"VACUUM (ANALYZE, PARALLEL 8) %s"
106+
% connection.ops.quote_name(fake_model._meta.db_table)
107+
]
108+
109+
110+
@pytest.mark.django_db(transaction=True)
111+
def test_schema_editor_vacuum_model_fields(fake_model):
112+
schema_editor = PostgresSchemaEditor(connection)
113+
114+
with CaptureQueriesContext(connection) as ctx:
115+
schema_editor.vacuum_model(
116+
fake_model,
117+
[fake_model._meta.get_field("name")],
118+
analyze=True,
119+
parallel=8,
120+
)
121+
122+
queries = [query["sql"] for query in ctx.captured_queries]
123+
assert queries == [
124+
'VACUUM (ANALYZE, PARALLEL 8) %s ("name")'
125+
% connection.ops.quote_name(fake_model._meta.db_table)
126+
]
127+
128+
129+
@pytest.mark.django_db(transaction=True)
130+
def test_schema_editor_vacuum_model_non_concrete_fields(
131+
fake_model, fake_model_non_concrete_field
132+
):
133+
schema_editor = PostgresSchemaEditor(connection)
134+
135+
with CaptureQueriesContext(connection) as ctx:
136+
schema_editor.vacuum_model(
137+
fake_model,
138+
[fake_model._meta.get_field("fakes")],
139+
analyze=True,
140+
parallel=8,
141+
)
142+
143+
queries = [query["sql"] for query in ctx.captured_queries]
144+
assert queries == [
145+
"VACUUM (ANALYZE, PARALLEL 8) %s"
146+
% connection.ops.quote_name(fake_model._meta.db_table)
147+
]

0 commit comments

Comments
 (0)