Update pyes/queryset.py

pyes/queryset.py

@@ -28,43 +28,50 @@ class DoesNotExist(Exception):
 class MultipleObjectsReturned(Exception):
     pass
 
-def get_es_connection():
-    return ES()
+
+def get_es_connection(es_url):
+    #import pdb;pdb.set_trace()
+    if es_url:
+        return ES(es_url)
+    else:
+        return ES()
+
 
 class ESModel(ElasticSearchModel):
 
-    def __init__(self, index, type):
-        self._index=index
-        self._type=type
-        self.objects = QuerySet(self)
+    def __init__(self, index, type, es_url=None):
+        self._index = index
+        self._type = type
+        self.objects = QuerySet(self, es_url=es_url)
         setattr(self, "DoesNotExist", DoesNotExist)
         setattr(self, "MultipleObjectsReturned", MultipleObjectsReturned)
 
-def generate_model(index, doc_type):
-    MyModel = type(
-           'MyModel',
-           (ElasticSearchModel,),
-           {}
-           )
-    setattr(MyModel, "objects", QuerySet(MyModel, index=index, type=doc_type))
+
+def generate_model(index, doc_type, es_url=None):
+    MyModel = type('MyModel', (ElasticSearchModel,), {})
+
+    setattr(MyModel, "objects", QuerySet(MyModel, index=index, type=doc_type, es_url=es_url))
     setattr(MyModel, "DoesNotExist", DoesNotExist)
     setattr(MyModel, "MultipleObjectsReturned", MultipleObjectsReturned)
     return MyModel
 
+
 class QuerySet(object):
     """
     Represents a lazy database lookup for a set of objects.
     """
-    def __init__(self, model=None, using=None, index=None, type=None):
+    def __init__(self, model=None, using=None, index=None, type=None, es_url=None):
+        self.es_url = es_url
         if model is None and index and type:
-            model = ESModel(index, type)
+            model = ESModel(index, type, es_url=self.es_url)
         self.model = model
+        self.es_url = es_url
+
         # EmptyQuerySet instantiates QuerySet with model as None
-        self._index=index
+        self._index = index
         if using:
             self._index = using
         self._type=type
-
         self._queries = []
         self._filters = []
         self._facets = []
@@ -152,8 +159,7 @@ def _build_search(self):
         return query
 
     def _do_query(self):
-
-        return get_es_connection().search(self._build_search(), indices=self.index, doc_types=self.type)
+        return get_es_connection(self.es_url).search(self._build_search(), indices=self.index, doc_types=self.type)
 
 
     def __len__(self):
@@ -162,6 +168,7 @@ def __len__(self):
         # whilst not messing up any existing iterators against the QuerySet.
         if self._result_cache is None:
             self._result_cache = self._do_query()
+
         return self._result_cache.total
 
     def __iter__(self):
@@ -396,7 +403,7 @@ def get_or_create(self, **kwargs):
             params.update(defaults)
             obj = self.model(**params)
             meta = obj.get_meta()
-            meta.connection = get_es_connection()
+            meta.connection = get_es_connection(self.es_url)
             meta.index=self.index
             meta.type=self.type
             obj.save(force=True)
@@ -437,7 +444,7 @@ def delete(self):
         # and one to delete. Make sure that the discovery of related
         # objects is performed on the same database as the deletion.
         del_query._clear_ordering()
-        get_es_connection().delete_by_query(self._build_query())
+        get_es_connection(self.es_url).delete_by_query(self._build_query())
         # Clear the result cache, in case this QuerySet gets reused.
         self._result_cache = None
 
@@ -447,7 +454,7 @@ def update(self, **kwargs):
         fields to the appropriate values.
         """
         query = self._build_query()
-        connection =get_es_connection()
+        connection = get_es_connection(self.es_url)
         results = connection.search(query, indices=self.index, doc_types=self.type,
                                              model=self.model, scan=True)
         for item in results:
@@ -472,7 +479,7 @@ def values(self, *fields):
         search = self._build_search()
         search.facet.reset()
         search.fields=fields
-        return get_es_connection().search(search, indices=self.index, doc_types=self.type)
+        return get_es_connection(self.es_url).search(search, indices=self.index, doc_types=self.type)
 
     def values_list(self, *fields, **kwargs):
         flat = kwargs.pop('flat', False)
@@ -486,10 +493,10 @@ def values_list(self, *fields, **kwargs):
         search.facet.reset()
         search.fields=fields
         if flat:
-            return get_es_connection().search(search, indices=self.index, doc_types=self.type,
+            return get_es_connection(self.es_url).search(search, indices=self.index, doc_types=self.type,
                                               model=lambda x,y: y.get("fields", {}).get(fields[0], None))
 
-        return get_es_connection().search(search, indices=self.index, doc_types=self.type)
+        return get_es_connection(self.es_url).search(search, indices=self.index, doc_types=self.type)
 
     def dates(self, field_name, kind, order='ASC'):
         """
@@ -507,7 +514,7 @@ def dates(self, field_name, kind, order='ASC'):
         search.facet.add_date_facet(name=field_name.replace("__", "."),
                  field=field_name, interval=kind)
         search.size=0
-        resulset= get_es_connection().search(search, indices=self.index, doc_types=self.type)
+        resulset = get_es_connection(self.es_url).search(search, indices=self.index, doc_types=self.type)
         resulset.fix_facets()
         entries = []
         for val in resulset.facets.get(field_name.replace("__", ".")).get("entries", []):
@@ -779,7 +786,7 @@ def facets(self):
     def _clone(self, klass=None, setup=False, **kwargs):
         if klass is None:
             klass = self.__class__
-        c = klass(model=self.model, using=self.index, index=self.index, type=self.type)
+        c = klass(model=self.model, using=self.index, index=self.index, type=self.type, es_url=self.es_url)
         #copy filters/queries/facet????
         c.__dict__.update(kwargs)
         c._queries=list(self._queries)