From 9dbabffc81afb6a0812a828701479e8a7dd83c2a Mon Sep 17 00:00:00 2001
From: clime <clime@redhat.com>
Date: Tue, 10 May 2016 18:44:27 +0200
Subject: [PATCH] added whoosheer option for whooshee_search method

---
 flask_whooshee.py | 32 ++++++++++++++++----------------
 test.py           | 30 ++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/flask_whooshee.py b/flask_whooshee.py
index 6a32fb2..46927b2 100644
--- a/flask_whooshee.py
+++ b/flask_whooshee.py
@@ -17,8 +17,7 @@ from sqlalchemy.orm.mapper import Mapper
 class WhoosheeQuery(BaseQuery):
     """An override for SQLAlchemy query used to do fulltext search."""
 
-    # TODO: add an option to override used Whoosheer
-    def whooshee_search(self, search_string, group=whoosh.qparser.OrGroup,
+    def whooshee_search(self, search_string, group=whoosh.qparser.OrGroup, whoosheer=None,
                         match_substrings=True, limit=None, order_by_relevance=10):
         """Do a fulltext search on the query.
 
@@ -32,21 +31,22 @@ class WhoosheeQuery(BaseQuery):
         Returns:
             query filtered with results of the fulltext search
         """
-        ### inspiration taken from flask-WhooshAlchemy
-        # find out all entities in join
-        entities = set()
-        # directly queried entities
-        for cd in self.column_descriptions:
-            entities.add(cd['type'])
-        # joined entities
-        if self._join_entities and isinstance(self._join_entities[0], Mapper):
-            # SQLAlchemy >= 0.8.0
-            entities.update(set([x.entity for x in self._join_entities]))
-        else:
-            # SQLAlchemy < 0.8.0
-            entities.update(set(self._join_entities))
+        if not whoosheer:
+            ### inspiration taken from flask-WhooshAlchemy
+            # find out all entities in join
+            entities = set()
+            # directly queried entities
+            for cd in self.column_descriptions:
+                entities.add(cd['type'])
+            # joined entities
+            if self._join_entities and isinstance(self._join_entities[0], Mapper):
+                # SQLAlchemy >= 0.8.0
+                entities.update(set([x.entity for x in self._join_entities]))
+            else:
+                # SQLAlchemy < 0.8.0
+                entities.update(set(self._join_entities))
 
-        whoosheer = next(w for w in Whooshee.whoosheers if set(w.models) == entities)
+            whoosheer = next(w for w in Whooshee.whoosheers if set(w.models) == entities)
 
         # TODO what if unique field doesn't exist or there are multiple?
         for fname, field in list(whoosheer.schema._fields.items()):
diff --git a/test.py b/test.py
index b8bc492..efef8f3 100644
--- a/test.py
+++ b/test.py
@@ -76,6 +76,7 @@ class BaseTestCases(object):
                                         title=entry.title,
                                         content=entry.content)
 
+
             self.User = User
             self.Entry = Entry
             self.EntryUserWhoosheer = EntryUserWhoosheer
@@ -189,6 +190,35 @@ class BaseTestCases(object):
             titles = [int(entry.title) for entry in found_entries]
             self.assertEqual(titles, sorted(titles))
 
+        def test_whoosheer_search_option(self):
+
+            # alternative whoosheer
+            @self.wh.register_whoosheer
+            class EntryWhoosheer(AbstractWhoosheer):
+                schema = whoosh.fields.Schema(
+                    entry_id = whoosh.fields.NUMERIC(stored=True, unique=True),
+                    title = whoosh.fields.TEXT()
+                )
+
+                models = [self.Entry]
+
+                @classmethod
+                def update_entry(cls, writer, entry):
+                    writer.update_document(entry_id=entry.id, title=entry.title+'cookie')
+
+                @classmethod
+                def insert_entry(cls, writer, entry):
+                    writer.add_document(entry_id=entry.id, title=entry.title+'cookie')
+
+            entry = self.Entry(title=u'secret_', content=u'blah blah blah', user=self.u1)
+            self.db.session.add(entry)
+            self.db.session.commit()
+
+            found = self.Entry.query.join(self.User).whooshee_search('secret_cookie').all()
+            self.assertEqual(len(found), 0)
+            found = self.Entry.query.join(self.User).whooshee_search('secret_cookie', whoosheer=EntryWhoosheer).all()
+            self.assertEqual(len(found), 1)
+
         def test_reindex(self):
             self.db.session.add_all(self.all_inst)
             self.db.session.commit()