From ab896e318a1310e95353136b2d66edd2c83d12c4 Mon Sep 17 00:00:00 2001
From: Karl Hobley <karl@torchbox.com>
Date: Wed, 17 Sep 2014 12:00:16 +0100
Subject: [PATCH] DB search now returns a SearchResults object

Previously, it returned a QuerySet which gave the database backend much
more API than the Elasticsearch backend has. Which could potentially
cause issues if the developer is using the DB backend locally and the
Elasticsearch backend on their production server.

Fixes #460
---
 wagtail/wagtailsearch/backends/db.py | 85 ++++++++++++++++++----------
 1 file changed, 54 insertions(+), 31 deletions(-)

diff --git a/wagtail/wagtailsearch/backends/db.py b/wagtail/wagtailsearch/backends/db.py
index d190e29e8..cbde8d58f 100644
--- a/wagtail/wagtailsearch/backends/db.py
+++ b/wagtail/wagtailsearch/backends/db.py
@@ -1,6 +1,57 @@
 from django.db import models
 
-from wagtail.wagtailsearch.backends.base import BaseSearch
+from wagtail.wagtailsearch.backends.base import BaseSearch, BaseSearchResults
+
+
+class DBSearchQuery(object):
+    def __init__(self, queryset, query_string, fields=None):
+        self.queryset = queryset
+        self.query_string = query_string
+        self.fields = fields
+
+    def get_queryset(self):
+        queryset = self.queryset
+        model = queryset.model
+
+        if self.query_string is not None:
+            # Get fields
+            fields = self.fields or [field.field_name for field in model.get_searchable_search_fields()]
+
+            # Get terms
+            terms = self.query_string.split()
+            if not terms:
+                return model.objects.none()
+
+            # Filter by terms
+            for term in terms:
+                term_query = models.Q()
+                for field_name in fields:
+                    # Check if the field exists (this will filter out indexed callables)
+                    try:
+                        model._meta.get_field_by_name(field_name)
+                    except:
+                        continue
+
+                    # Filter on this field
+                    term_query |= models.Q(**{'%s__icontains' % field_name: term})
+
+                queryset = queryset.filter(term_query)
+
+            # Distinct
+            queryset = queryset.distinct()
+
+        return queryset
+
+
+class DBSearchResults(BaseSearchResults):
+    def get_queryset(self):
+        return self.query.get_queryset()[self.start:self.stop]
+
+    def _do_search(self):
+        return self.get_queryset()
+
+    def _do_count(self):
+        return self.get_queryset().count()
 
 
 class DBSearch(BaseSearch):
@@ -20,38 +71,10 @@ class DBSearch(BaseSearch):
         pass # Not needed
 
     def add_bulk(self, model, obj_list):
-        pass # Not needed
+        return # Not needed
 
     def delete(self, obj):
         pass # Not needed
 
     def _search(self, queryset, query_string, fields=None):
-        if query_string is not None:
-            # Get fields
-            if fields is None:
-                fields = [field.field_name for field in queryset.model.get_searchable_search_fields()]
-
-            # Get terms
-            terms = query_string.split()
-            if not terms:
-                return queryset.model.objects.none()
-
-            # Filter by terms
-            for term in terms:
-                term_query = models.Q()
-                for field_name in fields:
-                    # Check if the field exists (this will filter out indexed callables)
-                    try:
-                        queryset.model._meta.get_field_by_name(field_name)
-                    except:
-                        continue
-
-                    # Filter on this field
-                    term_query |= models.Q(**{'%s__icontains' % field_name: term})
-
-                queryset = queryset.filter(term_query)
-
-            # Distinct
-            queryset = queryset.distinct()
-
-        return queryset
+        return DBSearchResults(self, DBSearchQuery(queryset, query_string, fields=fields))