From 6c0893b48731e885986d3ee747e233cc959dab54 Mon Sep 17 00:00:00 2001
From: Dave Dash <dd+github@davedash.com>
Date: Wed, 13 Jul 2011 11:29:22 -0700
Subject: [PATCH] A generic wrapper to require ES for a view.

---
 elasticutils/__init__.py | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/elasticutils/__init__.py b/elasticutils/__init__.py
index 079f3b4..4d3b892 100644
--- a/elasticutils/__init__.py
+++ b/elasticutils/__init__.py
@@ -2,7 +2,7 @@ import logging
 from functools import wraps
 from threading import local
 
-from pyes import ES
+from pyes import ES, exceptions
 
 try:
     from django.conf import settings
@@ -18,8 +18,9 @@ def get_es():
     """Return one es object."""
     if not hasattr(_local, 'es'):
         timeout = getattr(settings, 'ES_TIMEOUT', 1)
+        dump = getattr(settings, 'ES_DUMP_CURL', False)
         _local.es = ES(settings.ES_HOSTS, default_indexes=[settings.ES_INDEX],
-                       timeout=timeout)
+                       timeout=timeout, dump_curl=dump)
     return _local.es
 
 
@@ -34,6 +35,40 @@ def es_required(f):
     return wrapper
 
 
+def es_required_or_50x(disabled_msg, error_msg):
+    """
+    This takes a Django view that requires ElasticSearch.
+
+    If `ES_DISABLED` is `True` then we raise a 501 Not Implemented and display
+    the disabled_msg.  If we try the view and an ElasticSearch exception is
+    raised we raise a 503 error with the error_msg.
+
+    We use user-supplied templates in elasticutils/501.html and
+    elasticutils/503.html.
+    """
+    def wrap(f):
+        @wraps(f)
+        def wrapper(request, *args, **kw):
+            from django.shortcuts import render
+            if settings.ES_DISABLED:
+                response = render(request, 'elasticutils/501.html',
+                                  {'msg': disabled_msg})
+                response.status_code = 501
+                return response
+            else:
+                try:
+                    return f(request, *args, **kw)
+                except exceptions.ElasticSearchException as error:
+                    response = render(request, 'elasticutils/503.html',
+                            {'msg': error_msg, 'error': error})
+                    response.status_code = 503
+                    return response
+
+        return wrapper
+
+    return wrap
+
+
 def _process_filters(filters):
     qfilters = []
     for field, value in filters.iteritems():
@@ -95,6 +130,7 @@ class S(object):
 
     def _clone(self):
         new = self.__class__(self.type)
+        new.query = self.query
         new.filter_ = self.filter_
         new.results = list(self.results)
         new.facets = dict(self.facets)