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, elidable=False):
 53        self.sql = sql
 54        self.state_operations = state_operations or []
 55        self.elidable = elidable
 56
 57    def deconstruct(self):
 58        kwargs = {
 59            "sql": self.sql,
 60        }
 61        if self.state_operations:
 62            kwargs["state_operations"] = self.state_operations
 63        return (self.__class__.__qualname__, [], kwargs)
 64
 65    def state_forwards(self, package_label, state):
 66        for state_operation in self.state_operations:
 67            state_operation.state_forwards(package_label, state)
 68
 69    def database_forwards(self, package_label, schema_editor, from_state, to_state):
 70        self._run_sql(schema_editor, self.sql)
 71
 72    def describe(self):
 73        return "Raw SQL operation"
 74
 75    def _run_sql(self, schema_editor, sqls):
 76        if isinstance(sqls, list | tuple):
 77            for sql in sqls:
 78                params = None
 79                if isinstance(sql, list | tuple):
 80                    elements = len(sql)
 81                    if elements == 2:
 82                        sql, params = sql
 83                    else:
 84                        raise ValueError("Expected a 2-tuple but got %d" % elements)  # noqa: UP031
 85                schema_editor.execute(sql, params=params)
 86        else:
 87            statements = schema_editor.connection.ops.prepare_sql_script(sqls)
 88            for statement in statements:
 89                schema_editor.execute(statement, params=None)
 90
 91
 92class RunPython(Operation):
 93    """
 94    Run Python code in a context suitable for doing versioned ORM operations.
 95    """
 96
 97    reduces_to_sql = False
 98
 99    def __init__(self, code, *, atomic=None, elidable=False):
100        self.atomic = atomic
101        # Forwards code
102        if not callable(code):
103            raise ValueError("RunPython must be supplied with a callable")
104        self.code = code
105        self.elidable = elidable
106
107    def deconstruct(self):
108        kwargs = {
109            "code": self.code,
110        }
111        if self.atomic is not None:
112            kwargs["atomic"] = self.atomic
113        return (self.__class__.__qualname__, [], kwargs)
114
115    def state_forwards(self, package_label, state):
116        # RunPython objects have no state effect. To add some, combine this
117        # with SeparateDatabaseAndState.
118        pass
119
120    def database_forwards(self, package_label, schema_editor, from_state, to_state):
121        # RunPython has access to all models. Ensure that all models are
122        # reloaded in case any are delayed.
123        from_state.clear_delayed_models_cache()
124        # We now execute the Python code in a context that contains a 'models'
125        # object, representing the versioned models as an app registry.
126        # We could try to override the global cache, but then people will still
127        # use direct imports, so we go with a documentation approach instead.
128        self.code(from_state.models_registry, schema_editor)
129
130    def describe(self):
131        return "Raw Python operation"