From 930833324de095eb7e3bc4e9df5f75da60ee3cd2 Mon Sep 17 00:00:00 2001
From: Karl Hobley <karl@torchbox.com>
Date: Mon, 14 Jul 2014 10:50:30 +0100
Subject: [PATCH] Fixed backwards compatibility with Elasticsearch URLs setting

---
 .../wagtailsearch/backends/elasticsearch.py   | 23 +++++++--
 .../tests/test_elasticsearch_backend.py       | 51 +++++++++++++++++++
 2 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/wagtail/wagtailsearch/backends/elasticsearch.py b/wagtail/wagtailsearch/backends/elasticsearch.py
index 55f1971a2..8f8a78dbb 100644
--- a/wagtail/wagtailsearch/backends/elasticsearch.py
+++ b/wagtail/wagtailsearch/backends/elasticsearch.py
@@ -2,6 +2,8 @@ from __future__ import absolute_import
 
 import json
 
+from six.moves.urllib.parse import urlparse
+
 from django.db import models
 from django.db.models.sql.where import SubqueryConstraint
 
@@ -446,17 +448,30 @@ class ElasticSearch(BaseSearch):
         super(ElasticSearch, self).__init__(params)
 
         # Get settings
+        self.es_hosts = params.pop('HOSTS', None)
         self.es_urls = params.pop('URLS', ['http://localhost:9200'])
         self.es_index = params.pop('INDEX', 'wagtail')
-        self.es_timeout = params.pop('TIMEOUT', 5)
-        self.es_force_new = params.pop('FORCE_NEW', False)
+        self.es_timeout = params.pop('TIMEOUT', 10)
+
+        # If HOSTS is not set, convert URLS setting to HOSTS
+        if self.es_hosts is None:
+            self.es_hosts = []
+
+            for url in self.es_urls:
+                parsed_url = urlparse(url)
+
+                self.es_hosts.append({
+                    'host': parsed_url.hostname,
+                    'port': parsed_url.port or 9200,
+                    'url_prefix': parsed_url.path,
+                    'use_ssl': parsed_url.scheme == 'https',
+                })
 
         # Get ElasticSearch interface
         # Any remaining params are passed into the ElasticSearch constructor
         self.es = Elasticsearch(
-            urls=self.es_urls,
+            hosts=self.es_hosts,
             timeout=self.es_timeout,
-            force_new=self.es_force_new,
             **params)
 
     def reset_index(self):
diff --git a/wagtail/wagtailsearch/tests/test_elasticsearch_backend.py b/wagtail/wagtailsearch/tests/test_elasticsearch_backend.py
index a95d442ff..3ddce5618 100644
--- a/wagtail/wagtailsearch/tests/test_elasticsearch_backend.py
+++ b/wagtail/wagtailsearch/tests/test_elasticsearch_backend.py
@@ -397,3 +397,54 @@ class TestElasticSearchMappingInheritance(TestCase):
         }
 
         self.assertDictEqual(document, expected_result)
+
+
+class TestBackendConfiguration(TestCase):
+    def setUp(self):
+        # Import using a try-catch block to prevent crashes if the elasticsearch-py
+        # module is not installed
+        try:
+            from wagtail.wagtailsearch.backends.elasticsearch import ElasticSearch
+        except ImportError:
+            raise unittest.SkipTest("elasticsearch-py not installed")
+
+        self.ElasticSearch = ElasticSearch
+
+    def test_default_settings(self):
+        backend = self.ElasticSearch(params={})
+
+        self.assertEqual(len(backend.es_hosts), 1)
+        self.assertEqual(backend.es_hosts[0]['host'], 'localhost')
+        self.assertEqual(backend.es_hosts[0]['port'], 9200)
+        self.assertEqual(backend.es_hosts[0]['use_ssl'], False)
+
+    def test_hosts(self):
+        # This tests that HOSTS goes to es_hosts
+        backend = self.ElasticSearch(params={
+            'HOSTS': [
+                {
+                    'host': '127.0.0.1',
+                    'port': 9300,
+                    'use_ssl': True,
+                }
+            ]
+        })
+
+        self.assertEqual(len(backend.es_hosts), 1)
+        self.assertEqual(backend.es_hosts[0]['host'], '127.0.0.1')
+        self.assertEqual(backend.es_hosts[0]['port'], 9300)
+        self.assertEqual(backend.es_hosts[0]['use_ssl'], True)
+
+    def test_urls(self):
+        # This test backwards compatibility with old URLS setting
+        backend = self.ElasticSearch(params={
+            'URLS': ['http://localhost:12345', 'https://127.0.0.1:54321'],
+        })
+
+        self.assertEqual(len(backend.es_hosts), 2)
+        self.assertEqual(backend.es_hosts[0]['host'], 'localhost')
+        self.assertEqual(backend.es_hosts[0]['port'], 12345)
+        self.assertEqual(backend.es_hosts[0]['use_ssl'], False)
+        self.assertEqual(backend.es_hosts[1]['host'], '127.0.0.1')
+        self.assertEqual(backend.es_hosts[1]['port'], 54321)
+        self.assertEqual(backend.es_hosts[1]['use_ssl'], True)