1from .base import Operation
  2
  3
  4class SeparateDatabaseAndState(Operation):
  5    """
  6    Take two lists of operations - ones that will be used for the database,
  7    and ones that will be used for the state change. This allows operations
  8    that don't support state change to have it applied, or have operations
  9    that affect the state or not the database, or so on.
 10    """
 11
 12    serialization_expand_args = ["database_operations", "state_operations"]
 13
 14    def __init__(self, database_operations=None, state_operations=None):
 15        self.database_operations = database_operations or []
 16        self.state_operations = state_operations or []
 17
 18    def deconstruct(self):
 19        kwargs = {}
 20        if self.database_operations:
 21            kwargs["database_operations"] = self.database_operations
 22        if self.state_operations:
 23            kwargs["state_operations"] = self.state_operations
 24        return (self.__class__.__qualname__, [], kwargs)
 25
 26    def state_forwards(self, package_label, state):
 27        for state_operation in self.state_operations:
 28            state_operation.state_forwards(package_label, state)
 29
 30    def database_forwards(self, package_label, schema_editor, from_state, to_state):
 31        # We calculate state separately in here since our state functions aren't useful
 32        for database_operation in self.database_operations:
 33            to_state = from_state.clone()
 34            database_operation.state_forwards(package_label, to_state)
 35            database_operation.database_forwards(
 36                package_label, schema_editor, from_state, to_state
 37            )
 38            from_state = to_state
 39
 40    def describe(self):
 41        return "Custom state/database change combination"
 42
 43
 44class RunSQL(Operation):
 45    """
 46    Run some raw SQL.
 47
 48    Also accept a list of operations that represent the state change effected
 49    by this SQL change, in case it's custom column/table creation/deletion.
 50    """
 51
 52    def __init__(self, sql, *, state_operations=None, hints=None, elidable=False):
 53        self.sql = sql
 54        self.state_operations = state_operations or []
 55        self.hints = hints or {}
 56        self.elidable = elidable
 57
 58    def deconstruct(self):
 59        kwargs = {
 60            "sql": self.sql,
 61        }
 62        if self.state_operations:
 63            kwargs["state_operations"] = self.state_operations
 64        if self.hints:
 65            kwargs["hints"] = self.hints
 66        return (self.__class__.__qualname__, [], kwargs)
 67
 68    def state_forwards(self, package_label, state):
 69        for state_operation in self.state_operations:
 70            state_operation.state_forwards(package_label, state)
 71
 72    def database_forwards(self, package_label, schema_editor, from_state, to_state):
 73        self._run_sql(schema_editor, self.sql)
 74
 75    def describe(self):
 76        return "Raw SQL operation"
 77
 78    def _run_sql(self, schema_editor, sqls):
 79        if isinstance(sqls, list | tuple):
 80            for sql in sqls:
 81                params = None
 82                if isinstance(sql, list | tuple):
 83                    elements = len(sql)
 84                    if elements == 2:
 85                        sql, params = sql
 86                    else:
 87                        raise ValueError("Expected a 2-tuple but got %d" % elements)  # noqa: UP031
 88                schema_editor.execute(sql, params=params)
 89        else:
 90            statements = schema_editor.connection.ops.prepare_sql_script(sqls)
 91            for statement in statements:
 92                schema_editor.execute(statement, params=None)
 93
 94
 95class RunPython(Operation):
 96    """
 97    Run Python code in a context suitable for doing versioned ORM operations.
 98    """
 99
100    reduces_to_sql = False
101
102    def __init__(self, code, *, atomic=None, hints=None, elidable=False):
103        self.atomic = atomic
104        # Forwards code
105        if not callable(code):
106            raise ValueError("RunPython must be supplied with a callable")
107        self.code = code
108        self.hints = hints or {}
109        self.elidable = elidable
110
111    def deconstruct(self):
112        kwargs = {
113            "code": self.code,
114        }
115        if self.atomic is not None:
116            kwargs["atomic"] = self.atomic
117        if self.hints:
118            kwargs["hints"] = self.hints
119        return (self.__class__.__qualname__, [], kwargs)
120
121    def state_forwards(self, package_label, state):
122        # RunPython objects have no state effect. To add some, combine this
123        # with SeparateDatabaseAndState.
124        pass
125
126    def database_forwards(self, package_label, schema_editor, from_state, to_state):
127        # RunPython has access to all models. Ensure that all models are
128        # reloaded in case any are delayed.
129        from_state.clear_delayed_models_cache()
130        # We now execute the Python code in a context that contains a 'models'
131        # object, representing the versioned models as an app registry.
132        # We could try to override the global cache, but then people will still
133        # use direct imports, so we go with a documentation approach instead.
134        self.code(from_state.models_registry, schema_editor)
135
136    def describe(self):
137        return "Raw Python operation"