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 using(
    self,
    _positional: Any = _sentinel,
    *,
    database: str | Any | None | Database = Undefined,
    schema: str | Any | None | Literal[False] = Undefined,
) -> QuerySet:
    """
    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", self._clone())

    # Process the 'database' argument.
    if database is not Undefined:
        if isinstance(database, Database):
            connection: Database = database
        elif database is None:
            # Use the default database from the model's registry.
            connection = self.model_class.meta.registry.database
        else:
            # Assert that the database name exists in the extra connections.
            assert database is None or database in self.model_class.meta.registry.extra, (
                f"`{database}` is not in the connections extra of the model"
                f"`{self.model_class.__name__}` registry"
            )
            # Retrieve the database from extra connections.
            connection = self.model_class.meta.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:
    """
    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.
    """
    queryset: QuerySet = self._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 = [*self._prefetch_related, *prefetch]
    return queryset

def using_with_db(
    self, connection_name: str, schema: str | Any | None | Literal[False] = Undefined
) -> QuerySet:
    """
    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 _clone(self) -> QuerySet:
    """
    This is core to the builder pattern
    """
    queryset = self.__class__(
        self.model_class,
        database=getattr(self, "_database", None),
        filter_clauses=self.filter_clauses,
        prefetch_related=self._prefetch_related,
        limit=self.limit_count,
        offset=self._offset,
        batch_size=self._batch_size,
        order_by=self._order_by,
        group_by=self._group_by,
        distinct=self.distinct_on,
        only=self._only,
        defer=self._defer,
        embed_parent=self.embed_parent,
        using_schema=self.using_schema,
        table=getattr(self, "_table", None),
        exclude_secrets=self._exclude_secrets,
        reference_select=self._reference_select,
        extra_select=self._extra_select,
    )
    queryset.or_clauses.extend(self.or_clauses)
    queryset.embed_parent_filters = self.embed_parent_filters
    queryset._select_related.update(self._select_related)
    queryset._select_related_weak.update(self._select_related_weak)
    queryset._cached_select_related_expression = self._cached_select_related_expression
    queryset._for_update = self._for_update.copy() if self._for_update is not None else None
    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[BaseModelType, Any] | None = None
    self._cache_last: tuple[BaseModelType, Any] | 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)

def _validate_only_and_defer(self) -> None:
    if self._only and self._defer:
        raise QuerySetError("You cannot use .only() and .defer() at the same time.")

def _get_join_graph_data(self) -> tuple[Any, tables_and_models_type]:
    """
    Gets the join graph, building it via the compiler if needed.

    This is the new "bridge" that manages the
    _cached_select_related_expression variable to satisfy brittle tests,
    while keeping the compiler itself stateless.
    """
    if self._cached_select_related_expression is None:
        # Create a compiler just to build the join graph
        compiler = QueryCompiler(self)

        # Call the compiler's build method and cache the result
        self._cached_select_related_expression = compiler.build_join_graph()
    return self._cached_select_related_expression

async def as_select_with_tables(
    self,
) -> tuple[Any, tables_and_models_type]:
    """
    (Refactored: Now delegates to the Compiler)
    """
    if self._cached_select_with_tables is None:
        self._cached_select_with_tables = await self._as_select_with_tables()
    return self._cached_select_with_tables

async def as_select(
    self,
) -> Any:
    return (await self.as_select_with_tables())[0]

def _kwargs_to_clauses(
    self,
    kwargs: Any,
) -> tuple[list[Any], set[str]]:
    """
    This is part of the 'filter' builder logic
    """
    clauses = []
    select_related: set[str] = set()
    cleaned_kwargs = clauses_mod.clean_query_kwargs(
        self.model_class, kwargs, self.embed_parent_filters, model_database=self.database
    )

    for key, value in cleaned_kwargs.items():
        model_class, field_name, op, related_str, _, cross_db_remainder = crawl_relationship(
            self.model_class, key
        )
        if related_str:
            select_related.add(related_str)
        field = model_class.meta.fields.get(field_name, clauses_mod.generic_field)
        if cross_db_remainder:
            assert field is not clauses_mod.generic_field
            fk_field = cast(BaseForeignKey, field)
            sub_query = (
                fk_field.target.query.filter(**{cross_db_remainder: value})
                .only(*fk_field.related_columns.keys())
                .values_list(fields=fk_field.related_columns.keys())
            )

            async def wrapper(
                queryset: QuerySet,
                tables_and_models: tables_and_models_type,
                *,
                _field: BaseFieldType = field,
                _sub_query: QuerySet = sub_query,
                _prefix: str = related_str,
            ) -> Any:
                table = tables_and_models[_prefix][0]
                fk_tuple = sqlalchemy.tuple_(
                    *(getattr(table.columns, colname) for colname in _field.get_column_names())
                )
                return fk_tuple.in_(await _sub_query)

            clauses.append(wrapper)
        else:
            assert not isinstance(value, BaseModelType), (
                f"should be parsed in clean: {key}: {value}"
            )

            async def wrapper(
                queryset: QuerySet,
                tables_and_models: tables_and_models_type,
                *,
                _field: BaseFieldType = field,
                _value: Any = value,
                _op: str | None = op,
                _prefix: str = related_str,
                _field_name: str = field_name,
            ) -> Any:
                _value = await clauses_mod.parse_clause_arg(
                    _value, queryset, tables_and_models
                )
                table = tables_and_models[_prefix][0]
                return _field.operator_to_clause(_field_name, _op, table, _value)

            wrapper._edgy_force_callable_queryset_filter = True
            clauses.append(wrapper)

    return clauses, select_related

def _prepare_order_by(self, order_by: str, tables_and_models: tables_and_models_type) -> Any:
    """
    (Helper for 'order_by' builder logic, but called by compiler)
    """
    reverse = order_by.startswith("-")
    order_by = order_by.lstrip("-")
    crawl_result = clauses_mod.clean_path_to_crawl_result(
        self.model_class,
        path=order_by,
        embed_parent=self.embed_parent_filters,
        model_database=self.database,
    )
    order_col = tables_and_models[crawl_result.forward_path][0].columns[
        crawl_result.field_name
    ]
    return order_col.desc() if reverse else order_col

def _update_select_related_weak(self, fields: Iterable[str], *, clear: bool) -> bool:
    related: set[str] = set()
    for field_name in fields:
        field_name = field_name.lstrip("-")
        related_element = clauses_mod.clean_path_to_crawl_result(
            self.model_class,
            path=field_name,
            embed_parent=self.embed_parent_filters,
            model_database=self.database,
        ).forward_path
        if related_element:
            related.add(related_element)
    if related and not self._select_related.union(self._select_related_weak).issuperset(
        related
    ):
        self._cached_select_related_expression = None
        if clear:
            self._select_related_weak.clear()
        self._select_related_weak.update(related)
        return True
    return False

def _update_select_related(self, pathes: Iterable[str]) -> None:
    related: set[str] = set()
    for path in pathes:
        path = path.lstrip("-")
        crawl_result = clauses_mod.clean_path_to_crawl_result(
            self.model_class,
            path=path,
            embed_parent=self.embed_parent_filters,
            model_database=self.database,
        )
        related_element = (
            crawl_result.field_name
            if not crawl_result.forward_path
            else f"{crawl_result.forward_path}__{crawl_result.field_name}"
        )
        if crawl_result.cross_db_remainder:
            raise QuerySetError(
                detail=f'Selected path "{related_element}" is on another database.'
            )
        if related_element:
            related.add(related_element)
    if related and not self._select_related.issuperset(related):
        self._cached_select_related_expression = None
        self._select_related.update(related)

def _prepare_distinct(
    self, distinct_on: str, tables_and_models: tables_and_models_type
) -> sqlalchemy.Column:
    """Helper for 'distinct' builder, but called by compiler"""
    crawl_result = clauses_mod.clean_path_to_crawl_result(
        self.model_class,
        path=distinct_on,
        embed_parent=self.embed_parent_filters,
        model_database=self.database,
    )
    return tables_and_models[crawl_result.forward_path][0].columns[crawl_result.field_name]

async def _embed_parent_in_result(
    self, result: EdgyModel | Awaitable[EdgyModel]
) -> tuple[EdgyModel, Any]:
    """
    This is a result transformation, called by the Parser.
    """
    if isawaitable(result):
        result = await result
    if not self.embed_parent:
        return result, result
    token = MODEL_GETATTR_BEHAVIOR.set("coro")
    try:
        new_result: Any = result
        for part in self.embed_parent[0].split("__"):
            new_result = getattr(new_result, part)
            if isawaitable(new_result):
                new_result = await new_result
    finally:
        MODEL_GETATTR_BEHAVIOR.reset(token)
    if self.embed_parent[1]:
        setattr(new_result, self.embed_parent[1], result)
    return result, new_result

def get_schema(self) -> str | None:
    schema = self.using_schema
    if schema is Undefined:
        schema = get_schema()
    if schema is None:
        schema = self.model_class.get_db_schema()
    return schema

async def _execute_iterate(
    self, fetch_all_at_once: bool = False
) -> AsyncIterator[BaseModelType]:
    """
    (Refactored: Now delegates to the Executor)
    """
    # Create the specialists
    compiler = QueryCompiler(self)
    parser = ResultParser(self)
    executor = QueryExecutor(self, compiler, parser)

    # Delegate the work
    async for model in executor.iterate(fetch_all_at_once=fetch_all_at_once):  # type: ignore
        yield model

async def _execute_all(self) -> list[EdgyModel]:
    """
    Still relies on _execute_iterate.
    """
    results = [result async for result in self._execute_iterate(fetch_all_at_once=True)]

    # Only attempt dedupe for "normal" querysets.
    # Embedded querysets (e.g. album.tracks with embed_parent) must not go
    # through this path – their model_class and returned objects differ.
    if (
        len(results) > 1
        and self.embed_parent is None
        and not getattr(self, "_suppress_pk_deduplication", False)
    ):
        pk_names = tuple(self.model_class.pkcolumns)
        if not pk_names:
            # Nothing reliable to dedupe on
            return results  # type: ignore

        seen: set[tuple] = set()
        unique = []

        for obj in results:
            try:
                key = tuple(getattr(obj, name) for name in pk_names)
            except AttributeError:
                # The returned object does not expose the model_class PK attrs;
                # this can happen in advanced/embedded scenarios. In that case
                # we bail out and keep the original list to avoid breaking
                # existing behaviour.
                return results  # type: ignore

            if key not in seen:
                seen.add(key)
                unique.append(obj)

        results = unique

    return results  # type: ignore

def _filter_or_exclude(
    self,
    kwargs: Any,
    clauses: Sequence[
        sqlalchemy.sql.expression.BinaryExpression
        | Callable[
            [QuerySetType],
            sqlalchemy.sql.expression.BinaryExpression
            | Awaitable[sqlalchemy.sql.expression.BinaryExpression],
        ]
        | dict[str, Any]
        | QuerySet
    ],
    exclude: bool = False,
    or_: bool = False,
    allow_global_or: bool = True,
) -> QuerySet:
    """
    This is the core 'filter' builder logic.
    """
    from edgy.core.db.querysets.queryset import QuerySet

    queryset: QuerySet = self._clone()
    if kwargs:
        clauses = [*clauses, kwargs]
    converted_clauses: Sequence[
        sqlalchemy.sql.expression.BinaryExpression
        | Callable[
            [QuerySetType],
            sqlalchemy.sql.expression.BinaryExpression
            | Awaitable[sqlalchemy.sql.expression.BinaryExpression],
        ]
    ] = []
    for raw_clause in clauses:
        if isinstance(raw_clause, dict):
            extracted_clauses, related = queryset._kwargs_to_clauses(kwargs=raw_clause)
            if not queryset._select_related.issuperset(related):
                queryset._select_related.update(related)
                queryset._cached_select_related_expression = None
            if or_ and extracted_clauses:
                wrapper_and = clauses_mod.and_(*extracted_clauses, no_select_related=True)

                if allow_global_or and len(clauses) == 1:
                    # Global OR mode: promote existing AND filters into the OR group.
                    # This turns:
                    #   qs.filter(A).or_(B)
                    # into:
                    #   OR( AND(A), AND(B) )
                    # instead of: OR(B) AND A.
                    assert not exclude

                    if queryset.filter_clauses:
                        # Wrap existing filters into a single AND group and move them to or_clauses
                        existing_and = clauses_mod.and_(
                            *queryset.filter_clauses, no_select_related=True
                        )
                        queryset.or_clauses.append(existing_and)
                        # Clear filter_clauses so they are not ANDed again later
                        queryset.filter_clauses = []

                    # Add the new OR operand
                    queryset.or_clauses.append(wrapper_and)
                    return queryset

                # Non-global OR (e.g. local_or) or multiple clauses:
                # just collect and handle them at the end as a local OR group.
                converted_clauses.append(wrapper_and)
            else:
                converted_clauses.extend(extracted_clauses)
        elif isinstance(raw_clause, QuerySet):
            assert raw_clause.model_class is queryset.model_class, (
                f"QuerySet arg has wrong model_class {raw_clause.model_class}"
            )
            converted_clauses.append(raw_clause.build_where_clause)
            if not queryset._select_related.issuperset(raw_clause._select_related):
                queryset._select_related.update(raw_clause._select_related)
                queryset._cached_select_related_expression = None
        else:
            clause = raw_clause

            # Support global OR mode for non-dict clauses (e.g. Q objects, raw callables)
            if or_ and allow_global_or and len(clauses) == 1:
                # Global OR only makes sense for non-exclude queries
                assert not exclude

                # If there are existing AND filters, promote them into the OR group
                if queryset.filter_clauses:
                    existing_and = clauses_mod.and_(
                        *queryset.filter_clauses,
                        no_select_related=True,
                    )
                    queryset.or_clauses.append(existing_and)
                    queryset.filter_clauses = []

                # Propagate select_related coming from this clause, if any
                if hasattr(clause, "_edgy_calculate_select_related"):
                    select_related_calculated = clause._edgy_calculate_select_related(queryset)
                    if not queryset._select_related.issuperset(select_related_calculated):
                        queryset._select_related.update(select_related_calculated)
                        queryset._cached_select_related_expression = None

                # Add this clause as a new OR branch and return immediately
                queryset.or_clauses.append(clause)
                return queryset

            # Normal path (no global OR promotion)
            converted_clauses.append(clause)
            if hasattr(clause, "_edgy_calculate_select_related"):
                select_related_calculated = clause._edgy_calculate_select_related(queryset)
                if not queryset._select_related.issuperset(select_related_calculated):
                    queryset._select_related.update(select_related_calculated)
                    queryset._cached_select_related_expression = None
    if not converted_clauses:
        return queryset

    if exclude:
        op = clauses_mod.and_ if not or_ else clauses_mod.or_

        queryset.filter_clauses.append(
            clauses_mod.not_(
                op(*converted_clauses, no_select_related=True), no_select_related=True
            )
        )
    elif or_:
        queryset.filter_clauses.append(
            clauses_mod.or_(*converted_clauses, no_select_related=True)
        )
    else:
        queryset.filter_clauses.extend(converted_clauses)
    return queryset