from __future__ import annotations import types from collections import defaultdict from collections.abc import AsyncIterator, Callable, Collection, Generator, Iterable from copy import copy from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, cast, overload from pypika_tortoise import JoinType, Order, Table from pypika_tortoise.analytics import Count from pypika_tortoise.functions import Cast from pypika_tortoise.queries import QueryBuilder from pypika_tortoise.terms import Case, Field, Star, Term, ValueWrapper from typing_extensions import Literal, Protocol from tortoise.backends.base.client import BaseDBAsyncClient, Capabilities from tortoise.exceptions import ( DoesNotExist, FieldError, IntegrityError, MultipleObjectsReturned, ParamsError, ) from tortoise.expressions import Expression, Q, RawSQL, ResolveContext, ResolveResult from tortoise.fields.relational import ( ForeignKeyFieldInstance, OneToOneFieldInstance, RelationalField, ) from tortoise.filters import FilterInfoDict from tortoise.query_utils import ( Prefetch, QueryModifier, TableCriterionTuple, expand_lookup_expression, get_joins_for_related_field, ) from tortoise.router import router from tortoise.utils import chunk # Empty placeholder - Should never be edited. QUERY: QueryBuilder = QueryBuilder() if TYPE_CHECKING: # pragma: nocoverage from tortoise.models import Model MODEL = TypeVar("MODEL", bound="Model") T_co = TypeVar("T_co", covariant=True) SINGLE = TypeVar("SINGLE", bound=bool) class QuerySetSingle(Protocol[T_co]): """ Awaiting on this will resolve a single instance of the Model object, and not a sequence. """ # pylint: disable=W0104 def __await__(self) -> Generator[Any, None, T_co]: ... # pragma: nocoverage def prefetch_related( self, *args: str | Prefetch ) -> QuerySetSingle[T_co]: ... # pragma: nocoverage def select_related(self, *args: str) -> QuerySetSingle[T_co]: ... # pragma: nocoverage def annotate( self, **kwargs: Expression | Term ) -> QuerySetSingle[T_co]: ... # pragma: nocoverage def only(self, *fields_for_select: str) -> QuerySetSingle[T_co]: ... # pragma: nocoverage def values_list( self, *fields_: str, flat: bool = False ) -> ValuesListQuery[Literal[True]]: ... # pragma: nocoverage def values( self, *args: str, **kwargs: str ) -> ValuesQuery[Literal[True]]: ... # pragma: nocoverage class AwaitableQuery(Generic[MODEL]): __slots__ = ( "query", "model", "_joined_tables", "_db", "capabilities", "_annotations", "_custom_filters", "_q_objects", ) def __init__(self, model: type[MODEL]) -> None: self._joined_tables: list[Table] = [] self.model: type[MODEL] = model self.query: QueryBuilder = QUERY self._db: BaseDBAsyncClient = None # type: ignore self.capabilities: Capabilities = model._meta.db.capabilities self._annotations: dict[str, Expression | Term] = {} self._custom_filters: dict[str, FilterInfoDict] = {} self._q_objects: list[Q] = [] def _choose_db(self, for_write: bool = False) -> BaseDBAsyncClient: """ Return the connection that will be used if this query is executed now. :return: BaseDBAsyncClient: """ if self._db: return self._db if for_write: db = router.db_for_write(self.model) else: db = router.db_for_read(self.model) return db or self.model._meta.db def _choose_db_if_not_chosen(self, for_write: bool = False) -> None: if self._db is None: self._db = self._choose_db(for_write) # type: ignore def resolve_filters(self) -> None: """Builds the common filters for a QuerySet.""" has_aggregate = self._resolve_annotate() modifier = QueryModifier() for node in self._q_objects: modifier &= node.resolve( ResolveContext( model=self.model, table=self.model._meta.basetable, annotations=self._annotations, custom_filters=self._custom_filters, ) ) for join in modifier.joins: if join[0] not in self._joined_tables: self.query = self.query.join(join[0], how=JoinType.left_outer).on(join[1]) self._joined_tables.append(join[0]) self.query._havings = modifier.having_criterion self.query._wheres = modifier.where_criterion if has_aggregate and (self._joined_tables or self.query._havings or self.query._orderbys): self.query = self.query.groupby( *[self.model._meta.basetable[field] for field in self.model._meta.db_fields] ) def _join_table_by_field( self, table: Table, related_field_name: str, related_field: RelationalField ) -> Table: joins = get_joins_for_related_field(table, related_field, related_field_name) for join in joins: self._join_table(join) return joins[-1][0] def _join_table(self, table_criterio_tuple: TableCriterionTuple) -> None: if table_criterio_tuple[0] not in self._joined_tables: self.query = self.query.join(table_criterio_tuple[0], how=JoinType.left_outer).on( table_criterio_tuple[1] ) self._joined_tables.append(table_criterio_tuple[0]) @staticmethod def _resolve_ordering_string(ordering: str, reverse: bool = False) -> tuple[str, Order]: order_type = Order.asc if ordering[0] == "-": field_name = ordering[1:] order_type = Order.desc else: field_name = ordering if reverse: order_type = Order.desc if order_type == Order.asc else Order.asc return field_name, order_type def resolve_ordering( self, model: type[Model], table: Table, orderings: Iterable[tuple[str, str | Order]], annotations: dict[str, Term | Expression], fields_for_select: Collection[str] | None = None, ) -> None: """ Applies standard ordering to QuerySet. :param model: The Model this queryset is based on. :param table: ``pypika_tortoise.Table`` to keep track of the virtual SQL table (to allow self referential joins) :param orderings: What columns/order to order by :param annotations: Annotations that may be ordered on :param fields_for_select: Contains fields that are selected in the SELECT clause if .only(), .values() or .values_list() are used. :raises FieldError: If a field provided does not exist in model. """ # Do not apply default ordering for annotated queries to not mess them up if not orderings and self.model._meta.ordering and not annotations: orderings = self.model._meta.ordering for ordering in orderings: field_name = ordering[0] if field_name in model._meta.fetch_fields: raise FieldError( "Filtering by relation is not possible. Filter by nested field of related model" ) related_field_name, __, forwarded = field_name.partition("__") if related_field_name in model._meta.fetch_fields: related_field = cast(RelationalField, model._meta.fields_map[related_field_name]) related_table = self._join_table_by_field(table, related_field_name, related_field) self.resolve_ordering( related_field.related_model, related_table, [(forwarded, ordering[1])], {}, ) elif field_name in annotations: term: Term if not fields_for_select or field_name in fields_for_select: # The annotation is SELECTed, we can just reference it in the following cases: # - Empty fields_for_select means that all columns and annotations are selected, # hence we can reference the annotation. # - The annotation is in fields_for_select, hence we can reference it. term = Field(field_name) else: # The annotation is not in SELECT, resolve it annotation = annotations[field_name] if isinstance(annotation, Term): term = annotation else: term = annotation.resolve( ResolveContext( model=self.model, table=table, annotations=annotations, custom_filters={}, ) ).term self.query = self.query.orderby(term, order=ordering[1]) else: field_object = model._meta.fields_map.get(field_name) if not field_object: raise FieldError(f"Unknown field {field_name} for model {model.__name__}") field_name = field_object.source_field or field_name field = table[field_name] func = field_object.get_for_dialect( model._meta.db.capabilities.dialect, "function_cast" ) if func: field = func(field_object, field) self.query = self.query.orderby(field, order=ordering[1]) def _resolve_annotate(self) -> bool: if not self._annotations: return False annotation_info: dict[str, ResolveResult] = {} for key, annotation in self._annotations.items(): if isinstance(annotation, Term): annotation_info[key] = ResolveResult(term=annotation) else: annotation_info[key] = annotation.resolve( ResolveContext( model=self.model, table=self.model._meta.basetable, annotations=self._annotations, custom_filters=self._custom_filters, ) ) for key, info in annotation_info.items(): for join in info.joins: self._join_table(join) if key in self._annotations: self.query._select_other(info.term.as_(key)) # type:ignore[arg-type] return any(info.term.is_aggregate for info in annotation_info.values()) def sql(self, params_inline=False) -> str: """ Returns the SQL query that will be executed. By default, it will return the query with placeholders, but if you set `params_inline=True`, it will inline the parameters. :param params_inline: Whether to inline the parameters """ self._choose_db_if_not_chosen() self._make_query() if params_inline: sql = self.query.get_sql() else: sql, _ = self.query.get_parameterized_sql() return sql def _make_query(self) -> None: raise NotImplementedError() # pragma: nocoverage async def _execute(self) -> Any: raise NotImplementedError() # pragma: nocoverage class QuerySet(AwaitableQuery[MODEL]): __slots__ = ( "fields", "_prefetch_map", "_prefetch_queries", "_single", "_raise_does_not_exist", "_db", "_limit", "_offset", "_fields_for_select", "_filter_kwargs", "_orderings", "_distinct", "_having", "_group_bys", "_select_for_update", "_select_for_update_nowait", "_select_for_update_skip_locked", "_select_for_update_of", "_select_for_update_no_key", "_select_related", "_select_related_idx", "_use_indexes", "_force_indexes", ) def __init__(self, model: type[MODEL]) -> None: super().__init__(model) self.fields: set[str] = model._meta.db_fields self._prefetch_map: dict[str, set[str | Prefetch]] = {} self._prefetch_queries: dict[str, list[tuple[str | None, QuerySet]]] = {} self._single: bool = False self._raise_does_not_exist: bool = False self._limit: int | None = None self._offset: int | None = None self._filter_kwargs: dict[str, Any] = {} self._orderings: list[tuple[str, Any]] = [] self._distinct: bool = False self._having: dict[str, Any] = {} self._fields_for_select: tuple[str, ...] = () self._group_bys: tuple[str, ...] = () self._select_for_update: bool = False self._select_for_update_nowait: bool = False self._select_for_update_skip_locked: bool = False self._select_for_update_of: set[str] = set() self._select_for_update_no_key: bool = False self._select_related: set[str] = set() self._select_related_idx: list[ tuple[type[Model], int, Table | str, type[Model], Iterable[str | None]] ] = [] # format with: model,idx,model_name,parent_model self._force_indexes: set[str] = set() self._use_indexes: set[str] = set() def _clone(self) -> QuerySet[MODEL]: queryset = self.__class__.__new__(self.__class__) queryset.fields = self.fields queryset.model = self.model queryset.query = self.query queryset.capabilities = self.capabilities queryset._prefetch_map = copy(self._prefetch_map) queryset._prefetch_queries = copy(self._prefetch_queries) queryset._single = self._single queryset._raise_does_not_exist = self._raise_does_not_exist queryset._db = self._db queryset._limit = self._limit queryset._offset = self._offset queryset._fields_for_select = self._fields_for_select queryset._filter_kwargs = copy(self._filter_kwargs) queryset._orderings = copy(self._orderings) queryset._joined_tables = copy(self._joined_tables) queryset._q_objects = copy(self._q_objects) queryset._distinct = self._distinct queryset._annotations = copy(self._annotations) queryset._having = copy(self._having) queryset._custom_filters = copy(self._custom_filters) queryset._group_bys = copy(self._group_bys) queryset._select_for_update = self._select_for_update queryset._select_for_update_nowait = self._select_for_update_nowait queryset._select_for_update_skip_locked = self._select_for_update_skip_locked queryset._select_for_update_of = self._select_for_update_of queryset._select_for_update_no_key = self._select_for_update_no_key queryset._select_related = self._select_related queryset._select_related_idx = self._select_related_idx queryset._force_indexes = self._force_indexes queryset._use_indexes = self._use_indexes return queryset def _filter_or_exclude(self, *args: Q, negate: bool, **kwargs: Any) -> QuerySet[MODEL]: queryset = self._clone() for arg in args: if not isinstance(arg, Q): raise TypeError("expected Q objects as args") if negate: queryset._q_objects.append(~arg) else: queryset._q_objects.append(arg) for key, value in kwargs.items(): if negate: queryset._q_objects.append(~Q(**{key: value})) else: queryset._q_objects.append(Q(**{key: value})) return queryset def filter(self, *args: Q, **kwargs: Any) -> QuerySet[MODEL]: """ Filters QuerySet by given kwargs. You can filter by related objects like this: .. code-block:: python3 Team.filter(events__tournament__name='Test') You can also pass Q objects to filters as args. """ return self._filter_or_exclude(negate=False, *args, **kwargs) def exclude(self, *args: Q, **kwargs: Any) -> QuerySet[MODEL]: """ Same as .filter(), but with appends all args with NOT """ return self._filter_or_exclude(negate=True, *args, **kwargs) def _parse_orderings( self, orderings: tuple[str, ...], reverse=False ) -> list[tuple[str, Order]]: """ Convert ordering from strings to standard items for queryset. :param orderings: What columns/order to order by :param reverse: Whether reverse order :return: standard ordering for QuerySet. """ new_ordering = [] for ordering in orderings: field_name, order_type = self._resolve_ordering_string(ordering, reverse=reverse) if not ( field_name.split("__")[0] in self.model._meta.fields or field_name in self._annotations ): raise FieldError(f"Unknown field {field_name} for model {self.model.__name__}") new_ordering.append((field_name, order_type)) return new_ordering def order_by(self, *orderings: str) -> QuerySet[MODEL]: """ Accept args to filter by in format like this: .. code-block:: python3 .order_by('name', '-tournament__name') Supports ordering by related models too. A '-' before the name will result in descending sort order, default is ascending. :raises FieldError: If unknown field has been provided. """ queryset = self._clone() queryset._orderings = self._parse_orderings(orderings) return queryset def _as_single(self) -> QuerySetSingle[MODEL | None]: self._single = True self._limit = 1 return cast(QuerySetSingle[Optional[MODEL]], self) def latest(self, *orderings: str) -> QuerySetSingle[MODEL | None]: """ Returns the most recent object by ordering descending on the providers fields. :params orderings: Fields to order by. :raises FieldError: If unknown or no fields has been provided. """ if not orderings: raise FieldError("No fields passed") queryset = self._clone() queryset._orderings = self._parse_orderings(orderings, reverse=True) return queryset._as_single() def earliest(self, *orderings: str) -> QuerySetSingle[MODEL | None]: """ Returns the earliest object by ordering ascending on the specified field. :params orderings: Fields to order by. :raises FieldError: If unknown or no fields has been provided. """ if not orderings: raise FieldError("No fields passed") queryset = self._clone() queryset._orderings = self._parse_orderings(orderings) return queryset._as_single() def limit(self, limit: int) -> QuerySet[MODEL]: """ Limits QuerySet to given length. :raises ParamsError: Limit should be non-negative number. """ if limit < 0: raise ParamsError("Limit should be non-negative number") queryset = self._clone() queryset._limit = limit return queryset def offset(self, offset: int) -> QuerySet[MODEL]: """ Query offset for QuerySet. :raises ParamsError: Offset should be non-negative number. """ if offset < 0: raise ParamsError("Offset should be non-negative number") queryset = self._clone() queryset._offset = offset if self.capabilities.requires_limit and queryset._limit is None: queryset._limit = 1000000 return queryset def __getitem__(self, key: slice) -> QuerySet[MODEL]: """ Query offset and limit for Queryset. :raises ParamsError: QuerySet indices must be slices. :raises ParamsError: Slice steps should be 1 or None. :raises ParamsError: Slice start should be non-negative number or None. :raises ParamsError: Slice stop should be non-negative number greater that slice start, or None. """ if not isinstance(key, slice): raise ParamsError("QuerySet indices must be slices.") if not (key.step is None or (isinstance(key.step, int) and key.step == 1)): raise ParamsError("Slice steps should be 1 or None.") start = key.start if key.start is not None else 0 if not isinstance(start, int) or start < 0: raise ParamsError("Slice start should be non-negative number or None.") if key.stop is not None and (not isinstance(key.stop, int) or key.stop <= start): raise ParamsError( "Slice stop should be non-negative number greater that slice start, or None.", ) queryset = self.offset(start) if key.stop: queryset = queryset.limit(key.stop - start) return queryset def distinct(self) -> QuerySet[MODEL]: """ Make QuerySet distinct. Only makes sense in combination with a ``.values()`` or ``.values_list()`` as it precedes all the fetched fields with a distinct. """ queryset = self._clone() queryset._distinct = True return queryset def select_for_update( self, nowait: bool = False, skip_locked: bool = False, of: tuple[str, ...] = (), no_key: bool = False, ) -> QuerySet[MODEL]: """ Make QuerySet select for update. Returns a queryset that will lock rows until the end of the transaction, generating a SELECT ... FOR UPDATE SQL statement on supported databases. :param nowait: If `True`, raise an error if the lock cannot be obtained immediately. :param skip_locked: If `True`, skip rows that are already locked by other transactions instead of waiting. :param of: Specify the tables to lock when dealing with multiple related tables, e.g. when using `select_related`. Provide a tuple of table names to indicate which tables' rows should be locked. By default, all fetched rows are locked. :param no_key: If `True`, use the lower SELECT ... FOR NO KEY UPDATE lock strength on PostgreSQL to allow creating or deleting rows in other tables that reference the locked rows via foreign keys. The parameter is ignored on other backends. """ if self.capabilities.support_for_update: queryset = self._clone() queryset._select_for_update = True queryset._select_for_update_nowait = nowait queryset._select_for_update_skip_locked = skip_locked queryset._select_for_update_of = set(of) queryset._select_for_update_no_key = ( no_key and self.capabilities.support_for_no_key_update ) return queryset return self def annotate(self, **kwargs: Expression | Term) -> QuerySet[MODEL]: """ Annotate result with aggregation or function result. :raises TypeError: Value of kwarg is expected to be a ``Function`` instance. """ from tortoise.models import get_filters_for_field queryset = self._clone() for key, annotation in kwargs.items(): # if not isinstance(annotation, (Function, Term)): # raise TypeError("value is expected to be Function/Term instance") queryset._annotations[key] = annotation queryset._custom_filters.update(get_filters_for_field(key, None, key)) return queryset def group_by(self, *fields: str) -> QuerySet[MODEL]: """ Make QuerySet returns list of dict or tuple with group by. Must call before .values() or .values_list() """ queryset = self._clone() queryset._group_bys = fields return queryset def values_list(self, *fields_: str, flat: bool = False) -> ValuesListQuery[Literal[False]]: """ Make QuerySet returns list of tuples for given args instead of objects. If call after `.get()`, `.get_or_none()` or `.first()` return tuples for given args instead of object. If ```flat=True`` and only one arg is passed can return flat list or just scalar. If no arguments are passed it will default to a tuple containing all fields in order of declaration. """ if self._fields_for_select: raise ValueError(".values_list() cannot be used with .only()") fields_for_select_list = fields_ or [ field for field in self.model._meta.fields_map if field in self.model._meta.db_fields ] + list(self._annotations.keys()) return ValuesListQuery( db=self._db, model=self.model, q_objects=self._q_objects, single=self._single, raise_does_not_exist=self._raise_does_not_exist, flat=flat, fields_for_select_list=fields_for_select_list, distinct=self._distinct, limit=self._limit, offset=self._offset, orderings=self._orderings, annotations=self._annotations, custom_filters=self._custom_filters, group_bys=self._group_bys, force_indexes=self._force_indexes, use_indexes=self._use_indexes, ) def values(self, *args: str, **kwargs: str) -> ValuesQuery[Literal[False]]: """ Make QuerySet return dicts instead of objects. If called after `.get()`, `.get_or_none()` or `.first()`, returns a dict instead of an object. You can specify which fields to include by: - Passing field names as positional arguments - Using kwargs in the format `field_name='name_in_dict'` to customize the keys in the resulting dict If no arguments are passed, it will default to a dict containing all fields. :raises FieldError: If duplicate key has been provided. """ if self._fields_for_select: raise ValueError(".values() cannot be used with .only()") if args or kwargs: fields_for_select: dict[str, str] = {} for field in args: if field in fields_for_select: raise FieldError(f"Duplicate key {field}") fields_for_select[field] = field for return_as, field in kwargs.items(): if return_as in fields_for_select: raise FieldError(f"Duplicate key {return_as}") fields_for_select[return_as] = field else: _fields = [ field for field in self.model._meta.fields_map.keys() if field in self.model._meta.fields_db_projection.keys() ] + list(self._annotations.keys()) fields_for_select = {field: field for field in _fields} return ValuesQuery( db=self._db, model=self.model, q_objects=self._q_objects, single=self._single, raise_does_not_exist=self._raise_does_not_exist, fields_for_select=fields_for_select, distinct=self._distinct, limit=self._limit, offset=self._offset, orderings=self._orderings, annotations=self._annotations, custom_filters=self._custom_filters, group_bys=self._group_bys, force_indexes=self._force_indexes, use_indexes=self._use_indexes, ) def delete(self) -> DeleteQuery: """ Delete all objects in QuerySet. """ return DeleteQuery( db=self._db, model=self.model, q_objects=self._q_objects, annotations=self._annotations, custom_filters=self._custom_filters, limit=self._limit, orderings=self._orderings, ) def update(self, **kwargs: Any) -> UpdateQuery: """ Update all objects in QuerySet with given kwargs. .. admonition: Example: .. code-block:: py3 await Employee.filter(occupation='developer').update(salary=5000) Will instead of returning a resultset, update the data in the DB itself. """ return UpdateQuery( db=self._db, model=self.model, update_kwargs=kwargs, q_objects=self._q_objects, annotations=self._annotations, custom_filters=self._custom_filters, limit=self._limit, orderings=self._orderings, ) def count(self) -> CountQuery: """ Return count of objects in queryset instead of objects. """ return CountQuery( db=self._db, model=self.model, q_objects=self._q_objects, annotations=self._annotations, custom_filters=self._custom_filters, limit=self._limit, offset=self._offset, force_indexes=self._force_indexes, use_indexes=self._use_indexes, ) def exists(self) -> ExistsQuery: """ Return True/False whether queryset exists. """ return ExistsQuery( db=self._db, model=self.model, q_objects=self._q_objects, annotations=self._annotations, custom_filters=self._custom_filters, force_indexes=self._force_indexes, use_indexes=self._use_indexes, ) def all(self) -> QuerySet[MODEL]: """ Return the whole QuerySet. Essentially a no-op except as the only operation. """ return self._clone() def raw(self, sql: str) -> RawSQLQuery: """ Return the QuerySet from raw SQL """ return RawSQLQuery(model=self.model, db=self._db, sql=sql) def first(self) -> QuerySetSingle[MODEL | None]: """ Limit queryset to one object and return one object instead of list. """ queryset = self._clone() return queryset._as_single() def last(self) -> QuerySetSingle[MODEL | None]: """ Limit queryset to one object and return the last object instead of list. """ queryset = self._clone() if queryset._orderings: new_ordering = [ (field, Order.desc if order_type == Order.asc else Order.asc) for field, order_type in queryset._orderings ] elif pk := self.model._meta.pk: new_ordering = [(pk.model_field_name, Order.desc)] else: raise FieldError( f"QuerySet has no ordering and model {self.model.__name__} has no pk defined" ) queryset._orderings = new_ordering return queryset._as_single() def get(self, *args: Q, **kwargs: Any) -> QuerySetSingle[MODEL]: """ Fetch exactly one object matching the parameters. """ queryset = self.filter(*args, **kwargs) queryset._limit = 2 queryset._single = True queryset._raise_does_not_exist = True return queryset # type: ignore async def in_bulk(self, id_list: Iterable[str | int], field_name: str) -> dict[str, MODEL]: """ Return a dictionary mapping each of the given IDs to the object with that ID. If `id_list` isn't provided, evaluate the entire QuerySet. :param id_list: A list of field values :param field_name: Must be a unique field """ objs = await self.filter(**{f"{field_name}__in": id_list}) return {getattr(obj, field_name): obj for obj in objs} def bulk_create( self, objects: Iterable[MODEL], batch_size: int | None = None, ignore_conflicts: bool = False, update_fields: Iterable[str] | None = None, on_conflict: Iterable[str] | None = None, ) -> BulkCreateQuery[MODEL]: """ This method inserts the provided list of objects into the database in an efficient manner (generally only 1 query, no matter how many objects there are). :param on_conflict: On conflict index name :param update_fields: Update fields when conflicts :param ignore_conflicts: Ignore conflicts when inserting :param objects: List of objects to bulk create :param batch_size: How many objects are created in a single query :raises ValueError: If params do not meet specifications """ if ignore_conflicts and update_fields: raise ValueError( "ignore_conflicts and update_fields are mutually exclusive.", ) if not ignore_conflicts: if (update_fields and not on_conflict) or (on_conflict and not update_fields): raise ValueError("update_fields and on_conflict need set in same time.") return BulkCreateQuery( db=self._db, model=self.model, objects=objects, batch_size=batch_size, ignore_conflicts=ignore_conflicts, update_fields=update_fields, on_conflict=on_conflict, ) def bulk_update( self, objects: Iterable[MODEL], fields: Iterable[str], batch_size: int | None = None, ) -> BulkUpdateQuery[MODEL]: """ Update the given fields in each of the given objects in the database. :param objects: List of objects to bulk create :param fields: The fields to update :param batch_size: How many objects are created in a single query :raises ValueError: If objects have no primary key set """ if any(obj.pk is None for obj in objects): raise ValueError("All bulk_update() objects must have a primary key set.") return BulkUpdateQuery( db=self._db, model=self.model, q_objects=self._q_objects, annotations=self._annotations, custom_filters=self._custom_filters, limit=self._limit, orderings=self._orderings, objects=objects, fields=fields, batch_size=batch_size, ) def get_or_none(self, *args: Q, **kwargs: Any) -> QuerySetSingle[MODEL | None]: """ Fetch exactly one object matching the parameters. """ queryset = self.filter(*args, **kwargs) queryset._limit = 2 queryset._single = True return queryset # type: ignore def only(self, *fields_for_select: str) -> QuerySet[MODEL]: """ Fetch ONLY the specified fields to create a partial model. Persisting changes on the model is allowed only when: * All the fields you want to update is specified in ``.save(update_fields=[...])`` * You included the Model primary key in the `.only(...)`` To protect against common mistakes we ensure that errors get raised: * If you access a field that is not specified, you will get an ``AttributeError``. * If you do a ``.save()`` a ``IncompleteInstanceError`` will be raised as the model is, as requested, incomplete. * If you do a ``.save(update_fields=[...])`` and you didn't include the primary key in the ``.only(...)``, then ``IncompleteInstanceError`` will be raised indicating that updates can't be done without the primary key being known. * If you do a ``.save(update_fields=[...])`` and one of the fields in ``update_fields`` was not in the ``.only(...)``, then ``IncompleteInstanceError`` as that field is not available to be updated. """ if not fields_for_select: raise ValueError(".only() requires at least one field") queryset = self._clone() queryset._fields_for_select = fields_for_select return queryset def select_related(self, *fields: str) -> QuerySet[MODEL]: """ Return a new QuerySet instance that will select related objects. If fields are specified, they must be ForeignKey fields and only those related objects are included in the selection. """ queryset = self._clone() for field in fields: queryset._select_related.add(field) return queryset def force_index(self, *index_names: str) -> QuerySet[MODEL]: """ The FORCE INDEX hint acts like USE INDEX (index_list), with the addition that a table scan is assumed to be very expensive. """ if self.capabilities.support_index_hint: queryset = self._clone() for index_name in index_names: queryset._force_indexes.add(index_name) return queryset return self def use_index(self, *index_names: str) -> QuerySet[MODEL]: """ The USE INDEX (index_list) hint tells MySQL to use only one of the named indexes to find rows in the table. """ if self.capabilities.support_index_hint: queryset = self._clone() for index_name in index_names: queryset._use_indexes.add(index_name) return queryset return self def prefetch_related(self, *args: str | Prefetch) -> QuerySet[MODEL]: """ Like ``.fetch_related()`` on instance, but works on all objects in QuerySet. :raises FieldError: If the field to prefetch on is not a relation, or not found. """ queryset = self._clone() queryset._prefetch_map = {} for relation in args: if isinstance(relation, Prefetch): relation.resolve_for_queryset(queryset) continue first_level_field, __, forwarded_prefetch = relation.partition("__") if first_level_field not in self.model._meta.fetch_fields: if first_level_field in self.model._meta.fields: raise FieldError( f"Field {first_level_field} on {self.model._meta.full_name} is not a relation" ) raise FieldError( f"Relation {first_level_field} for {self.model._meta.full_name} not found" ) if first_level_field not in queryset._prefetch_map.keys(): queryset._prefetch_map[first_level_field] = set() if forwarded_prefetch: queryset._prefetch_map[first_level_field].add(forwarded_prefetch) return queryset async def explain(self) -> Any: """Fetch and return information about the query execution plan. This is done by executing an ``EXPLAIN`` query whose exact prefix depends on the database backend, as documented below. - PostgreSQL: ``EXPLAIN (FORMAT JSON, VERBOSE) ...`` - SQLite: ``EXPLAIN QUERY PLAN ...`` - MySQL: ``EXPLAIN FORMAT=JSON ...`` .. note:: This is only meant to be used in an interactive environment for debugging and query optimization. **The output format may (and will) vary greatly depending on the database backend.** """ self._choose_db_if_not_chosen() self._make_query() return await self._db.executor_class(model=self.model, db=self._db).execute_explain( self.query.get_sql() ) def using_db(self, _db: BaseDBAsyncClient | None) -> QuerySet[MODEL]: """ Executes query in provided db client. Useful for transactions workaround. """ queryset = self._clone() queryset._db = _db if _db else queryset._db return queryset def _join_select_related(self, lookup_expression: str) -> tuple[type[Model], Table]: fields = expand_lookup_expression(self.model, lookup_expression) model = self.model table = self.model._meta.basetable path: tuple[str | None, ...] = (None,) for field in fields: field = cast(RelationalField, field) path = path + (field.model_field_name,) table = self._join_table_by_field(table, field.model_field_name, field) # do not select related fields if we are only selecting a subset of fields if self._fields_for_select: continue related_fields = field.related_model._meta.db_fields append_item = ( field.related_model, len(related_fields), field.model_field_name, model, path, ) model = field.related_model if append_item not in self._select_related_idx: self._select_related_idx.append(append_item) self.query = self.query.select( *[ table[related_field].as_(f"{table.get_table_name()}.{related_field}") for related_field in related_fields ] ) return model, table def _resolve_only(self, only_lookup_expressions: tuple[str, ...]) -> None: # Group fields by fetch fields, e.g. ["a__b", "a__c"] -> {"a": ["b", "c"]}. # The direct fields of the model are the ones that would have the key "". fetch_to_fields = defaultdict(list) # the order is important here, we need to process the shallowest fields first # because we want to populate _select_related_idx with actual items that need to be # selected, not "filler" items tha just tell the executor that an empty instance has # to be created for expression in sorted(only_lookup_expressions, key=lambda x: x.count("__")): fetch_fields_lookup, __, field_name = expression.rpartition("__") fetch_to_fields[fetch_fields_lookup].append(field_name) # select direct model fields which would have the key "": {"": ["a", "b"]} data_fields = fetch_to_fields.pop("", None) if data_fields: table = self.model._meta.basetable self._select_related_idx.append( ( self.model, len(data_fields), table, self.model, (None,), ) ) try: self.query = self.query.select( *[ table[self.model._meta.fields_db_projection[field]].as_(field) for field in data_fields if field not in self._annotations ] ) except KeyError as e: raise FieldError( f'Unknown field "{e.args[0]}" for model "{self.model.__name__}"' ) from e else: # even though no data fields are selected, we need to let the executor know # that an empty instance of the model has to be created self._select_related_idx.append( ( self.model, 0, self.model._meta.basetable, self.model, (None,), ) ) # Select fields of related models, e.g. {"a": ["b", "c"]} added_paths = set() for fetch_fields_lookup, data_fields in fetch_to_fields.items(): fetch_fields = expand_lookup_expression(self.model, fetch_fields_lookup) referring_model = model = self.model table = self.model._meta.basetable path: tuple[str | None, ...] = (None,) for i, fetch_field in enumerate(fetch_fields): field = cast(RelationalField, fetch_field) path = path + (field.model_field_name,) table = self._join_table_by_field(table, field.model_field_name, field) referring_model = model model = field.related_model if path in added_paths: continue self._select_related_idx.append( ( model, # we need 0 items for letting know the executor that instances need to # be created even though no their fields are selected, e.g. # .only("a__b__field") len(data_fields) if i == len(fetch_fields) - 1 else 0, table, referring_model, path, ) ) added_paths.add(path) self.query = self.query.select( *[table[field].as_(f"{table.get_table_name()}.{field}") for field in data_fields] ) def _make_query(self) -> None: # clean tmp records first self._select_related_idx = [] self._joined_tables = [] table = self.model._meta.basetable if self._fields_for_select: # select .only() fields self.query = self.model._meta.basequery.select() self._resolve_only(self._fields_for_select) else: # select all fields self.query = copy(self.model._meta.basequery_all_fields) # type:ignore[assignment] append_item = ( self.model, len(self.model._meta.db_fields) + len(self._annotations), table, self.model, (None,), ) self._select_related_idx.append(append_item) self.resolve_ordering( self.model, self.model._meta.basetable, self._orderings, self._annotations, self._fields_for_select, ) self.resolve_filters() if self._limit is not None: self.query._limit = self.query._wrapper_cls(self._limit) if self._offset is not None: self.query._offset = self.query._wrapper_cls(self._offset) if self._distinct: self.query._distinct = True if self._select_for_update: self.query = self.query.for_update( self._select_for_update_nowait, self._select_for_update_skip_locked, self._select_for_update_of, self._select_for_update_no_key, ) if self._select_related: for select_related in self._select_related: self._join_select_related(select_related) if self._force_indexes: self.query._force_indexes = [] self.query = self.query.force_index(*self._force_indexes) if self._use_indexes: self.query._use_indexes = [] self.query = self.query.use_index(*self._use_indexes) def __await__(self) -> Generator[Any, None, list[MODEL]]: if self._db is None: self._db = self._choose_db(self._select_for_update) # type: ignore self._make_query() return self._execute().__await__() async def __aiter__(self) -> AsyncIterator[MODEL]: for val in await self: yield val async def _execute(self) -> list[MODEL]: instance_list = await self._db.executor_class( model=self.model, db=self._db, prefetch_map=self._prefetch_map, prefetch_queries=self._prefetch_queries, select_related_idx=self._select_related_idx, # type: ignore ).execute_select( *self.query.get_parameterized_sql(), custom_fields=list(self._annotations.keys()), ) if self._single: if len(instance_list) == 1: return instance_list[0] if not instance_list: if self._raise_does_not_exist: raise DoesNotExist(self.model) return None # type: ignore raise MultipleObjectsReturned(self.model) return instance_list class UpdateQuery(AwaitableQuery): __slots__ = ( "update_kwargs", "_orderings", "_limit", "values", ) def __init__( self, model: type[MODEL], update_kwargs: dict[str, Any], db: BaseDBAsyncClient, q_objects: list[Q], annotations: dict[str, Any], custom_filters: dict[str, FilterInfoDict], limit: int | None, orderings: list[tuple[str, str]], ) -> None: super().__init__(model) self.update_kwargs = update_kwargs self._q_objects = q_objects self._annotations = annotations self._custom_filters = custom_filters self._db = db self._limit = limit self._orderings = orderings def _make_query(self) -> None: table = self.model._meta.basetable self.query = self._db.query_class.update(table) if self.capabilities.support_update_limit_order_by and self._limit: self.query._limit = self.query._wrapper_cls(self._limit) self.resolve_ordering(self.model, table, self._orderings, self._annotations) self.resolve_filters() for key, value in self.update_kwargs.items(): field_object = self.model._meta.fields_map.get(key) if not field_object: raise FieldError(f"Unknown keyword argument {key} for model {self.model}") if field_object.pk: raise IntegrityError(f"Field {key} is PK and can not be updated") if isinstance(field_object, (ForeignKeyFieldInstance, OneToOneFieldInstance)): self.model._validate_relation_type(key, value) fk_field: str = field_object.source_field # type: ignore db_field = self.model._meta.fields_map[fk_field].source_field value = self.model._meta.fields_map[fk_field].to_db_value( getattr(value, field_object.to_field_instance.model_field_name), None, ) else: try: db_field = self.model._meta.fields_db_projection[key] except KeyError: raise FieldError(f"Field {key} is virtual and can not be updated") if isinstance(value, Expression): value = value.resolve( ResolveContext( model=self.model, table=table, annotations=self._annotations, custom_filters=self._custom_filters, ) ).term else: value = self.model._meta.fields_map[key].to_db_value(value, None) self.query = self.query.set(db_field, value) def __await__(self) -> Generator[Any, None, int]: self._choose_db_if_not_chosen(True) self._make_query() return self._execute().__await__() async def _execute(self) -> int: return (await self._db.execute_query(*self.query.get_parameterized_sql()))[0] class DeleteQuery(AwaitableQuery): __slots__ = ( "_annotations", "_custom_filters", "_orderings", "_limit", ) def __init__( self, model: type[MODEL], db: BaseDBAsyncClient, q_objects: list[Q], annotations: dict[str, Any], custom_filters: dict[str, FilterInfoDict], limit: int | None, orderings: list[tuple[str, str]], ) -> None: super().__init__(model) self._q_objects = q_objects self._annotations = annotations self._custom_filters = custom_filters self._db = db self._limit = limit self._orderings = orderings def _make_query(self) -> None: self.query = copy(self.model._meta.basequery) if self.capabilities.support_update_limit_order_by and self._limit: self.query._limit = self.query._wrapper_cls(self._limit) self.resolve_ordering( model=self.model, table=self.model._meta.basetable, orderings=self._orderings, annotations=self._annotations, ) self.resolve_filters() self.query._delete_from = True return def __await__(self) -> Generator[Any, None, int]: self._choose_db_if_not_chosen(True) self._make_query() return self._execute().__await__() async def _execute(self) -> int: return (await self._db.execute_query(*self.query.get_parameterized_sql()))[0] class ExistsQuery(AwaitableQuery): __slots__ = ( "_force_indexes", "_use_indexes", ) def __init__( self, model: type[MODEL], db: BaseDBAsyncClient, q_objects: list[Q], annotations: dict[str, Any], custom_filters: dict[str, FilterInfoDict], force_indexes: set[str], use_indexes: set[str], ) -> None: super().__init__(model) self._q_objects = q_objects self._db = db self._annotations = annotations self._custom_filters = custom_filters self._force_indexes = force_indexes self._use_indexes = use_indexes def _make_query(self) -> None: self.query = copy(self.model._meta.basequery) self.resolve_filters() self.query._limit = self.query._wrapper_cls(1) self.query._select_other(ValueWrapper(1, allow_parametrize=False)) # type:ignore[arg-type] if self._force_indexes: self.query._force_indexes = [] self.query = self.query.force_index(*self._force_indexes) if self._use_indexes: self.query._use_indexes = [] self.query = self.query.use_index(*self._use_indexes) def __await__(self) -> Generator[Any, None, bool]: self._choose_db_if_not_chosen() self._make_query() return self._execute().__await__() async def _execute( self, ) -> bool: result, _ = await self._db.execute_query(*self.query.get_parameterized_sql()) return bool(result) class CountQuery(AwaitableQuery): __slots__ = ( "_limit", "_offset", "_force_indexes", "_use_indexes", ) def __init__( self, model: type[MODEL], db: BaseDBAsyncClient, q_objects: list[Q], annotations: dict[str, Any], custom_filters: dict[str, FilterInfoDict], limit: int | None, offset: int | None, force_indexes: set[str], use_indexes: set[str], ) -> None: super().__init__(model) self._q_objects = q_objects self._annotations = annotations self._custom_filters = custom_filters self._limit = limit self._offset = offset or 0 self._db = db self._force_indexes = force_indexes self._use_indexes = use_indexes def _make_query(self) -> None: self.query = copy(self.model._meta.basequery) self.resolve_filters() count_term = Count(Star()) if self.query._groupbys: count_term = count_term.over() # remove annotations self.query._selects = [] self.query._select_other(count_term) if self._force_indexes: self.query._force_indexes = [] self.query = self.query.force_index(*self._force_indexes) if self._use_indexes: self.query._use_indexes = [] self.query = self.query.use_index(*self._use_indexes) def __await__(self) -> Generator[Any, None, int]: self._choose_db_if_not_chosen() self._make_query() return self._execute().__await__() async def _execute(self) -> int: _, result = await self._db.execute_query(*self.query.get_parameterized_sql()) if not result: return 0 count = list(dict(result[0]).values())[0] - self._offset if self._limit and count > self._limit: return self._limit return count class FieldSelectQuery(AwaitableQuery): # pylint: disable=W0223 def __init__(self, model: type[MODEL], annotations: dict[str, Any]) -> None: super().__init__(model) self._annotations = annotations def _join_table_with_forwarded_fields( self, model: type[MODEL], table: Table, field: str, forwarded_fields: str ) -> tuple[Table, str]: if field in model._meta.fields_db_projection and not forwarded_fields: return table, model._meta.fields_db_projection[field] if field in model._meta.fields_db_projection and forwarded_fields: raise FieldError(f'Field "{field}" for model "{model.__name__}" is not relation') if field in self.model._meta.fetch_fields and not forwarded_fields: raise ValueError( f'Selecting relation "{field}" is not possible, select concrete ' "field on related model" ) field_object = cast(RelationalField, model._meta.fields_map.get(field)) if not field_object: raise FieldError(f'Unknown field "{field}" for model "{model.__name__}"') table = self._join_table_by_field(table, field, field_object) field, __, forwarded_fields_ = forwarded_fields.partition("__") return self._join_table_with_forwarded_fields( model=field_object.related_model, table=table, field=field, forwarded_fields=forwarded_fields_, ) def add_field_to_select_query(self, field: str, return_as: str) -> None: table = self.model._meta.basetable if field in self._annotations: self._annotations[return_as] = self._annotations[field] return if field in self.model._meta.fields_db_projection: db_field = self.model._meta.fields_db_projection[field] self.query._select_field(table[db_field].as_(return_as)) return if field in self.model._meta.fetch_fields: raise ValueError( f'Selecting relation "{field}" is not possible, select ' "concrete field on related model" ) field_, __, forwarded_fields = field.partition("__") if field_ in self.model._meta.fetch_fields: related_table, related_db_field = self._join_table_with_forwarded_fields( model=self.model, table=table, field=field_, forwarded_fields=forwarded_fields, ) self.query._select_field(related_table[related_db_field].as_(return_as)) return raise FieldError(f'Unknown field "{field}" for model "{self.model.__name__}"') def resolve_to_python_value(self, model: type[MODEL], field: str) -> Callable: if field in model._meta.fetch_fields: # return as is to get whole model objects return lambda x: x if field in (x[1] for x in model._meta.db_native_fields): return lambda x: x if field in self._annotations: annotation = self._annotations[field] field_object = getattr(annotation, "field_object", None) if field_object: return field_object.to_python_value return lambda x: x if field in model._meta.fields_map: return model._meta.fields_map[field].to_python_value field_, __, forwarded_fields = field.partition("__") if field_ in model._meta.fetch_fields: new_model = model._meta.fields_map[field_].related_model # type: ignore return self.resolve_to_python_value(new_model, forwarded_fields) raise FieldError(f'Unknown field "{field}" for model "{model}"') def _resolve_group_bys(self, *field_names: str) -> list: group_bys = [] for field_name in field_names: if field_name in self._annotations: group_bys.append(Term(field_name)) continue field, __, forwarded_fields = field_name.partition("__") related_table, related_db_field = self._join_table_with_forwarded_fields( model=self.model, table=self.model._meta.basetable, field=field, forwarded_fields=forwarded_fields, ) field = related_table[related_db_field].as_( f"{related_table.get_table_name()}__{field_name}" ) group_bys.append(field) return group_bys class ValuesListQuery(FieldSelectQuery, Generic[SINGLE]): __slots__ = ( "fields", "_limit", "_offset", "_distinct", "_orderings", "_single", "_raise_does_not_exist", "_fields_for_select_list", "_flat", "_group_bys", "_force_indexes", "_use_indexes", ) def __init__( self, model: type[MODEL], db: BaseDBAsyncClient, q_objects: list[Q], single: bool, raise_does_not_exist: bool, fields_for_select_list: tuple[str, ...] | list[str], limit: int | None, offset: int | None, distinct: bool, orderings: list[tuple[str, str]], flat: bool, annotations: dict[str, Any], custom_filters: dict[str, FilterInfoDict], group_bys: tuple[str, ...], force_indexes: set[str], use_indexes: set[str], ) -> None: super().__init__(model, annotations) if flat and (len(fields_for_select_list) != 1): raise TypeError("You can flat value_list only if contains one field") fields_for_select = {str(i): field for i, field in enumerate(fields_for_select_list)} self.fields = fields_for_select self._limit = limit self._offset = offset self._distinct = distinct self._orderings = orderings self._custom_filters = custom_filters self._q_objects = q_objects self._single = single self._raise_does_not_exist = raise_does_not_exist self._fields_for_select_list = fields_for_select_list self._flat = flat self._db = db self._group_bys = group_bys self._force_indexes = force_indexes self._use_indexes = use_indexes def _make_query(self) -> None: self._joined_tables = [] self.query = copy(self.model._meta.basequery) for positional_number, field in self.fields.items(): self.add_field_to_select_query(field, positional_number) self.resolve_ordering( model=self.model, table=self.model._meta.basetable, orderings=self._orderings, annotations=self._annotations, fields_for_select=self._fields_for_select_list, ) self.resolve_filters() if self._limit: self.query._limit = self.query._wrapper_cls(self._limit) if self._offset: self.query._offset = self.query._wrapper_cls(self._offset) if self._distinct: self.query._distinct = True if self._group_bys: self.query._groupbys = self._resolve_group_bys(*self._group_bys) if self._force_indexes: self.query._force_indexes = [] self.query = self.query.force_index(*self._force_indexes) if self._use_indexes: self.query._use_indexes = [] self.query = self.query.use_index(*self._use_indexes) @overload def __await__( self: ValuesListQuery[Literal[False]], ) -> Generator[Any, None, list[tuple[Any, ...]]]: ... @overload def __await__( self: ValuesListQuery[Literal[True]], ) -> Generator[Any, None, tuple[Any, ...]]: ... def __await__(self) -> Generator[Any, None, list[Any] | tuple[Any, ...]]: self._choose_db_if_not_chosen() self._make_query() return self._execute().__await__() # pylint: disable=E1101 async def __aiter__(self: ValuesListQuery[Any]) -> AsyncIterator[Any]: for val in await self: yield val async def _execute(self) -> list[Any] | tuple: _, result = await self._db.execute_query(*self.query.get_parameterized_sql()) columns = [ (key, self.resolve_to_python_value(self.model, name)) for key, name in self.fields.items() ] if self._flat: func = columns[0][1] flatmap = lambda entry: func(entry["0"]) # noqa lst_values = list(map(flatmap, result)) else: listmap = lambda entry: tuple(func(entry[column]) for column, func in columns) # noqa lst_values = list(map(listmap, result)) if self._single: if len(lst_values) == 1: return lst_values[0] if not lst_values: if self._raise_does_not_exist: raise DoesNotExist(self.model) return None # type: ignore raise MultipleObjectsReturned(self.model) return lst_values class ValuesQuery(FieldSelectQuery, Generic[SINGLE]): __slots__ = ( "_fields_for_select", "_limit", "_offset", "_distinct", "_orderings", "_single", "_raise_does_not_exist", "_group_bys", "_force_indexes", "_use_indexes", ) def __init__( self, model: type[MODEL], db: BaseDBAsyncClient, q_objects: list[Q], single: bool, raise_does_not_exist: bool, fields_for_select: dict[str, str], limit: int | None, offset: int | None, distinct: bool, orderings: list[tuple[str, str]], annotations: dict[str, Any], custom_filters: dict[str, FilterInfoDict], group_bys: tuple[str, ...], force_indexes: set[str], use_indexes: set[str], ) -> None: super().__init__(model, annotations) self._fields_for_select = fields_for_select self._limit = limit self._offset = offset self._distinct = distinct self._orderings = orderings self._custom_filters = custom_filters self._q_objects = q_objects self._single = single self._raise_does_not_exist = raise_does_not_exist self._db = db self._group_bys = group_bys self._force_indexes = force_indexes self._use_indexes = use_indexes def _make_query(self) -> None: self._joined_tables = [] self.query = copy(self.model._meta.basequery) for return_as, field in self._fields_for_select.items(): self.add_field_to_select_query(field, return_as) self.resolve_ordering( model=self.model, table=self.model._meta.basetable, orderings=self._orderings, annotations=self._annotations, fields_for_select=self._fields_for_select.keys(), ) self.resolve_filters() # remove annotations that are not in fields_for_select self.query._selects = [ select for select in self.query._selects if select.alias in self._fields_for_select ] if self._limit: self.query._limit = self.query._wrapper_cls(self._limit) if self._offset: self.query._offset = self.query._wrapper_cls(self._offset) if self._distinct: self.query._distinct = True if self._group_bys: self.query._groupbys = self._resolve_group_bys(*self._group_bys) if self._force_indexes: self.query._force_indexes = [] self.query = self.query.force_index(*self._force_indexes) if self._use_indexes: self.query._use_indexes = [] self.query = self.query.use_index(*self._use_indexes) @overload def __await__( self: ValuesQuery[Literal[False]], ) -> Generator[Any, None, list[dict[str, Any]]]: ... @overload def __await__( self: ValuesQuery[Literal[True]], ) -> Generator[Any, None, dict[str, Any]]: ... def __await__( self, ) -> Generator[Any, None, list[dict[str, Any]] | dict[str, Any]]: self._choose_db_if_not_chosen() self._make_query() return self._execute().__await__() # pylint: disable=E1101 async def __aiter__(self: ValuesQuery[Any]) -> AsyncIterator[dict[str, Any]]: for val in await self: yield val async def _execute(self) -> list[dict] | dict: result = await self._db.execute_query_dict(*self.query.get_parameterized_sql()) columns = [ val for val in [ (alias, self.resolve_to_python_value(self.model, field_name)) for alias, field_name in self._fields_for_select.items() ] if not isinstance(val[1], types.LambdaType) ] if columns: for row in result: for col, func in columns: row[col] = func(row[col]) if self._single: if len(result) == 1: return result[0] if not result: if self._raise_does_not_exist: raise DoesNotExist(self.model) return None # type: ignore raise MultipleObjectsReturned(self.model) return result class RawSQLQuery(AwaitableQuery): __slots__ = ("_sql", "_db") def __init__(self, model: type[MODEL], db: BaseDBAsyncClient, sql: str) -> None: super().__init__(model) self._sql = sql self._db = db async def _execute(self) -> Any: instance_list = await self._db.executor_class( model=self.model, db=self._db, ).execute_select(RawSQL(self._sql).get_sql(self._db.query_class.SQL_CONTEXT), []) return instance_list def __await__(self) -> Generator[Any, None, list[MODEL]]: self._choose_db_if_not_chosen() return self._execute().__await__() class BulkUpdateQuery(UpdateQuery, Generic[MODEL]): __slots__ = ("fields", "_objects", "_batch_size", "_queries") def __init__( self, model: type[MODEL], db: BaseDBAsyncClient, q_objects: list[Q], annotations: dict[str, Any], custom_filters: dict[str, FilterInfoDict], limit: int | None, orderings: list[tuple[str, str]], objects: Iterable[MODEL], fields: Iterable[str], batch_size: int | None = None, ): super().__init__( model, update_kwargs={}, db=db, q_objects=q_objects, annotations=annotations, custom_filters=custom_filters, limit=limit, orderings=orderings, ) self.fields = fields self._objects = objects self._batch_size = batch_size self._queries: list[QueryBuilder] = [] def _make_queries(self) -> list[tuple[str, list[Any]]]: table = self.model._meta.basetable self.query = self._db.query_class.update(table) if self.capabilities.support_update_limit_order_by and self._limit: self.query._limit = self.query._wrapper_cls(self._limit) self.resolve_ordering( model=self.model, table=table, orderings=self._orderings, annotations=self._annotations, ) self.resolve_filters() pk_attr = self.model._meta.pk_attr source_pk_attr = self.model._meta.fields_map[pk_attr].source_field or pk_attr pk = Field(source_pk_attr) for objects_item in chunk(self._objects, self._batch_size): query = copy(self.query) for field in self.fields: case = Case() pk_list = [] for obj in objects_item: pk_value = self.model._meta.fields_map[pk_attr].to_db_value(obj.pk, None) field_obj = obj._meta.fields_map[field] field_value = field_obj.to_db_value(getattr(obj, field), obj) case.when( pk == pk_value, ( Cast( self.query._wrapper_cls(field_value), field_obj.get_for_dialect( self._db.schema_generator.DIALECT, "SQL_TYPE" ), ) if self._db.schema_generator.DIALECT == "postgres" else self.query._wrapper_cls(field_value) ), ) pk_list.append(pk_value) query = query.set(field, case) query = query.where(pk.isin(pk_list)) self._queries.append(query) return [query.get_parameterized_sql() for query in self._queries] async def _execute_many(self, queries_with_params: list[tuple[str, list[Any]]]) -> int: count = 0 for sql, values in queries_with_params: count += (await self._db.execute_query(sql, values))[0] return count def __await__(self) -> Generator[Any, Any, int]: self._choose_db_if_not_chosen(True) queries = self._make_queries() return self._execute_many(queries).__await__() def sql(self, params_inline=False) -> str: self._choose_db_if_not_chosen() queries = self._make_queries() return ";".join([sql for sql, _ in queries]) class BulkCreateQuery(AwaitableQuery, Generic[MODEL]): __slots__ = ( "_objects", "_ignore_conflicts", "_batch_size", "_db", "_executor", "_update_fields", "_on_conflict", ) def __init__( self, model: type[MODEL], db: BaseDBAsyncClient, objects: Iterable[MODEL], batch_size: int | None = None, ignore_conflicts: bool = False, update_fields: Iterable[str] | None = None, on_conflict: Iterable[str] | None = None, ): super().__init__(model) self._objects = objects self._ignore_conflicts = ignore_conflicts self._batch_size = batch_size self._db = db self._update_fields = update_fields self._on_conflict = on_conflict def _make_queries(self) -> tuple[str, str]: self._executor = self._db.executor_class(model=self.model, db=self._db) if self._ignore_conflicts or self._update_fields: _, columns = self._executor._prepare_insert_columns() insert_query = self._executor._prepare_insert_statement( columns, ignore_conflicts=self._ignore_conflicts ) insert_query_all = insert_query if self.model._meta.generated_db_fields: _, columns_all = self._executor._prepare_insert_columns(include_generated=True) insert_query_all = self._executor._prepare_insert_statement( columns_all, has_generated=False, ignore_conflicts=self._ignore_conflicts, ) if self._update_fields: alias = f"new_{self.model._meta.db_table}" insert_query_all = insert_query_all.as_(alias).on_conflict( *(self._on_conflict or []) ) insert_query = insert_query.as_(alias).on_conflict(*(self._on_conflict or [])) for update_field in self._update_fields: insert_query_all = insert_query_all.do_update(update_field) insert_query = insert_query.do_update(update_field) return insert_query.get_sql(), insert_query_all.get_sql() else: return self._executor.insert_query, self._executor.insert_query_all async def _execute_many(self, insert_sql: str, insert_sql_all: str) -> None: fields_map = self.model._meta.fields_map for instance_chunk in chunk(self._objects, self._batch_size): values_lists_all = [] values_lists = [] for instance in instance_chunk: if instance._custom_generated_pk: values_lists_all.append( [ fields_map[field_name].to_db_value( getattr(instance, field_name), instance ) for field_name in self._executor.regular_columns_all ] ) else: values_lists.append( [ fields_map[field_name].to_db_value( getattr(instance, field_name), instance ) for field_name in self._executor.regular_columns ] ) if values_lists_all: await self._db.execute_many(insert_sql_all, values_lists_all) if values_lists: await self._db.execute_many(insert_sql, values_lists) def __await__(self) -> Generator[Any, None, None]: self._choose_db_if_not_chosen(True) insert_sql, insert_sql_all = self._make_queries() return self._execute_many(insert_sql, insert_sql_all).__await__() def sql(self, params_inline=False) -> str: self._choose_db_if_not_chosen() insert_sql, insert_sql_all = self._make_queries() if all(o._custom_generated_pk for o in self._objects): return insert_sql_all if all(not o._custom_generated_pk for o in self._objects): return insert_sql return ";".join([insert_sql, insert_sql_all])