QuerySet class¶

async def build_where_clause(
    self, _: Any = None, tables_and_models: tables_and_models_type | None = None
) -> Any:
    """
    (This method is now a simple forwarder to the Compiler.
     It's kept for API compatibility, e.g. for QuerySet(QuerySet) filters)
    """
    compiler = QueryCompiler(self)
    joins: Any | None = None
    if tables_and_models is None:
        joins, tables_and_models = self._get_join_graph_data()

    return await compiler.build_where_clause(tables_and_models, joins=joins)

def update_embed_parent(
    self, embed_parent: tuple[str, str] | None
) -> QuerySet[EdgyModel, EdgyEmbedTarget] | QuerySet[EdgyModel, EdgyModel]:
    """
    Update or remove (provide None) embed_parent applied on instances.
    Note: this doesn't affect embed_parent for filters.

    Args:
        embed_parent: define the new embed_parent.
    Returns:
        QuerySetType: A new QuerySet instance with the new embedding.
    """
    self_queryset = cast("QuerySet[EdgyModel, EdgyEmbedTarget]", self)
    if embed_parent is not None:
        if len(cast(tuple, embed_parent)) != 2:
            raise ValueError("Provided tuple has not exactly two elements.")
        if not embed_parent[0]:
            raise ValueError("First argument in `embed_parent` tuple can't be empty.")
    queryset = self_queryset._clone()
    queryset._embed_parent = embed_parent
    select_pathes: set[str] = set()
    if queryset._embed_parent:
        # can be cross db, so only add the first expanded part
        crawl_result = crawl_relationship(
            queryset.model_class,
            queryset._embed_parent[0],
            model_database=queryset.database,
            traverse_last=True,
        )
        select_pathes.add(crawl_result.forward_path)

    if (
        queryset._update_related_weak(
            select_pathes,
            cache_name="_select_related_embedding",
            clear=True,
        )
        and queryset._prefetch_related
    ):
        # regenerate prefetch pathes

        queryset._update_related_weak(
            chain(
                *(
                    prefetch._generate_select_related_pathes(queryset)
                    for prefetch in queryset._prefetch_related
                )
            ),
            cache_name="_select_related_embedding",
            clear=False,
        )
    return queryset

async def raw_delete(
    self, *, use_models: bool = False, remove_referenced_call: str | bool = False
) -> int | None:
    """
    Internal delete method.

    Exposes remove_referenced_call and doesn't raise an extra signal.

    Delegates to QueryExecutor.delete.
    """
    # We must create new executors *every time* because the queryset
    # state might have changed (e.g., in _model_based_delete)
    executor = QueryExecutor(self)

    return await executor.delete(
        use_models=use_models,
        remove_referenced_call=remove_referenced_call,
        injected_filters=_injected_filters_deletion.get(),
    )

def using(
    self,
    _positional: Any = _sentinel,
    *,
    database: str | Any | None | Database = Undefined,
    schema: str | Any | None | Literal[False] = Undefined,
) -> QuerySet[EdgyModel, EdgyEmbedTarget]:
    """
    Enables and switches the database schema and/or database connection for
    the queryset.

    This method creates a new QuerySet instance that operates within the
    specified database and/or schema context. It's designed to support
    multi-tenancy by allowing dynamic selection of the database or schema
    without modifying the original QuerySet.

    Args:
        _positional (Any): Deprecated positional argument for `schema`. This
                           argument is maintained for backward compatibility but
                           its use is discouraged. It will be treated as the
                           `schema` argument.
        database (str | Database | None): Specifies the database to use.
            Use an extra connection name (`str`), a `Database` instance, `None`
            (for the default registry database), or keep it `Undefined` to retain
            the current queryset database.
        schema (str | None | Literal[False]): Specifies the database schema to use.
            Use a schema name (`str`), `False` (unset schema and use model default),
            `None` (no specific schema), or keep it `Undefined` to retain
            the current queryset schema.

    Returns:
        QuerySet: A new QuerySet instance configured with the specified database
                  and/or schema settings. This new instance allows chaining
                  operations within the new context.

    Warnings:
        DeprecationWarning: If positional arguments are passed to this method for
                            `_positional`. Users should explicitly use `schema=`
                            keyword argument instead.
    """
    # Check for deprecated positional argument usage.
    if _positional is not _sentinel:
        warnings.warn(
            "Passing positional arguments to using is deprecated. Use schema= instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        schema = _positional
        # If schema was not explicitly set as Undefined (meaning it was passed
        # positionally) and is still Undefined, set it to False to unset the schema.
        if schema is Undefined:
            schema = False

    # Create a clone of the current queryset to avoid modifying the original.
    queryset = cast("QuerySet[EdgyModel, EdgyEmbedTarget]", self)._clone()

    # Process the 'database' argument.
    if database is not Undefined:
        registry = cast("Registry", queryset.model_class.meta.registry)
        if isinstance(database, Database):
            connection: Database = database
        elif database is None:
            # Use the default database from the model's registry.
            connection = registry.database
        else:
            # Assert that the database name exists in the extra connections.
            assert database in registry.extra, (
                f"`{database}` is not in the connections extra of the model"
                f"`{queryset.model_class.__name__}` registry"
            )
            # Retrieve the database from extra connections.
            connection = registry.extra[database]
        # Assign the determined connection to the new queryset's database.
        queryset.database = connection

    # Process the 'schema' argument.
    if schema is not Undefined:
        # If schema is False, set using_schema to Undefined to unset it;
        # otherwise, use the provided schema.
        queryset.using_schema = schema if schema is not False else Undefined
        # Get the new schema based on the current queryset's configuration.
        new_schema = queryset.get_schema()
        # If the new schema is different from the active schema, update and
        # invalidate table cache.
        if new_schema != queryset.active_schema:
            queryset.active_schema = new_schema
            queryset.table = None  # Force regeneration of the table with the new schema.

    return queryset

def prefetch_related(self, *prefetch: Prefetch) -> QuerySet[EdgyModel, EdgyEmbedTarget]:
    """
    Performs a reverse lookup for foreign keys and other relationships,
    populating results onto the main model instances.

    This method is distinct from `select_related` in that `select_related`
    performs a SQL JOIN to fetch related data in the same query, whereas
    `prefetch_related` executes separate queries for each relationship
    and then joins the results in Python. This is particularly useful for
    many-to-many relationships or reverse foreign key lookups, or when
    preloading related objects for a large set of parent objects.

    Args:
        *prefetch (Prefetch): One or more `Prefetch` objects, each defining
                               a relationship to prefetch, including the
                               `related_name` and the `to_attr` where results
                               will be stored. An optional custom `QuerySet`
                               can also be provided within the `Prefetch` object.

    Returns:
        QuerySet: A new `QuerySet` instance with the specified prefetch
                  relationships configured. This new QuerySet can then be
                  further filtered, ordered, or executed.

    Raises:
        QuerySetError: If any argument passed to `prefetch` is not an
                       instance of the `Prefetch` class.
    """
    self_queryset = cast("QuerySet[EdgyModel, EdgyEmbedTarget]", self)
    queryset: QuerySet = self_queryset._clone()

    # Validate that all provided arguments are instances of Prefetch.
    if any(not isinstance(value, Prefetch) for value in prefetch):
        raise QuerySetError("The prefetch_related must have Prefetch type objects only.")

    # Append the new prefetch objects to the queryset's internal list.
    queryset._prefetch_related = tuple(chain(self_queryset._prefetch_related, prefetch))
    select_pathes: set[str] = set()
    # this one extra doesn't matter much from performance perspective, is maybe even cheaper
    if queryset._embed_parent:
        # can be cross db, so only add the first expanded part
        crawl_result = crawl_relationship(
            queryset.model_class,
            queryset._embed_parent[0],
            model_database=queryset.database,
            traverse_last=True,
        )
        select_pathes.add(crawl_result.forward_path)
    # now add the forward pathes, we need to resolve them because embed_parent is resolved
    select_pathes.update(
        *(
            prefetch._generate_select_related_pathes(queryset)
            for prefetch in queryset._prefetch_related
        )
    )
    # they are sanitized and analyzed later in _update_related_weak
    queryset._update_related_weak(
        select_pathes, cache_name="_select_related_embedding", clear=True
    )
    return queryset

async def _apply_prefetches_self_and_related(
    self,
    *,
    instance: EdgyModel,
    prepared_prefetches: Sequence[Prefetch],
    mapping: Mapping,
    tables_and_models: tables_and_models_type | None,
) -> None:
    """Apply prefetches on select related branches.

    This method applies prefetching logic to the current queryset by considering
    select related relationships and embedding targets. It handles fetching related
    data concurrently and embedding parents in the result set.

    Args:
        instance: The current EdgyModel instance.
        prepared_prefetches: A Sequence of prepared Prefetch objects.
        mapping: The mapping object used for database operations.
        tables_and_models: Optional tables and models information.

    Returns:
        None
    """
    self_queryset = cast("QuerySet[EdgyModel, EdgyEmbedTarget]", self)
    seen: set[str] = set()
    target_dict: dict[str, list[Prefetch]] = {}
    anchor_dict: dict[str, Mapping | None] = {}
    for prefetch in prepared_prefetches:
        target_dict.setdefault(
            prefetch.to_attr.rsplit("__", 1)[0] if "__" in prefetch.to_attr else "", []
        ).append(prefetch)
        if prefetch.from_anchor and prefetch._anchor.cross_db_remainder:
            anchor_dict[prefetch.from_anchor] = None

    token = MODEL_GETATTR_BEHAVIOR.set("coro")
    try:
        # we need to check the prefetches dicts for the select pathes
        for path in anchor_dict:
            prefix = ""
            current_instance: BaseModelType | None | Awaitable[BaseModelType | None] = instance
            for part in path.split("__"):
                prefix = f"{prefix}__{part}" if prefix else part
                if prefix in seen:
                    continue
                seen.add(prefix)
                # part is empty when path="" (empty, only main model)
                if part:
                    current_instance = cast(
                        "BaseModelType | None | Awaitable[BaseModelType | None]",
                        getattr(current_instance, part, None),
                    )
                if isawaitable(current_instance):
                    current_instance = await current_instance
                if current_instance is None:
                    break
                if prefix in anchor_dict:
                    # we need the pks of the current instance and can't just select them (different db)
                    await current_instance.load(only_needed=True)
                    mapping = current_instance.extract_column_values(
                        current_instance.extract_db_fields(only=instance.pkcolumns),
                        instance=current_instance,
                        phase="",
                    )
                    anchor_dict[prefetch.from_anchor] = mapping
        # reset
        seen.clear()
        for path in target_dict:
            prefix = ""
            current_instance = instance
            for part in path.split("__"):
                prefix = f"{prefix}__{part}" if prefix else part
                if prefix in seen:
                    continue
                seen.add(prefix)
                # part is empty when path="" (empty, only main model)
                if part:
                    current_instance = cast(
                        "BaseModelType | None | Awaitable[BaseModelType | None]",
                        getattr(current_instance, part, None),
                    )
                if isawaitable(current_instance):
                    current_instance = await current_instance
                if current_instance is None:
                    break
                if prefetches_list := target_dict.get(prefix):
                    await run_concurrently(
                        [
                            prefetch._init_bake()
                            for prefetch in prefetches_list
                            if prefetch._bake_without_mapping
                        ],
                        limit=1
                        if getattr(self_queryset.database, "force_rollback", False)
                        else None,
                    )
                    await _apply_prefetches_helper(
                        instance=current_instance,
                        mapping=mapping,
                        tables_and_models=tables_and_models,
                        prefetches=prefetches_list,
                        anchor_dict=anchor_dict,
                    )
    finally:
        MODEL_GETATTR_BEHAVIOR.reset(token)

async def _embed_parent_in_result(
    self,
    result: EdgyModel | Awaitable[EdgyModel] | None,
    *,
    prepared_prefetches: Sequence[Prefetch] | None = None,
    mapping: Mapping | None = None,
    tables_and_models: tables_and_models_type | None = None,
) -> tuple[EdgyModel, EdgyEmbedTarget] | tuple[None, None]:
    """
    This is a result transformation and apply prefetches.
    """
    self_queryset = cast("QuerySet[EdgyModel, EdgyEmbedTarget]", self)
    if isawaitable(result):
        result = await result
    if result is None:
        return None, None
    if mapping is not None and prepared_prefetches:
        if tables_and_models is None:
            tables_and_models = (await self_queryset.as_select_with_tables())[1]
        # apply embedding
        prefetched_instance = result
        if self_queryset._embed_parent_filters:
            token = MODEL_GETATTR_BEHAVIOR.set("coro")
            try:
                for part in self_queryset._embed_parent_filters[0].split("__"):
                    prefetched_instance = getattr(prefetched_instance, part)
                    if isawaitable(prefetched_instance):
                        prefetched_instance = await prefetched_instance
            finally:
                MODEL_GETATTR_BEHAVIOR.reset(token)
        await self._apply_prefetches_self_and_related(
            instance=prefetched_instance,
            prepared_prefetches=prepared_prefetches,
            mapping=mapping,
            tables_and_models=tables_and_models,
        )
    if not self_queryset._embed_parent:
        return result, cast("EdgyEmbedTarget", result)
    prefix = ""
    token = MODEL_GETATTR_BEHAVIOR.set("coro")
    try:
        new_result: Any = result
        for part in self_queryset._embed_parent[0].split("__"):
            prefix = f"{prefix}__{part}" if prefix else part
            new_result = getattr(new_result, part)
            if isawaitable(new_result):
                new_result = await new_result
    finally:
        MODEL_GETATTR_BEHAVIOR.reset(token)
    if self_queryset._embed_parent[1]:
        # this works also on strict models
        object.__setattr__(new_result, self_queryset._embed_parent[1], result)
    return result, new_result

def _prepare_prefetches_for_rows(
    self, *, rows: Sequence[sqlalchemy.Row], tables_and_models: tables_and_models_type
) -> Sequence[Prefetch]:
    """
    Builds the Prefetch objects for a given batch of results.
    This is the *prefetch building* half of the original _handle_batch.

    Args:
        rows: The current batch of SQLAlchemy Row objects.
        tables_and_models: The table/model mapping from the compiler.

    Returns:
        A list of populated Prefetch objects, ready to be executed.

    Raises:
        NotImplementedError: If a prefetch crosses database boundaries.
        QuerySetError: If a prefetch path is invalid (e.g., unidirectional).
    """
    self_queryset = cast("QuerySet[EdgyModel, EdgyEmbedTarget]", self)
    prepared_prefetches: dict[str, Prefetch] = {}
    seen_prefetches: set[tuple[str, str, str]] = set()

    for prefetch in self_queryset._prefetch_related:
        compare_tuple = (
            prefetch.related_name,
            prefetch.from_anchor,
            prefetch.to_attr,
        )
        if compare_tuple in seen_prefetches:
            continue
        else:
            seen_prefetches.add(compare_tuple)
        new_prefetch = Prefetch(
            related_name=prefetch.related_name,
            to_attr=prefetch.to_attr,
            from_anchor=prefetch.from_anchor,
        )

        anchor_crawl_result = crawl_relationship(
            self_queryset.model_class,
            prefetch.from_anchor,
            embed_parent=self_queryset._embed_parent_filters,
            model_database=self_queryset.database,
            # allow_crossing_db=True,
            traverse_last=True,
            allow_crossing_db=True,
        )
        if anchor_crawl_result.field_name or anchor_crawl_result.operator:
            raise QuerySetError(
                detail=f"`from_anchor` path points to a non-relation field: `{anchor_crawl_result.field_name}`."
            )
        new_prefetch._anchor = anchor_crawl_result

        prefetch_crawl_result = crawl_relationship(
            anchor_crawl_result.model_class,
            prefetch.related_name,
            traverse_last=True,
        )
        if prefetch_crawl_result.cross_db_remainder:
            raise QuerySetError(
                detail=(
                    f"The `related_name` crosses from: `{prefetch_crawl_result.forward_path}` the database. "
                    "Please use an appropiate `anchor_from`."
                )
            )

        if prefetch_crawl_result.field_name or prefetch_crawl_result.operator:
            raise QuerySetError(
                detail=f"The `related_name` path points to non-relation field: `{prefetch_crawl_result.field_name}`."
            )
        if prefetch_crawl_result.reverse_path is False:
            raise QuerySetError(
                detail=(
                    "Creating a reverse path from `related_name` is not possible, unidirectional fields were used."
                )
            )

        crawl_relationship(
            self_queryset.model_class,
            prefetch.to_attr,
            traverse_last=True,
            allow_crossing_db=True,
            callback_fn=_target_crawl_fn,
        )
        prefetch_queryset: QuerySet | None = prefetch.queryset
        if prefetch_queryset is None:
            prefetch_queryset = prefetch_crawl_result.model_class.query.all()
        else:
            prefetch_queryset = prefetch_queryset.all()
        # just add the resolved path to _select_related_embedding, so get less cruft
        prefetch_queryset._select_related_embedding.add(prefetch_crawl_result.reverse_path)
        # the assigned queryset has an empty cache
        new_prefetch.queryset = prefetch_queryset
        # now check key if the prefetch configuration is already registered
        # if yes, raise.
        if new_prefetch.to_attr in prepared_prefetches:
            raise QuerySetError(
                f"Conflicting prefetches to_attr='{new_prefetch.to_attr}'"
                f"'on {prefetch_crawl_result.model_class.__name__}"
            )
        # check collision on target model
        new_prefetch.check_for_collision(prefetch_crawl_result.model_class)
        new_prefetch._reverse_path_to_anchor = prefetch_crawl_result.reverse_path
        new_prefetch._baking_finished = asyncio.Event()
        new_prefetch._baked_results = {}
        if not new_prefetch._anchor.cross_db_remainder:
            new_prefetch._set_clauses_by_mappings(
                mappings=(row._mapping for row in rows), tables_and_models=tables_and_models
            )
        prepared_prefetches[new_prefetch.to_attr] = new_prefetch
    return list(prepared_prefetches.values())

def using_with_db(
    self, connection_name: str, schema: str | Any | None | Literal[False] = Undefined
) -> QuerySet[EdgyModel, EdgyEmbedTarget]:
    """
    Switches the database connection and optionally the schema for the queryset.

    This method is deprecated in favor of the more flexible `using` method, which
    accepts both `database` and `schema` as keyword arguments.

    Args:
        connection_name (str): The name of the database connection (registered
                               in the model's registry's `extra` connections)
                               to switch to.
        schema (str | None | Literal[False]): The schema name to use.
            Use a schema name (`str`), `False` (unset schema and use default),
            `None` (no specific schema), or `Undefined` (retain current schema).

    Returns:
        QuerySet: A new QuerySet instance configured with the specified database
                  connection and schema.

    Warnings:
        DeprecationWarning: This method is deprecated. Users should migrate to
                            `using(database=..., schema=...)` for future
                            compatibility.
    """
    warnings.warn(
        "'using_with_db' is deprecated in favor of 'using' with schema, database arguments.",
        DeprecationWarning,
        stacklevel=2,
    )
    # Delegate to the `using` method with the appropriate keyword arguments.
    return self.using(database=connection_name, schema=schema)

def _create_clone_instance(self) -> QuerySet[EdgyModel, EdgyEmbedTarget]:
    """Base instance which is decorated later in clone."""
    return cast("type[QuerySet]", type(self))(
        model_class=self.model_class, using_schema=self.using_schema
    )

def _clone(self) -> QuerySet[EdgyModel, EdgyEmbedTarget]:
    """
    This is core to the builder pattern. Most cache is refreshed

    Note: the _cached_select_related_expression is transferred.
    """
    queryset = self._create_clone_instance()
    queryset._database = getattr(self, "_database", None)
    queryset._table = getattr(self, "_table", None)
    queryset._prefetch_related = self._prefetch_related
    queryset._exclude_secrets = self._exclude_secrets
    # copying won't work, we would need a deep copy but not necessary anyway
    queryset._reference_select = self._reference_select
    queryset._offset = self._offset
    # tuple, so we can just move it
    queryset._order_by = self._order_by
    # tuple, so we can just move it
    queryset._group_by = self._group_by
    # tuple, so we can just move it
    queryset.distinct_on = self.distinct_on
    # tuple, so we can just move it
    queryset._extra_select = self._extra_select
    queryset.limit_count = self.limit_count
    queryset._batch_size = self._batch_size
    queryset.filter_clauses.extend(self.filter_clauses)
    queryset.or_clauses.extend(self.or_clauses)
    # this handles the create arguments
    queryset._injected_create_handler = self._injected_create_handler
    queryset._embed_parent = self._embed_parent
    queryset._embed_parent_filters = self._embed_parent_filters
    queryset._only.update(self._only)
    queryset._defer.update(self._defer)
    queryset._select_related.update(self._select_related)
    queryset._select_related_g_and_o.update(self._select_related_g_and_o)
    queryset._select_related_embedding.update(self._select_related_embedding)
    # by default this is copied, we need to clear it when select_related caches are changing
    queryset._cached_select_related_expression = self._cached_select_related_expression
    queryset._for_update = self._for_update
    return cast("QuerySet", queryset)

async def _as_select_with_tables(
    self,
) -> tuple[Any, tables_and_models_type]:
    """
    (This is the new internal method for the base class)
    Builds the query select by delegating to the QueryCompiler.
    """
    compiler = QueryCompiler(self)
    self._get_join_graph_data()
    expression, tables_and_models = await compiler.build_select()
    return expression, tables_and_models

def _clear_cache(
    self, *, keep_result_cache: bool = False, keep_cached_selected: bool = False
) -> None:
    if not keep_result_cache:
        self._cache.clear()
    if not keep_cached_selected:
        self._cached_select_with_tables: (
            tuple[Any, dict[str, tuple[sqlalchemy.Table, type[BaseModelType]]]] | None
        ) = None
    self._cache_count: int | None = None
    self._cache_first: tuple[EdgyModel, EdgyEmbedTarget] | None = None
    self._cache_last: tuple[EdgyModel, EdgyEmbedTarget] | None = None
    self._cache_fetch_all: bool = False

def _build_order_by_iterable(
    self, order_by: Iterable[str], tables_and_models: tables_and_models_type
) -> Iterable:
    """
    This is a helper for the *compiler* but is called by it,
     so it's okay for it to live here as it's part of the 'builder' logic.
    """
    return (self._prepare_order_by(entry, tables_and_models) for entry in order_by)