1from plain.models.db import router
2
3from .base import Operation
4
5
6class SeparateDatabaseAndState(Operation):
7 """
8 Take two lists of operations - ones that will be used for the database,
9 and ones that will be used for the state change. This allows operations
10 that don't support state change to have it applied, or have operations
11 that affect the state or not the database, or so on.
12 """
13
14 serialization_expand_args = ["database_operations", "state_operations"]
15
16 def __init__(self, database_operations=None, state_operations=None):
17 self.database_operations = database_operations or []
18 self.state_operations = state_operations or []
19
20 def deconstruct(self):
21 kwargs = {}
22 if self.database_operations:
23 kwargs["database_operations"] = self.database_operations
24 if self.state_operations:
25 kwargs["state_operations"] = self.state_operations
26 return (self.__class__.__qualname__, [], kwargs)
27
28 def state_forwards(self, package_label, state):
29 for state_operation in self.state_operations:
30 state_operation.state_forwards(package_label, state)
31
32 def database_forwards(self, package_label, schema_editor, from_state, to_state):
33 # We calculate state separately in here since our state functions aren't useful
34 for database_operation in self.database_operations:
35 to_state = from_state.clone()
36 database_operation.state_forwards(package_label, to_state)
37 database_operation.database_forwards(
38 package_label, schema_editor, from_state, to_state
39 )
40 from_state = to_state
41
42 def database_backwards(self, package_label, schema_editor, from_state, to_state):
43 # We calculate state separately in here since our state functions aren't useful
44 to_states = {}
45 for dbop in self.database_operations:
46 to_states[dbop] = to_state
47 to_state = to_state.clone()
48 dbop.state_forwards(package_label, to_state)
49 # to_state now has the states of all the database_operations applied
50 # which is the from_state for the backwards migration of the last
51 # operation.
52 for database_operation in reversed(self.database_operations):
53 from_state = to_state
54 to_state = to_states[database_operation]
55 database_operation.database_backwards(
56 package_label, schema_editor, from_state, to_state
57 )
58
59 def describe(self):
60 return "Custom state/database change combination"
61
62
63class RunSQL(Operation):
64 """
65 Run some raw SQL. A reverse SQL statement may be provided.
66
67 Also accept a list of operations that represent the state change effected
68 by this SQL change, in case it's custom column/table creation/deletion.
69 """
70
71 noop = ""
72
73 def __init__(
74 self, sql, reverse_sql=None, state_operations=None, hints=None, elidable=False
75 ):
76 self.sql = sql
77 self.reverse_sql = reverse_sql
78 self.state_operations = state_operations or []
79 self.hints = hints or {}
80 self.elidable = elidable
81
82 def deconstruct(self):
83 kwargs = {
84 "sql": self.sql,
85 }
86 if self.reverse_sql is not None:
87 kwargs["reverse_sql"] = self.reverse_sql
88 if self.state_operations:
89 kwargs["state_operations"] = self.state_operations
90 if self.hints:
91 kwargs["hints"] = self.hints
92 return (self.__class__.__qualname__, [], kwargs)
93
94 @property
95 def reversible(self):
96 return self.reverse_sql is not None
97
98 def state_forwards(self, package_label, state):
99 for state_operation in self.state_operations:
100 state_operation.state_forwards(package_label, state)
101
102 def database_forwards(self, package_label, schema_editor, from_state, to_state):
103 if router.allow_migrate(
104 schema_editor.connection.alias, package_label, **self.hints
105 ):
106 self._run_sql(schema_editor, self.sql)
107
108 def database_backwards(self, package_label, schema_editor, from_state, to_state):
109 if self.reverse_sql is None:
110 raise NotImplementedError("You cannot reverse this operation")
111 if router.allow_migrate(
112 schema_editor.connection.alias, package_label, **self.hints
113 ):
114 self._run_sql(schema_editor, self.reverse_sql)
115
116 def describe(self):
117 return "Raw SQL operation"
118
119 def _run_sql(self, schema_editor, sqls):
120 if isinstance(sqls, list | tuple):
121 for sql in sqls:
122 params = None
123 if isinstance(sql, list | tuple):
124 elements = len(sql)
125 if elements == 2:
126 sql, params = sql
127 else:
128 raise ValueError("Expected a 2-tuple but got %d" % elements)
129 schema_editor.execute(sql, params=params)
130 elif sqls != RunSQL.noop:
131 statements = schema_editor.connection.ops.prepare_sql_script(sqls)
132 for statement in statements:
133 schema_editor.execute(statement, params=None)
134
135
136class RunPython(Operation):
137 """
138 Run Python code in a context suitable for doing versioned ORM operations.
139 """
140
141 reduces_to_sql = False
142
143 def __init__(
144 self, code, reverse_code=None, atomic=None, hints=None, elidable=False
145 ):
146 self.atomic = atomic
147 # Forwards code
148 if not callable(code):
149 raise ValueError("RunPython must be supplied with a callable")
150 self.code = code
151 # Reverse code
152 if reverse_code is None:
153 self.reverse_code = None
154 else:
155 if not callable(reverse_code):
156 raise ValueError("RunPython must be supplied with callable arguments")
157 self.reverse_code = reverse_code
158 self.hints = hints or {}
159 self.elidable = elidable
160
161 def deconstruct(self):
162 kwargs = {
163 "code": self.code,
164 }
165 if self.reverse_code is not None:
166 kwargs["reverse_code"] = self.reverse_code
167 if self.atomic is not None:
168 kwargs["atomic"] = self.atomic
169 if self.hints:
170 kwargs["hints"] = self.hints
171 return (self.__class__.__qualname__, [], kwargs)
172
173 @property
174 def reversible(self):
175 return self.reverse_code is not None
176
177 def state_forwards(self, package_label, state):
178 # RunPython objects have no state effect. To add some, combine this
179 # with SeparateDatabaseAndState.
180 pass
181
182 def database_forwards(self, package_label, schema_editor, from_state, to_state):
183 # RunPython has access to all models. Ensure that all models are
184 # reloaded in case any are delayed.
185 from_state.clear_delayed_packages_cache()
186 if router.allow_migrate(
187 schema_editor.connection.alias, package_label, **self.hints
188 ):
189 # We now execute the Python code in a context that contains a 'models'
190 # object, representing the versioned models as an app registry.
191 # We could try to override the global cache, but then people will still
192 # use direct imports, so we go with a documentation approach instead.
193 self.code(from_state.packages, schema_editor)
194
195 def database_backwards(self, package_label, schema_editor, from_state, to_state):
196 if self.reverse_code is None:
197 raise NotImplementedError("You cannot reverse this operation")
198 if router.allow_migrate(
199 schema_editor.connection.alias, package_label, **self.hints
200 ):
201 self.reverse_code(from_state.packages, schema_editor)
202
203 def describe(self):
204 return "Raw Python operation"
205
206 @staticmethod
207 def noop(packages, schema_editor):
208 return None