1"""
  2Query subclasses which provide extra functionality beyond simple data retrieval.
  3"""
  4
  5from __future__ import annotations
  6
  7from typing import TYPE_CHECKING, Any
  8
  9from plain.models.exceptions import FieldError
 10from plain.models.expressions import ResolvableExpression
 11from plain.models.sql.constants import CURSOR, GET_ITERATOR_CHUNK_SIZE, NO_RESULTS
 12from plain.models.sql.query import Query
 13
 14if TYPE_CHECKING:
 15    from plain.models.fields import Field
 16
 17__all__ = ["DeleteQuery", "UpdateQuery", "InsertQuery", "AggregateQuery"]
 18
 19
 20class DeleteQuery(Query):
 21    """A DELETE SQL query."""
 22
 23    def do_query(self, table: str, where: Any) -> int:
 24        self.alias_map = {table: self.alias_map[table]}
 25        self.where = where
 26        cursor = self.get_compiler().execute_sql(CURSOR)
 27        if cursor:
 28            with cursor:
 29                return cursor.rowcount
 30        return 0
 31
 32    def delete_batch(self, id_list: list[Any]) -> int:
 33        """
 34        Set up and execute delete queries for all the objects in id_list.
 35
 36        More than one physical query may be executed if there are a
 37        lot of values in id_list.
 38        """
 39        # number of objects deleted
 40        num_deleted = 0
 41        assert self.model is not None, "DELETE requires a model"
 42        meta = self.model._model_meta
 43        field = meta.get_forward_field("id")
 44        for offset in range(0, len(id_list), GET_ITERATOR_CHUNK_SIZE):
 45            self.clear_where()
 46            self.add_filter(
 47                f"{field.attname}__in",
 48                id_list[offset : offset + GET_ITERATOR_CHUNK_SIZE],
 49            )
 50            num_deleted += self.do_query(self.model.model_options.db_table, self.where)
 51        return num_deleted
 52
 53
 54class UpdateQuery(Query):
 55    """An UPDATE SQL query."""
 56
 57    def __init__(self, *args: Any, **kwargs: Any) -> None:
 58        super().__init__(*args, **kwargs)
 59        self._setup_query()
 60
 61    def _setup_query(self) -> None:
 62        """
 63        Run on initialization and at the end of chaining. Any attributes that
 64        would normally be set in __init__() should go here instead.
 65        """
 66        self.values: list[tuple[Any, Any, Any]] = []
 67        self.related_ids: dict[Any, list[Any]] | None = None
 68        self.related_updates: dict[Any, list[tuple[Any, Any, Any]]] = {}
 69
 70    def clone(self) -> UpdateQuery:
 71        obj = super().clone()
 72        obj.related_updates = self.related_updates.copy()
 73        return obj
 74
 75    def update_batch(self, id_list: list[Any], values: dict[str, Any]) -> None:
 76        self.add_update_values(values)
 77        for offset in range(0, len(id_list), GET_ITERATOR_CHUNK_SIZE):
 78            self.clear_where()
 79            self.add_filter(
 80                "id__in", id_list[offset : offset + GET_ITERATOR_CHUNK_SIZE]
 81            )
 82            self.get_compiler().execute_sql(NO_RESULTS)
 83
 84    def add_update_values(self, values: dict[str, Any]) -> None:
 85        """
 86        Convert a dictionary of field name to value mappings into an update
 87        query. This is the entry point for the public update() method on
 88        querysets.
 89        """
 90
 91        assert self.model is not None, "UPDATE requires model metadata"
 92        meta = self.model._model_meta
 93        values_seq = []
 94        for name, val in values.items():
 95            field = meta.get_field(name)
 96            direct = (
 97                not (field.auto_created and not field.concrete) or not field.concrete
 98            )
 99            model = field.model
100            from plain.models.fields.related import ManyToManyField
101
102            if not direct or isinstance(field, ManyToManyField):
103                raise FieldError(
104                    f"Cannot update model field {field!r} (only non-relations and "
105                    "foreign keys permitted)."
106                )
107            if model is not meta.model:
108                self.add_related_update(model, field, val)
109                continue
110            values_seq.append((field, model, val))
111        return self.add_update_fields(values_seq)
112
113    def add_update_fields(self, values_seq: list[tuple[Any, Any, Any]]) -> None:
114        """
115        Append a sequence of (field, model, value) triples to the internal list
116        that will be used to generate the UPDATE query. Might be more usefully
117        called add_update_targets() to hint at the extra information here.
118        """
119        for field, model, val in values_seq:
120            if isinstance(val, ResolvableExpression):
121                # Resolve expressions here so that annotations are no longer needed
122                val = val.resolve_expression(self, allow_joins=False, for_save=True)
123            self.values.append((field, model, val))
124
125    def add_related_update(self, model: Any, field: Any, value: Any) -> None:
126        """
127        Add (name, value) to an update query for an ancestor model.
128
129        Update are coalesced so that only one update query per ancestor is run.
130        """
131        self.related_updates.setdefault(model, []).append((field, None, value))
132
133    def get_related_updates(self) -> list[UpdateQuery]:
134        """
135        Return a list of query objects: one for each update required to an
136        ancestor model. Each query will have the same filtering conditions as
137        the current query but will only update a single table.
138        """
139        if not self.related_updates:
140            return []
141        result = []
142        for model, values in self.related_updates.items():
143            query = UpdateQuery(model)
144            query.values = values
145            if self.related_ids is not None:
146                query.add_filter("id__in", self.related_ids[model])
147            result.append(query)
148        return result
149
150
151class InsertQuery(Query):
152    def __str__(self) -> str:
153        raise NotImplementedError(
154            "InsertQuery does not support __str__(). "
155            "Use get_compiler().as_sql() which returns a list of SQL statements."
156        )
157
158    def sql_with_params(self) -> Any:
159        raise NotImplementedError(
160            "InsertQuery does not support sql_with_params(). "
161            "Use get_compiler().as_sql() which returns a list of SQL statements."
162        )
163
164    def __init__(
165        self,
166        *args: Any,
167        on_conflict: str | None = None,
168        update_fields: list[Field] | None = None,
169        unique_fields: list[Field] | None = None,
170        **kwargs: Any,
171    ) -> None:
172        super().__init__(*args, **kwargs)
173        self.fields: list[Field] = []
174        self.objs: list[Any] = []
175        self.on_conflict = on_conflict
176        self.update_fields: list[Field] = update_fields or []
177        self.unique_fields: list[Field] = unique_fields or []
178
179    def insert_values(
180        self, fields: list[Any], objs: list[Any], raw: bool = False
181    ) -> None:
182        self.fields = fields
183        self.objs = objs
184        self.raw = raw
185
186
187class AggregateQuery(Query):
188    """
189    Take another query as a parameter to the FROM clause and only select the
190    elements in the provided list.
191    """
192
193    def __init__(self, model: Any, inner_query: Any) -> None:
194        self.inner_query = inner_query
195        super().__init__(model)