From 4144b7aac9b6d85bf328a4c4189fbdf1d69c85dc Mon Sep 17 00:00:00 2001
From: Przemek Kaminski <przemyslaw.kaminski@sensisoft.com>
Date: Wed, 7 Aug 2013 08:44:29 +0200
Subject: [PATCH] MULTISEARCH: add routing_list option

---
 pyes/es.py | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/pyes/es.py b/pyes/es.py
index 16a88c1b..f82ae89b 100644
--- a/pyes/es.py
+++ b/pyes/es.py
@@ -1025,7 +1025,7 @@ def index(self, doc, index, doc_type, id=None, parent=None, force_insert=False,
             if not isinstance(ttl, basestring):
                 ttl = str(ttl)
             querystring_args['ttl'] = ttl
-            
+
         if id is None:
             request_method = 'POST'
         else:
@@ -1252,24 +1252,31 @@ def search_raw(self, query, indices=None, doc_types=None, **query_params):
         path = self._make_path(indices, doc_types, "_search")
         return self._send_request('GET', path, body, params=query_params)
 
-    def search_raw_multi(self, queries, indices_list=None, doc_types_list=None):
+    def search_raw_multi(self, queries, indices_list=None, doc_types_list=None,
+                         routing_list=None):
         if indices_list is None:
             indices_list = [None] * len(queries)
 
         if doc_types_list is None:
             doc_types_list = [None] * len(queries)
 
+        if routing_list is None:
+            routing_list = [None] * len(queries)
+
         queries = [query.search() if isinstance(query, Query)
                    else query.serialize() for query in queries]
         queries = map(self._encode_query, queries)
         headers = []
-        for index_name, doc_type in zip(indices_list,
-                                        doc_types_list):
+        for index_name, doc_type, routing in zip(indices_list,
+                                                 doc_types_list,
+                                                 routing_list):
             d = {}
             if index_name is not None:
                 d['index'] = index_name
             if doc_type is not None:
                 d['type'] = doc_type
+            if routing is not None:
+                d['routing'] = routing
 
             if d:
                 headers.append(d)
@@ -1306,11 +1313,13 @@ def search(self, query, indices=None, doc_types=None, model=None, scan=False, **
         return ResultSet(self, search, indices=indices, doc_types=doc_types,
                          model=model, query_params=query_params)
 
-    def search_multi(self, queries, indices_list=None, doc_types_list=None, models=None, scans=None):
+    def search_multi(self, queries, indices_list=None, doc_types_list=None,
+                     routing_list=None, models=None, scans=None):
         searches = [query if isinstance(query, Search) else Search(query) for query in queries]
 
-        return ResultSetMulti(self, searches, indices_list=indices_list, doc_types_list=doc_types_list,
-                              models=models)
+        return ResultSetMulti(self, searches, indices_list=indices_list,
+                              doc_types_list=doc_types_list,
+                              routing_list=routing_list, models=models)
 
 
     #    scan method is no longer working due to change in ES.search behavior.  May no longer warrant its own method.
@@ -1725,7 +1734,7 @@ def _search_raw(self, start=None, size=None):
 
 class ResultSetMulti(object):
     def __init__(self, connection, searches, indices_list=None,
-                 doc_types_list=None, models=None):
+                 doc_types_list=None, routing_list=None, models=None):
         """
         results: an es query results dict
         fix_keys: remove the "_" from every key, useful for django views
@@ -1749,6 +1758,10 @@ def __init__(self, connection, searches, indices_list=None,
             self.doc_types_list = [None] * num_searches
         else:
             self.doc_types_list = doc_types_list
+        if routing_list is None:
+            self.routing_list = [None] * num_searches
+        else:
+            self.routing_list = routing_list
         self._results_list = None
         self.models = models or self.connection.model
         self._total = None
@@ -1768,9 +1781,11 @@ def _do_search(self):
         if 'responses' in response:
             responses = response['responses']
             self._results_list = [ResultSet(self.connection, search,
-                                            indices=indices, query_params={})
-                for search, indices in
-                    zip(self.searches, self.indices_list)]
+                                            indices=indices, query_params={},
+                                            doc_types=doc_types)
+                for search, indices, doc_types in
+                    zip(self.searches, self.indices_list,
+                        self.doc_types_list)]
 
             for rs, rsp in zip(self._results_list, responses):
                 if 'error' in rsp:
@@ -1788,7 +1803,7 @@ def _do_search(self):
     def _search_raw_multi(self):
         self.multi_search_query, result = self.connection.search_raw_multi(
             self.searches, indices_list=self.indices_list,
-            doc_types_list=self.doc_types_list)
+            doc_types_list=self.doc_types_list, routing_list=self.routing_list)
 
         return result