Python haystack.query.SearchQuerySet() Examples

def solr_counts():
    total_q = SearchQuerySet()
    latest_q = total_q.filter(latest=True)
    registrations_q = latest_q.filter(category="entity_status::ACT")
    last_24h = datetime.now() - timedelta(days=1)
    last_week = datetime.now() - timedelta(days=7)
    last_month = datetime.now() - timedelta(days=30)
    last_24h_q = total_q.filter(create_timestamp__gte=last_24h)
    last_week_q = total_q.filter(create_timestamp__gte=last_week)
    last_month_q = total_q.filter(create_timestamp__gte=last_month)
    try:
        return {
            "total": total_q.count(),
            "active": latest_q.count(),
            "registrations": registrations_q.count(),
            "last_month": last_month_q.count(),
            "last_week": last_week_q.count(),
            "last_24h": last_24h_q.count(),
        }
    except SolrError:
        LOGGER.exception("Error when retrieving quickload counts from Solr")
        return False 

def test_from_queryset(self):
        """ Verify that a DistinctCountsSearchQuerySet can be built from an existing SearchQuerySet."""
        course_1 = CourseFactory()
        CourseRunFactory(title='foo', course=course_1)
        CourseRunFactory(title='foo', course=course_1)

        course_2 = CourseFactory()
        CourseRunFactory(title='foo', course=course_2)
        CourseRunFactory(title='bar', course=course_2)

        queryset = SearchQuerySet().filter(title='foo').models(CourseRun)
        dc_queryset = DistinctCountsSearchQuerySet.from_queryset(queryset)

        expected = sorted([run.key for run in queryset])
        actual = sorted([run.key for run in dc_queryset])
        assert expected == actual 

def search(self):
        # First, store the SearchQuerySet received from other processing.
        sqs = super().search()

        if not self.is_valid():
            return self.no_query_found()

        # Custom date range filter
        # TODO can this be done with native-haystack?
        if 'start_date' in self.data:
            try:
                sqs = sqs.filter(date__gte=datetime.datetime.strptime(self.data.get('start_date'), '%Y-%m-%d'))
            except ValueError:
                logger.error('Invalid start_date')

        if 'end_date' in self.data:
            try:
                sqs = sqs.filter(date__lte=datetime.datetime.strptime(self.data.get('end_date'), '%Y-%m-%d'))
            except ValueError as e:
                logger.error('Invalid end_date')

        return sqs 

def autocomplete_view(request):
    """Stub for auto-complete feature(title for all objects missing)

    """
    suggestions_limit = 5
    sqs = SearchQuerySet().autocomplete(title=request.GET.get('q', ''))[:suggestions_limit]

    # for result in sqs:  # type: SearchResult
    #     print(result.object)
    #     print(result.title)

    suggestions = [result.title for result in sqs]

    return JsonResponse({
        'results': suggestions
    }) 

def handle(self, *args, **options):
        sqs = SearchQuerySet()

        users = sqs.models(User).all().count()
        images = sqs.models(Image).all().count()

        integration = 0
        for i in sqs.models(Image).all():
            integration += i.integration
        integration = int(integration / 3600.0)

        gs = GlobalStat(
            users = users,
            images = images,
            integration = integration)
        gs.save() 

def handle_item(self, item):
        """Perform MLT-query on item"""
        sqs = SearchQuerySet().models(self.model)
        mlt_results = sqs.more_like_this(item)
        saved = []
        for result in mlt_results:  # type: SearchResult

            rel = self.model_relation(score=result.score)
            rel.set_relation(seed_id=item.pk, related_id=result.pk)
            rel.save()

            logger.debug('Saved %s' % rel)

            saved.append(rel)

        return saved 

def test_enrollable_paid_seat_boosting(self, has_enrollable_paid_seats, paid_seat_enrollment_end, expects_boost):
        """
        Verify that CourseRuns for which an unenrolled user may enroll and
        purchase a paid Seat are boosted.
        """

        # Create a control record (one that should never be boosted).
        with patch.object(CourseRun, 'has_enrollable_paid_seats', return_value=False):
            with patch.object(CourseRun, 'get_paid_seat_enrollment_end', return_value=None):
                self.build_normalized_course_run(title='test1')

        # Create the test record (may be boosted).
        with patch.object(CourseRun, 'has_enrollable_paid_seats', return_value=has_enrollable_paid_seats):
            with patch.object(CourseRun, 'get_paid_seat_enrollment_end', return_value=paid_seat_enrollment_end):
                test_record = self.build_normalized_course_run(title='test2')

        search_results = SearchQuerySet().models(CourseRun).all()
        assert len(search_results) == 2

        if expects_boost:
            assert search_results[0].score > search_results[1].score
            assert test_record.title == search_results[0].title
        else:
            assert search_results[0].score == search_results[1].score 

def test_viewset_get_queryset_with_DjangoModelPermissions_permission(self):
        from rest_framework.permissions import DjangoModelPermissions
        setattr(self.view, "permission_classes", (DjangoModelPermissions,))

        # The `DjangoModelPermissions` is not supported and should raise an
        # AssertionError from rest_framework.permissions.
        request = factory.get(path="/", data="", content_type="application/json")
        try:
            self.view.as_view(actions={"get": "list"})(request)
            self.fail("Did not fail with AssertionError or AttributeError "
                      "when calling HaystackView with DjangoModelPermissions")
        except (AttributeError, AssertionError) as e:
            if isinstance(e, AttributeError):
                self.assertEqual(str(e), "'SearchQuerySet' object has no attribute 'model'")
            else:
                self.assertEqual(str(e), "Cannot apply DjangoModelPermissions on a view that does "
                                         "not have `.model` or `.queryset` property.") 

def test_viewset_get_queryset_with_DjangoModelPermissionsOrAnonReadOnly_permission(self):
        from rest_framework.permissions import DjangoModelPermissionsOrAnonReadOnly
        setattr(self.view, "permission_classes", (DjangoModelPermissionsOrAnonReadOnly,))

        # The `DjangoModelPermissionsOrAnonReadOnly` is not supported and should raise an
        # AssertionError from rest_framework.permissions.
        request = factory.get(path="/", data="", content_type="application/json")
        try:
            self.view.as_view(actions={"get": "list"})(request)
            self.fail("Did not fail with AssertionError when calling HaystackView "
                      "with DjangoModelPermissionsOrAnonReadOnly")
        except (AttributeError, AssertionError) as e:
            if isinstance(e, AttributeError):
                self.assertEqual(str(e), "'SearchQuerySet' object has no attribute 'model'")
            else:
                self.assertEqual(str(e), "Cannot apply DjangoModelPermissions on a view that does "
                                         "not have `.model` or `.queryset` property.") 

def test_with_distinct_counts_raises_when_queryset_includes_unsupported_options(self):
        """
        Verify that an error is raised if the original queryset includes options that are not supported by our
        custom Query class.
        """
        dc_queryset = DistinctCountsSearchQuerySet.from_queryset(SearchQuerySet())

        with pytest.raises(RuntimeError) as err:
            now = datetime.datetime.now()
            ten_days = datetime.timedelta(days=10)
            start = now - ten_days
            end = now + ten_days
            dc_queryset.date_facet('start', start, end, 'day').with_distinct_counts('aggregation_key')
        assert str(err.value) == 'DistinctCountsSearchQuery does not support date facets.'

        with pytest.raises(RuntimeError) as err:
            dc_queryset.facet('pacing_type', order='term').with_distinct_counts('aggregation_key')
        assert 'DistinctCountsSearchQuery only supports a limited set of field facet options.' in str(err.value) 

def test_enrollable_course_run_boosting(self, enrollment_start, enrollment_end, expects_boost):
        """Verify that enrollable CourseRuns are boosted."""

        # Create a control record that should never be boosted
        self.build_normalized_course_run(title='test1')

        # Create the test record
        test_record = self.build_normalized_course_run(
            title='test2',
            enrollment_start=enrollment_start,
            enrollment_end=enrollment_end
        )

        search_results = SearchQuerySet().models(CourseRun).all()
        assert len(search_results) == 2

        if expects_boost:
            assert search_results[0].score > search_results[1].score
            assert test_record.title == search_results[0].title
        else:
            assert search_results[0].score == search_results[1].score 

def search(cls, query, queryset=None):
        """ Queries the search index.

        Args:
            query (str) -- Elasticsearch querystring (e.g. `title:intro*`)
            queryset (models.QuerySet) -- base queryset to search, defaults to objects.all()

        Returns:
            QuerySet
        """
        query = clean_query(query)

        if queryset is None:
            queryset = cls.objects.all()

        if query == '(*)':
            # Early-exit optimization. Wildcard searching is very expensive in elasticsearch. And since we just
            # want everything, we don't need to actually query elasticsearch at all.
            return queryset

        results = SearchQuerySet().models(cls).raw_search(query)
        ids = {result.pk for result in results}

        return queryset.filter(pk__in=ids) 

def list(self, request):
        search_string = request.query_params.get('q')
        if search_string:
            try:
                book = models.Book.objects.get_for_user(self.request.user)
            except models.Book.DoesNotExist:
                # rethink the logic here - how could a user not have a book?
                return Response(
                    "No contacts for user",
                    status=status.HTTP_400_BAD_REQUEST,
                )
            results = SearchQuerySet().filter(book=book.id).filter(
                SQ(name=AutoQuery(search_string)) | SQ(content=AutoQuery(search_string))
            )
            results = [result.object for result in results]
        else:
            results = models.Contact.objects.get_contacts_for_user(
                self.request.user,
            )
        serializer = serializers.ContactSerializer(results, many=True)
        return Response(serializer.data) 

def global_stats():
    from astrobin.models import Image, GlobalStat
    sqs = SearchQuerySet()

    users = sqs.models(User).all().count()
    images = sqs.models(Image).all().count()

    integration = 0
    for i in sqs.models(Image).all():
        integration += i.integration
    integration = int(integration / 3600.0)

    gs = GlobalStat(
        users=users,
        images=images,
        integration=integration)
    gs.save() 

def test_serializer_multiple_index_data(self):
        objs = SearchQuerySet().filter(text="John")
        serializer = self.serializer1(instance=objs, many=True)
        data = serializer.data

        self.assertEqual(len(data), 4)
        for result in data:
            if "name" in result:
                self.assertTrue("species" in result, "Pet results should have 'species' and 'name' fields")
                self.assertTrue("firstname" not in result, "Pet results should have 'species' and 'name' fields")
                self.assertTrue("lastname" not in result, "Pet results should have 'species' and 'name' fields")
            elif "firstname" in result:
                self.assertTrue("lastname" in result, "Person results should have 'firstname' and 'lastname' fields")
                self.assertTrue("name" not in result, "Person results should have 'firstname' and 'lastname' fields")
                self.assertTrue("species" not in result, "Person results should have 'firstname' and 'lastname' fields")
            else:
                self.fail("Result should contain either Pet or Person fields") 

def test_serializer_multiple_index_declared_fields(self):
        objs = SearchQuerySet().filter(text="John")
        serializer = self.serializer2(instance=objs, many=True)
        data = serializer.data

        self.assertEqual(len(data), 4)
        for result in data:
            if "name" in result:
                self.assertTrue("extra" in result, "'extra' should be present in Pet results")
                self.assertEqual(result["extra"], 1, "The value of 'extra' should be 1")
                self.assertTrue("hair_color" not in result, "'hair_color' should not be present in Pet results")
            elif "lastname" in result:
                self.assertTrue("extra" in result, "'extra' should be present in Person results")
                self.assertEqual(result["extra"], 1, "The value of 'extra' should be 1")
                self.assertTrue("hair_color" in result, "'hair_color' should be present in Person results")
                self.assertEqual(result["hair_color"], "black", "The value of 'hair_color' should be 'black'")
            else:
                self.fail("Result should contain either Pet or Person fields") 

def handle_delete(self, sender, instance, **kwargs):
        # modified from BaseSignalProcessor.handle_delete to force checking existence before removing from the index
        # (and getting an error message if it's not there).
        using_backends = self.connection_router.for_write(instance=instance)

        for using in using_backends:
            try:
                index = self.connections[using].get_unified_index().get_index(sender)

                # check to see if the object is actually in the index before removing:
                index.prepare(instance)
                ct = index.prepared_data['django_ct']
                obj_id = index.prepared_data['id']
                existing = SearchQuerySet().models(sender).filter(django_ct=ct, id=obj_id)

                if existing.count() > 0:
                    index.remove_object(instance, using=using)
            except NotHandled:
                pass 

def solr_counts():
    total_q = SearchQuerySet()
    latest_q = total_q.filter(latest=True)
    registrations_q = latest_q.filter(category="entity_status::ACT")
    last_24h = datetime.now() - timedelta(days=1)
    last_week = datetime.now() - timedelta(days=7)
    last_month = datetime.now() - timedelta(days=30)
    last_24h_q = total_q.filter(create_timestamp__gte=last_24h)
    last_week_q = total_q.filter(create_timestamp__gte=last_week)
    last_month_q = total_q.filter(create_timestamp__gte=last_month)
    try:
        return {
            "total_indexed_items": total_q.count(),
            "active": latest_q.count(),
            "registrations": registrations_q.count(),
            "last_month": last_month_q.count(),
            "last_week": last_week_q.count(),
            "last_24h": last_24h_q.count(),
        }
    except SolrError:
        LOGGER.exception("Error when retrieving quickload counts from Solr")
        return False 

def test_multi_serializer(self):
        objs = SearchQuerySet().filter(text="Zane")
        serializer = self.serializer1(instance=objs, many=True)
        self.assertEqual(
            json.loads(json.dumps(serializer.data)),
            [{
                "has_rabies": True,
                "text": "Zane",
                "name": "Zane",
                "species": "Dog"
            },
            {
                "text": "Zane Griffith\n",
                "firstname": "Zane",
                "lastname": "Griffith",
                "description": "Zane is a nice chap!"
            }]
        ) 

def update_top100_ids():
    from astrobin.models import Image

    LOCK_EXPIRE = 60 * 5  # Lock expires in 5 minutes
    lock_id = 'top100_ids_lock'

    # cache.add fails if the key already exists
    acquire_lock = lambda: cache.add(lock_id, 'true', LOCK_EXPIRE)
    # memcache delete is very slow, but we have to use it to take
    # advantage of using add() for atomic locking
    release_lock = lambda: cache.delete(lock_id)

    logger.debug('Building Top100 ids...')
    if acquire_lock():
        try:
            sqs = SearchQuerySet().models(Image).order_by('-likes')
            top100_ids = [int(x.pk) for x in sqs][:100]
            cache.set('top100_ids', top100_ids, 60 * 60 * 24)
        finally:
            release_lock()
        return

    logger.debug(
        'Top100 ids task is already being run by another worker') 

def test_serializer_get_fields_with_ignore_fields(self):
        obj = SearchQuerySet().filter(lastname="Foreman")[0]
        serializer = self.serializer3(instance=obj)
        fields = serializer.get_fields()

        self.assertIsInstance(fields, dict)
        self.assertIsInstance(fields["text"], CharField)
        self.assertIsInstance(fields["firstname"], CharField)
        self.assertIsInstance(fields["lastname"], CharField)
        self.assertFalse("autocomplete" in fields) 

def test_serialize_field_is_correct_type(self):
        obj = SearchQuerySet().models(MockAllField).latest('datetimefield')
        serializer = self.serializer1(instance=obj, many=False)

        self.assertIsInstance(serializer.fields['charfield'], fields.HaystackCharField)
        self.assertIsInstance(serializer.fields['integerfield'], fields.HaystackIntegerField)
        self.assertIsInstance(serializer.fields['floatfield'], fields.HaystackFloatField)
        self.assertIsInstance(serializer.fields['decimalfield'], fields.HaystackDecimalField)
        self.assertIsInstance(serializer.fields['boolfield'], fields.HaystackBooleanField)
        self.assertIsInstance(serializer.fields['datefield'], fields.HaystackDateField)
        self.assertIsInstance(serializer.fields['datetimefield'], fields.HaystackDateTimeField)
        self.assertIsInstance(serializer.fields['multivaluefield'], fields.HaystackMultiValueField) 

def test_serializer_get_fields_with_exclude(self):
        obj = SearchQuerySet().filter(lastname="Foreman")[0]
        serializer = self.serializer2(instance=obj)
        fields = serializer.get_fields()

        self.assertIsInstance(fields, dict)
        self.assertIsInstance(fields["text"], CharField)
        self.assertIsInstance(fields["lastname"], CharField)
        self.assertIsInstance(fields["autocomplete"], CharField)
        self.assertFalse("firstname" in fields) 

def test_serializer_boolean_field(self):
        dog = self.serializer7(instance=SearchQuerySet().filter(species="Dog")[0])
        iguana = self.serializer7(instance=SearchQuerySet().filter(species="Iguana")[0])
        self.assertTrue(dog.data["has_rabies"])
        self.assertFalse(iguana.data["has_rabies"]) 

def get(self, request):
        deals = SearchQuerySet().autocomplete(
            content_auto=request.GET.get('q', '')
        )
        return TemplateResponse(request, self.template_name, {'deals': deals}) 

def stats(request):
    response_dict = {}

    sqs = SearchQuerySet()
    gs = GlobalStat.objects.first()

    if gs:
        response_dict['total_users'] = gs.users
        response_dict['total_images'] = gs.images
        response_dict['total_integration'] = gs.integration

    sort = '-user_integration'
    if 'sort' in request.GET:
        sort = request.GET.get('sort')
        if sort == 'tot_integration':
            sort = '-user_integration'
        elif sort == 'avg_integration':
            sort = '-user_avg_integration'
        elif sort == 'images':
            sort = '-user_images'

    queryset = sqs.filter(user_images__gt=0).models(User).order_by(sort)

    return object_list(
        request,
        queryset=queryset,
        template_name='stats.html',
        template_object_name='user',
        extra_context=response_dict,
    ) 

def test_live_query(self):
        engine = ElasticsearchMultilingualSearchEngine()
        es = engine.backend('default', **Data.connection_options)
        es.setup()
        sqs = SearchQuerySet()
        self.assertFalse(sqs.query.has_run())
        self.assertIsInstance(sqs.query, ElasticsearchMultilingualSearchQuery)
        all_results = sqs.all()
        self.assertEqual(list(all_results), [])  # execute a query
        self.assertTrue(isinstance(all_results, SearchQuerySet)) 

def test_query(self, mock_es):
        sqs = SearchQuerySet()
        self.assertFalse(sqs.query.has_run())
        self.assertIsInstance(sqs.query, ElasticsearchMultilingualSearchQuery)
        all_results = sqs.all()
        all_results.query.backend = mock_backend()
        list(all_results)
        self.assertTrue(all_results.query.backend.search.called)
        self.assertEqual('*:*', all_results.query.backend.search.call_args[0][0]) 

def get_search_contacts(self):
        book = self.request.current_book
        self.query_raw = self.request.GET.get('q')
        results = re.split(
            r'(?P<tag>\w+\:(?:\"[\w\s]+\"|\w+\b))',
            self.query_raw,
        )
        self.search_tags = []
        parts = []
        for result in results:
            if result.startswith('tag:'):
                tag_str = result.strip().split(':')[1].strip('"')
                try:
                    self.search_tags.append(Tag.objects.get(book=book, tag=tag_str))
                except Tag.DoesNotExist:
                    pass
            else:
                parts.append(result.strip())
        searchqueryset = SearchQuerySet().filter(book=book.id)
        if self.search_tags:
            searchqueryset = searchqueryset.filter(
                tags_ids__in=[tag.id for tag in self.search_tags],
            )

        self.query = ' '.join(parts).strip()
        sqs = searchqueryset.filter(
            SQ(name=AutoQuery(self.query)) | SQ(content=AutoQuery(self.query))
        )
        try:
            contact_ids = [result.object.id for result in sqs]
        except:
            contact_ids = []
        return contact_ids 

def get(self, *args, **kwargs):

        query = self.request.GET.get('q', '')

        sqs = SearchQuerySet().autocomplete(
            content_auto=query, title=query)[:5]

        suggestions = []

        for result in sqs:

            r = {
                'title': result.title,
            }
            if result.content_type == "Image":
                r['url'] = result.object.url
            else:
                r['url'] = result.object.get_absolute_url()

            suggestions.append(r)

        # Make sure you return a JSON object, not a bare list.
        # Otherwise, you could be vulnerable to an XSS attack.
        the_data = json.dumps({
            'results': suggestions
        })
        return HttpResponse(the_data, content_type='application/json')