Skip to content

Commit 99b2c32

Browse files
quertenmontPhotonios
authored andcommitted
Gracefully handle the primary key already being included in the partioning key
1 parent 557e94b commit 99b2c32

File tree

2 files changed

+40
-4
lines changed

2 files changed

+40
-4
lines changed

psqlextra/backend/schema.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -171,10 +171,13 @@ def create_partitioned_model(self, model: Model) -> None:
171171

172172
# create a composite key that includes the partitioning key
173173
sql = sql.replace(" PRIMARY KEY", "")
174-
sql = sql[:-1] + ", PRIMARY KEY (%s, %s))" % (
175-
self.quote_name(model._meta.pk.name),
176-
partitioning_key_sql,
177-
)
174+
if model._meta.pk.name not in meta.key:
175+
sql = sql[:-1] + ", PRIMARY KEY (%s, %s))" % (
176+
self.quote_name(model._meta.pk.name),
177+
partitioning_key_sql,
178+
)
179+
else:
180+
sql = sql[:-1] + ", PRIMARY KEY (%s))" % (partitioning_key_sql,)
178181

179182
# extend the standard CREATE TABLE statement with
180183
# 'PARTITION BY ...'

tests/test_schema_editor_partitioning.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,39 @@ def test_schema_editor_create_delete_partitioned_model_list():
7676
assert len(partitions) == 0
7777

7878

79+
@pytest.mark.postgres_version(lt=110000)
80+
@pytest.mark.parametrize("key", [["name"], ["id", "name"]])
81+
def test_schema_editor_create_delete_partitioned_model_hash(key):
82+
"""Tests whether creating a partitioned model and adding a hash partition
83+
to it using the :see:PostgresSchemaEditor works."""
84+
85+
method = PostgresPartitioningMethod.HASH
86+
87+
model = define_fake_partitioned_model(
88+
{"name": models.TextField()},
89+
{"method": method, "key": key},
90+
)
91+
92+
schema_editor = PostgresSchemaEditor(connection)
93+
schema_editor.create_partitioned_model(model)
94+
95+
schema_editor.add_hash_partition(model, "pt1", modulus=1, remainder=0)
96+
97+
table = db_introspection.get_partitioned_table(model._meta.db_table)
98+
assert table.name == model._meta.db_table
99+
assert table.method == method
100+
assert table.key == key
101+
assert table.partitions[0].full_name == model._meta.db_table + "_pt1"
102+
103+
schema_editor.delete_partitioned_model(model)
104+
105+
table = db_introspection.get_partitioned_table(model._meta.db_table)
106+
assert not table
107+
108+
partitions = db_introspection.get_partitions(model._meta.db_table)
109+
assert len(partitions) == 0
110+
111+
79112
@pytest.mark.postgres_version(lt=110000)
80113
def test_schema_editor_create_delete_partitioned_model_default():
81114
"""Tests whether creating a partitioned model and adding a default

0 commit comments

Comments
 (0)