wagtail 7.1.1__py3-none-any.whl → 7.2rc1__py3-none-any.whl

wagtail/search/backends/database/postgres/postgres.py

@@ -1,749 +1 @@
-import warnings
-from collections import OrderedDict
-from functools import reduce
-
-from django.contrib.postgres.search import SearchQuery, SearchRank, SearchVector
-from django.db import (
-    NotSupportedError,
-    connections,
-    router,
-    transaction,
-)
-from django.db.models import Avg, Count, F, Manager, Q, TextField, Value
-from django.db.models.constants import LOOKUP_SEP
-from django.db.models.functions import Cast, Length
-from django.db.models.sql.subqueries import InsertQuery
-from django.utils.encoding import force_str
-from django.utils.functional import cached_property
-
-from ....index import AutocompleteField, RelatedFields, SearchField, get_indexed_models
-from ....models import IndexEntry
-from ....query import And, Boost, MatchAll, Not, Or, Phrase, PlainText
-from ....utils import (
-    ADD,
-    MUL,
-    OR,
-    get_content_type_pk,
-    get_descendants_content_types_pks,
-)
-from ...base import (
-    BaseSearchBackend,
-    BaseSearchQueryCompiler,
-    BaseSearchResults,
-    FilterFieldError,
-)
-from .query import Lexeme
-from .weights import get_sql_weights, get_weight
-
-EMPTY_VECTOR = SearchVector(Value("", output_field=TextField()))
-
-
-class ObjectIndexer:
-    """
-    Responsible for extracting data from an object to be inserted into the index.
-    """
-
-    def __init__(self, obj, backend):
-        self.obj = obj
-        self.search_fields = obj.get_search_fields()
-        self.config = backend.config
-        self.autocomplete_config = backend.autocomplete_config
-
-    def prepare_value(self, value):
-        if isinstance(value, str):
-            return value
-
-        elif isinstance(value, list):
-            return ", ".join(self.prepare_value(item) for item in value)
-
-        elif isinstance(value, dict):
-            return ", ".join(self.prepare_value(item) for item in value.values())
-
-        return force_str(value)
-
-    def prepare_field(self, obj, field):
-        if isinstance(field, SearchField):
-            yield (
-                field,
-                get_weight(field.boost),
-                self.prepare_value(field.get_value(obj)),
-            )
-
-        elif isinstance(field, AutocompleteField):
-            # AutocompleteField does not define a boost parameter, so use a base weight of 'D'
-            yield (field, "D", self.prepare_value(field.get_value(obj)))
-
-        elif isinstance(field, RelatedFields):
-            sub_obj = field.get_value(obj)
-            if sub_obj is None:
-                return
-
-            if isinstance(sub_obj, Manager):
-                sub_objs = sub_obj.all()
-
-            else:
-                if callable(sub_obj):
-                    sub_obj = sub_obj()
-
-                sub_objs = [sub_obj]
-
-            for sub_obj in sub_objs:
-                for sub_field in field.fields:
-                    yield from self.prepare_field(sub_obj, sub_field)
-
-    def as_vector(self, texts, for_autocomplete=False):
-        """
-        Converts an array of strings into a SearchVector that can be indexed.
-        """
-        texts = [(text.strip(), weight) for text, weight in texts]
-        texts = [(text, weight) for text, weight in texts if text]
-
-        if not texts:
-            return EMPTY_VECTOR
-
-        search_config = self.autocomplete_config if for_autocomplete else self.config
-
-        return ADD(
-            [
-                SearchVector(
-                    Value(text, output_field=TextField()),
-                    weight=weight,
-                    config=search_config,
-                )
-                for text, weight in texts
-            ]
-        )
-
-    @cached_property
-    def id(self):
-        """
-        Returns the value to use as the ID of the record in the index
-        """
-        return force_str(self.obj.pk)
-
-    @cached_property
-    def title(self):
-        """
-        Returns all values to index as "title". This is the value of all SearchFields that have the field_name 'title'
-        """
-        texts = []
-        for field in self.search_fields:
-            for current_field, boost, value in self.prepare_field(self.obj, field):
-                if (
-                    isinstance(current_field, SearchField)
-                    and current_field.field_name == "title"
-                ):
-                    texts.append((value, boost))
-
-        return self.as_vector(texts)
-
-    @cached_property
-    def body(self):
-        """
-        Returns all values to index as "body". This is the value of all SearchFields excluding the title
-        """
-        texts = []
-        for field in self.search_fields:
-            for current_field, boost, value in self.prepare_field(self.obj, field):
-                if (
-                    isinstance(current_field, SearchField)
-                    and not current_field.field_name == "title"
-                ):
-                    texts.append((value, boost))
-
-        return self.as_vector(texts)
-
-    @cached_property
-    def autocomplete(self):
-        """
-        Returns all values to index as "autocomplete". This is the value of all AutocompleteFields
-        """
-        texts = []
-        for field in self.search_fields:
-            for current_field, boost, value in self.prepare_field(self.obj, field):
-                if isinstance(current_field, AutocompleteField):
-                    texts.append((value, boost))
-
-        return self.as_vector(texts, for_autocomplete=True)
-
-
-class Index:
-    def __init__(self, backend):
-        self.backend = backend
-        self.name = self.backend.index_name
-
-        self.read_connection = connections[router.db_for_read(IndexEntry)]
-        self.write_connection = connections[router.db_for_write(IndexEntry)]
-
-        if (
-            self.read_connection.vendor != "postgresql"
-            or self.write_connection.vendor != "postgresql"
-        ):
-            raise NotSupportedError(
-                "You must select a PostgreSQL database to use PostgreSQL search."
-            )
-
-        self.entries = IndexEntry._default_manager.all()
-
-    def add_model(self, model):
-        pass
-
-    def refresh(self):
-        pass
-
-    def _refresh_title_norms(self, full=False):
-        """
-        Refreshes the value of the title_norm field.
-
-        This needs to be set to 'lavg/ld' where:
-         - lavg is the average length of titles in all documents (also in terms)
-         - ld is the length of the title field in this document (in terms)
-        """
-
-        lavg = (
-            self.entries.annotate(title_length=Length("title"))
-            .filter(title_length__gt=0)
-            .aggregate(Avg("title_length"))["title_length__avg"]
-        )
-
-        if full:
-            # Update the whole table
-            # This is the most accurate option but requires a full table rewrite
-            # so we can't do it too often as it could lead to locking issues.
-            entries = self.entries
-
-        else:
-            # Only update entries where title_norm is 1.0
-            # This is the default value set on new entries.
-            # It's possible that other entries could have this exact value but there shouldn't be too many of those
-            entries = self.entries.filter(title_norm=1.0)
-
-        entries.annotate(title_length=Length("title")).filter(
-            title_length__gt=0
-        ).update(title_norm=lavg / F("title_length"))
-
-    def delete_stale_model_entries(self, model):
-        existing_pks = model._default_manager.annotate(
-            object_id=Cast("pk", TextField())
-        ).values("object_id")
-        content_types_pks = get_descendants_content_types_pks(model)
-        stale_entries = self.entries.filter(
-            content_type_id__in=content_types_pks
-        ).exclude(object_id__in=existing_pks)
-        stale_entries.delete()
-
-    def delete_stale_entries(self):
-        for model in get_indexed_models():
-            # We don’t need to delete stale entries for non-root models,
-            # since we already delete them by deleting roots.
-            if not model._meta.parents:
-                self.delete_stale_model_entries(model)
-
-    def add_item(self, obj):
-        self.add_items(obj._meta.model, [obj])
-
-    def add_items(self, model, objs):
-        search_fields = model.get_search_fields()
-        if not search_fields:
-            return
-
-        indexers = [ObjectIndexer(obj, self.backend) for obj in objs]
-
-        # TODO: Delete unindexed objects while dealing with proxy models.
-        if not indexers:
-            return
-
-        content_type_pk = get_content_type_pk(model)
-        compiler = InsertQuery(IndexEntry).get_compiler(
-            connection=self.write_connection
-        )
-        title_sql = []
-        autocomplete_sql = []
-        body_sql = []
-        data_params = []
-
-        for indexer in indexers:
-            data_params.extend((content_type_pk, indexer.id))
-
-            # Compile title value
-            value = compiler.prepare_value(
-                IndexEntry._meta.get_field("title"), indexer.title
-            )
-            sql, params = value.as_sql(compiler, self.write_connection)
-            title_sql.append(sql)
-            data_params.extend(params)
-
-            # Compile autocomplete value
-            value = compiler.prepare_value(
-                IndexEntry._meta.get_field("autocomplete"), indexer.autocomplete
-            )
-            sql, params = value.as_sql(compiler, self.write_connection)
-            autocomplete_sql.append(sql)
-            data_params.extend(params)
-
-            # Compile body value
-            value = compiler.prepare_value(
-                IndexEntry._meta.get_field("body"), indexer.body
-            )
-            sql, params = value.as_sql(compiler, self.write_connection)
-            body_sql.append(sql)
-            data_params.extend(params)
-
-        data_sql = ", ".join(
-            [
-                f"(%s, %s, {a}, {b}, {c}, 1.0)"
-                for a, b, c in zip(title_sql, autocomplete_sql, body_sql)
-            ]
-        )
-
-        with self.write_connection.cursor() as cursor:
-            cursor.execute(
-                """
-                INSERT INTO %s (content_type_id, object_id, title, autocomplete, body, title_norm)
-                (VALUES %s)
-                ON CONFLICT (content_type_id, object_id)
-                DO UPDATE SET title = EXCLUDED.title,
-                              title_norm = 1.0,
-                              autocomplete = EXCLUDED.autocomplete,
-                              body = EXCLUDED.body
-                """
-                % (IndexEntry._meta.db_table, data_sql),
-                data_params,
-            )
-
-        self._refresh_title_norms()
-
-    def delete_item(self, item):
-        item.index_entries.all()._raw_delete(using=self.write_connection.alias)
-
-    def __str__(self):
-        return self.name
-
-
-class PostgresSearchQueryCompiler(BaseSearchQueryCompiler):
-    DEFAULT_OPERATOR = "and"
-    LAST_TERM_IS_PREFIX = False
-    TARGET_SEARCH_FIELD_TYPE = SearchField
-    HANDLES_ORDER_BY_EXPRESSIONS = True
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        local_search_fields = self.get_search_fields_for_model()
-
-        # Due to a Django bug, arrays are not automatically converted
-        # when we use WEIGHTS_VALUES.
-        self.sql_weights = get_sql_weights()
-
-        if self.fields is None:
-            # search over the fields defined on the current model
-            self.search_fields = local_search_fields
-        else:
-            # build a search_fields set from the passed definition,
-            # which may involve traversing relations
-            self.search_fields = {
-                field_lookup: self.get_search_field(
-                    field_lookup, fields=local_search_fields
-                )
-                for field_lookup in self.fields
-            }
-
-    def get_config(self, backend):
-        return backend.config
-
-    def get_search_fields_for_model(self):
-        return self.queryset.model.get_searchable_search_fields()
-
-    def get_search_field(self, field_lookup, fields=None):
-        if fields is None:
-            fields = self.search_fields
-
-        if LOOKUP_SEP in field_lookup:
-            field_lookup, sub_field_name = field_lookup.split(LOOKUP_SEP, 1)
-        else:
-            sub_field_name = None
-
-        for field in fields:
-            if (
-                isinstance(field, self.TARGET_SEARCH_FIELD_TYPE)
-                and field.field_name == field_lookup
-            ):
-                return field
-
-            # Note: Searching on a specific related field using
-            # `.search(fields=…)` is not yet supported by Wagtail.
-            # This method anticipates by already implementing it.
-            if isinstance(field, RelatedFields) and field.field_name == field_lookup:
-                return self.get_search_field(sub_field_name, field.fields)
-
-    def build_tsquery_content(self, query, config=None, invert=False):
-        if isinstance(query, PlainText):
-            terms = query.query_string.split()
-            if not terms:
-                return None
-
-            last_term = terms.pop()
-
-            lexemes = Lexeme(last_term, invert=invert, prefix=self.LAST_TERM_IS_PREFIX)
-            for term in terms:
-                new_lexeme = Lexeme(term, invert=invert)
-
-                if query.operator == "and":
-                    lexemes &= new_lexeme
-                else:
-                    lexemes |= new_lexeme
-
-            return SearchQuery(lexemes, search_type="raw", config=config)
-
-        elif isinstance(query, Phrase):
-            return SearchQuery(query.query_string, search_type="phrase", config=config)
-
-        elif isinstance(query, Boost):
-            # Not supported
-            msg = "The Boost query is not supported by the PostgreSQL search backend."
-            warnings.warn(msg, RuntimeWarning)
-
-            return self.build_tsquery_content(
-                query.subquery, config=config, invert=invert
-            )
-
-        elif isinstance(query, Not):
-            return self.build_tsquery_content(
-                query.subquery, config=config, invert=not invert
-            )
-
-        elif isinstance(query, (And, Or)):
-            # If this part of the query is inverted, we swap the operator and
-            # pass down the inversion state to the child queries.
-            # This works thanks to De Morgan's law.
-            #
-            # For example, the following query:
-            #
-            #   Not(And(Term("A"), Term("B")))
-            #
-            # Is equivalent to:
-            #
-            #   Or(Not(Term("A")), Not(Term("B")))
-            #
-            # It's simpler to code it this way as we only need to store the
-            # invert status of the terms rather than all the operators.
-
-            subquery_lexemes = [
-                self.build_tsquery_content(subquery, config=config, invert=invert)
-                for subquery in query.subqueries
-            ]
-
-            is_and = isinstance(query, And)
-
-            if invert:
-                is_and = not is_and
-
-            if is_and:
-                return reduce(lambda a, b: a & b, subquery_lexemes)
-            else:
-                return reduce(lambda a, b: a | b, subquery_lexemes)
-
-        raise NotImplementedError(
-            "`%s` is not supported by the PostgreSQL search backend."
-            % query.__class__.__name__
-        )
-
-    def build_tsquery(self, query, config=None):
-        return self.build_tsquery_content(query, config=config)
-
-    def build_tsrank(self, vector, query, config=None, boost=1.0):
-        if isinstance(query, (Phrase, PlainText, Not)):
-            rank_expression = SearchRank(
-                vector,
-                self.build_tsquery(query, config=config),
-                weights=self.sql_weights,
-            )
-
-            if boost != 1.0:
-                rank_expression *= boost
-
-            return rank_expression
-
-        elif isinstance(query, Boost):
-            boost *= query.boost
-            return self.build_tsrank(vector, query.subquery, config=config, boost=boost)
-
-        elif isinstance(query, And):
-            return (
-                MUL(
-                    1 + self.build_tsrank(vector, subquery, config=config, boost=boost)
-                    for subquery in query.subqueries
-                )
-                - 1
-            )
-
-        elif isinstance(query, Or):
-            return ADD(
-                self.build_tsrank(vector, subquery, config=config, boost=boost)
-                for subquery in query.subqueries
-            ) / (len(query.subqueries) or 1)
-
-        raise NotImplementedError(
-            "`%s` is not supported by the PostgreSQL search backend."
-            % query.__class__.__name__
-        )
-
-    def get_index_vectors(self, search_query):
-        return [
-            (F("index_entries__title"), F("index_entries__title_norm")),
-            (F("index_entries__body"), 1.0),
-        ]
-
-    def get_fields_vectors(self, search_query):
-        return [
-            (
-                SearchVector(
-                    field_lookup,
-                    config=search_query.config,
-                ),
-                search_field.boost,
-            )
-            for field_lookup, search_field in self.search_fields.items()
-        ]
-
-    def get_search_vectors(self, search_query):
-        if self.fields is None:
-            return self.get_index_vectors(search_query)
-
-        else:
-            return self.get_fields_vectors(search_query)
-
-    def _build_rank_expression(self, vectors, config):
-        rank_expressions = [
-            self.build_tsrank(vector, self.query, config=config) * boost
-            for vector, boost in vectors
-        ]
-
-        rank_expression = rank_expressions[0]
-        for other_rank_expression in rank_expressions[1:]:
-            rank_expression += other_rank_expression
-
-        return rank_expression
-
-    def search(self, config, start, stop, score_field=None):
-        # TODO: Handle MatchAll nested inside other search query classes.
-        if isinstance(self.query, MatchAll):
-            return self.queryset[start:stop]
-
-        elif isinstance(self.query, Not) and isinstance(self.query.subquery, MatchAll):
-            return self.queryset.none()
-
-        search_query = self.build_tsquery(self.query, config=config)
-        vectors = self.get_search_vectors(search_query)
-        rank_expression = self._build_rank_expression(vectors, config)
-
-        combined_vector = vectors[0][0]
-        for vector, boost in vectors[1:]:
-            combined_vector = combined_vector._combine(vector, "||", False)
-
-        queryset = self.queryset.annotate(_vector_=combined_vector).filter(
-            _vector_=search_query
-        )
-
-        if self.order_by_relevance:
-            queryset = queryset.order_by(rank_expression.desc(), "-pk")
-
-        elif not queryset.query.order_by:
-            # Adds a default ordering to avoid issue #3729.
-            queryset = queryset.order_by("-pk")
-            rank_expression = F("pk")
-
-        if score_field is not None:
-            queryset = queryset.annotate(**{score_field: rank_expression})
-
-        return queryset[start:stop]
-
-    def _process_lookup(self, field, lookup, value):
-        lhs = field.get_attname(self.queryset.model) + "__" + lookup
-        return Q(**{lhs: value})
-
-    def _process_match_none(self):
-        return Q(pk__in=[])
-
-    def _connect_filters(self, filters, connector, negated):
-        if connector == "AND":
-            q = Q(*filters)
-
-        elif connector == "OR":
-            q = OR([Q(fil) for fil in filters])
-
-        else:
-            return
-
-        if negated:
-            q = ~q
-
-        return q
-
-
-class PostgresAutocompleteQueryCompiler(PostgresSearchQueryCompiler):
-    LAST_TERM_IS_PREFIX = True
-    TARGET_SEARCH_FIELD_TYPE = AutocompleteField
-
-    def get_config(self, backend):
-        return backend.autocomplete_config
-
-    def get_search_fields_for_model(self):
-        return self.queryset.model.get_autocomplete_search_fields()
-
-    def get_index_vectors(self, search_query):
-        return [(F("index_entries__autocomplete"), 1.0)]
-
-    def get_fields_vectors(self, search_query):
-        return [
-            (
-                SearchVector(
-                    field_lookup,
-                    config=search_query.config,
-                    weight="D",
-                ),
-                1.0,
-            )
-            for field_lookup, search_field in self.search_fields.items()
-        ]
-
-
-class PostgresSearchResults(BaseSearchResults):
-    def get_queryset(self, for_count=False):
-        if for_count:
-            start = None
-            stop = None
-        else:
-            start = self.start
-            stop = self.stop
-
-        return self.query_compiler.search(
-            self.query_compiler.get_config(self.backend),
-            start,
-            stop,
-            score_field=self._score_field,
-        )
-
-    def _do_search(self):
-        return list(self.get_queryset())
-
-    def _do_count(self):
-        return self.get_queryset(for_count=True).count()
-
-    supports_facet = True
-
-    def facet(self, field_name):
-        # Get field
-        field = self.query_compiler._get_filterable_field(field_name)
-        if field is None:
-            raise FilterFieldError(
-                'Cannot facet search results with field "'
-                + field_name
-                + "\". Please add index.FilterField('"
-                + field_name
-                + "') to "
-                + self.query_compiler.queryset.model.__name__
-                + ".search_fields.",
-                field_name=field_name,
-            )
-
-        query = self.query_compiler.search(
-            self.query_compiler.get_config(self.backend), None, None
-        )
-        results = (
-            query.values(field_name).annotate(count=Count("pk")).order_by("-count")
-        )
-
-        return OrderedDict(
-            [(result[field_name], result["count"]) for result in results]
-        )
-
-
-class PostgresSearchRebuilder:
-    def __init__(self, index):
-        self.index = index
-
-    def start(self):
-        self.index.delete_stale_entries()
-        return self.index
-
-    def finish(self):
-        self.index._refresh_title_norms(full=True)
-
-
-class PostgresSearchAtomicRebuilder(PostgresSearchRebuilder):
-    def __init__(self, index):
-        super().__init__(index)
-        self.transaction = transaction.atomic(using=index.write_connection.alias)
-        self.transaction_opened = False
-
-    def start(self):
-        self.transaction.__enter__()
-        self.transaction_opened = True
-        return super().start()
-
-    def finish(self):
-        self.index._refresh_title_norms(full=True)
-
-        self.transaction.__exit__(None, None, None)
-        self.transaction_opened = False
-
-    def __del__(self):
-        # TODO: Implement a cleaner way to close the connection on failure.
-        if self.transaction_opened:
-            self.transaction.needs_rollback = True
-            self.finish()
-
-
-class PostgresSearchBackend(BaseSearchBackend):
-    query_compiler_class = PostgresSearchQueryCompiler
-    autocomplete_query_compiler_class = PostgresAutocompleteQueryCompiler
-    results_class = PostgresSearchResults
-    rebuilder_class = PostgresSearchRebuilder
-    atomic_rebuilder_class = PostgresSearchAtomicRebuilder
-
-    def __init__(self, params):
-        super().__init__(params)
-        self.index_name = params.get("INDEX", "default")
-        self.config = params.get("SEARCH_CONFIG")
-
-        # Use 'simple' config for autocomplete to disable stemming
-        # A good description for why this is important can be found at:
-        # https://www.postgresql.org/docs/9.1/datatype-textsearch.html#DATATYPE-TSQUERY
-        self.autocomplete_config = params.get("AUTOCOMPLETE_SEARCH_CONFIG", "simple")
-
-        if params.get("ATOMIC_REBUILD", False):
-            self.rebuilder_class = self.atomic_rebuilder_class
-
-    def get_index_for_model(self, model):
-        return Index(self)
-
-    def get_index_for_object(self, obj):
-        return self.get_index_for_model(obj._meta.model)
-
-    def reset_index(self):
-        for connection in [
-            connection
-            for connection in connections.all()
-            if connection.vendor == "postgresql"
-        ]:
-            IndexEntry._default_manager.all()._raw_delete(using=connection.alias)
-
-    def add_type(self, model):
-        pass  # Not needed.
-
-    def refresh_index(self):
-        pass  # Not needed.
-
-    def add(self, obj):
-        self.get_index_for_object(obj).add_item(obj)
-
-    def add_bulk(self, model, obj_list):
-        if obj_list:
-            self.get_index_for_object(obj_list[0]).add_items(model, obj_list)
-
-    def delete(self, obj):
-        self.get_index_for_object(obj).delete_item(obj)
-
-
-SearchBackend = PostgresSearchBackend
+from wagtailmodelsearch.backends.database.postgres.postgres import *  # noqa: F403