1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
|
from typing import Dict, List, Type
from psqlextra.models import PostgresPartitionedModel
from .model import PostgresModelState
class PostgresPartitionState:
"""Represents the state of a partition for a :see:PostgresPartitionedModel
during a migration."""
def __init__(self, app_label: str, model_name: str, name: str) -> None:
self.app_label = app_label
self.model_name = model_name
self.name = name
class PostgresRangePartitionState(PostgresPartitionState):
"""Represents the state of a range partition for a
:see:PostgresPartitionedModel during a migration."""
def __init__(
self, app_label: str, model_name: str, name: str, from_values, to_values
):
super().__init__(app_label, model_name, name)
self.from_values = from_values
self.to_values = to_values
class PostgresListPartitionState(PostgresPartitionState):
"""Represents the state of a list partition for a
:see:PostgresPartitionedModel during a migration."""
def __init__(self, app_label: str, model_name: str, name: str, values):
super().__init__(app_label, model_name, name)
self.values = values
class PostgresHashPartitionState(PostgresPartitionState):
"""Represents the state of a hash partition for a
:see:PostgresPartitionedModel during a migration."""
def __init__(
self,
app_label: str,
model_name: str,
name: str,
modulus: int,
remainder: int,
):
super().__init__(app_label, model_name, name)
self.modulus = modulus
self.remainder = remainder
class PostgresPartitionedModelState(PostgresModelState):
"""Represents the state of a :see:PostgresPartitionedModel in the
migrations."""
def __init__(
self,
*args,
partitions: List[PostgresPartitionState] = [],
partitioning_options={},
**kwargs
):
"""Initializes a new instance of :see:PostgresPartitionedModelState.
Arguments:
partitioning_options:
Dictionary of options for partitioning.
See: PostgresPartitionedModelMeta for a list.
"""
super().__init__(*args, **kwargs)
self.partitions: Dict[str, PostgresPartitionState] = {
partition.name: partition for partition in partitions
}
self.partitioning_options = dict(partitioning_options)
def add_partition(self, partition: PostgresPartitionState):
"""Adds a partition to this partitioned model state."""
self.partitions[partition.name] = partition
def delete_partition(self, name: str):
"""Deletes a partition from this partitioned model state."""
del self.partitions[name]
@classmethod
def _pre_new( # type: ignore[override]
cls,
model: PostgresPartitionedModel,
model_state: "PostgresPartitionedModelState",
) -> "PostgresPartitionedModelState":
"""Called when a new model state is created from the specified
model."""
model_state.partitions = dict()
model_state.partitioning_options = dict(
model._partitioning_meta.original_attrs
)
return model_state
def _pre_clone( # type: ignore[override]
self, model_state: "PostgresPartitionedModelState"
) -> "PostgresPartitionedModelState":
"""Called when this model state is cloned."""
model_state.partitions = dict(self.partitions)
model_state.partitioning_options = dict(self.partitioning_options)
return model_state
def _pre_render(self, name: str, bases, attributes):
"""Called when this model state is rendered into a model."""
partitioning_meta = type(
"PartitioningMeta", (), dict(self.partitioning_options)
)
return (
name,
bases,
{**attributes, "PartitioningMeta": partitioning_meta},
)
@classmethod
def _get_base_model_class(self) -> Type[PostgresPartitionedModel]:
"""Gets the class to use as a base class for rendered models."""
return PostgresPartitionedModel
|
