Added support for secure Elasticsearch connections

joke2k · Apr 11, 2023 · 44bb096 · 44bb096
1 parent d1494c6
commit 44bb096
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 14 deletions.
diff --git a/docs/types.rst b/docs/types.rst
@@ -156,10 +156,10 @@ For more detailed example see ":ref:`complex_dict_format`".
 
 :py:meth:`~.environ.Env.search_url` supports the following URL schemas:
 
-* Elasticsearch: ``elasticsearch://``
-* Elasticsearch2: ``elasticsearch2://``
-* Elasticsearch5: ``elasticsearch5://``
-* Elasticsearch7: ``elasticsearch7://``
+* Elasticsearch: ``elasticsearch://`` (http) or ``elasticsearchs://`` (https)
+* Elasticsearch2: ``elasticsearch2://`` (http) or ``elasticsearch2s://`` (https)
+* Elasticsearch5: ``elasticsearch5://`` (http) or ``elasticsearch5s://`` (https)
+* Elasticsearch7: ``elasticsearch7://`` (http) or ``elasticsearch7s://`` (https)
 * Solr: ``solr://``
 * Whoosh: ``whoosh://``
 * Xapian: ``xapian://``

diff --git a/environ/environ.py b/environ/environ.py
@@ -113,7 +113,6 @@ class Env:
     URL_CLASS = ParseResult
 
     POSTGRES_FAMILY = ['postgres', 'postgresql', 'psql', 'pgsql', 'postgis']
-    ELASTICSEARCH_FAMILY = ['elasticsearch' + x for x in ['', '2', '5', '7']]
 
     DEFAULT_DATABASE_ENV = 'DATABASE_URL'
     DB_SCHEMES = {
@@ -191,6 +190,8 @@ class Env:
         "xapian": "haystack.backends.xapian_backend.XapianEngine",
         "simple": "haystack.backends.simple_backend.SimpleEngine",
     }
+    ELASTICSEARCH_FAMILY = [scheme + s for scheme in SEARCH_SCHEMES if scheme.startswith("elasticsearch")
+                            for s in ('', 's')]
     CLOUDSQL = 'cloudsql'
 
     def __init__(self, **scheme):
@@ -760,9 +761,15 @@ def search_url_config(cls, url, engine=None):
         path = url.path[1:]
         path = unquote_plus(path.split('?', 2)[0])
 
-        if url.scheme not in cls.SEARCH_SCHEMES:
+        scheme = url.scheme
+        secure = False
+        # elasticsearch supports secure schemes, similar to http -> https
+        if scheme in cls.ELASTICSEARCH_FAMILY and scheme.endswith('s'):
+            scheme = scheme[:-1]
+            secure = True
+        if scheme not in cls.SEARCH_SCHEMES:
             raise ImproperlyConfigured(f'Invalid search schema {url.scheme}')
-        config["ENGINE"] = cls.SEARCH_SCHEMES[url.scheme]
+        config["ENGINE"] = cls.SEARCH_SCHEMES[scheme]
 
         # check commons params
         params = {}
@@ -811,7 +818,7 @@ def search_url_config(cls, url, engine=None):
                 index = split[0]
 
             config['URL'] = urlunparse(
-                ('http',) + url[1:2] + (path,) + ('', '', '')
+                ('https' if secure else 'http', url[1], path, '', '', '')
             )
             if 'TIMEOUT' in params:
                 config['TIMEOUT'] = cls.parse_value(params['TIMEOUT'][0], int)

diff --git a/tests/test_search.py b/tests/test_search.py
@@ -33,25 +33,45 @@ def test_solr_multicore_parsing(solr_url):
 
 
 @pytest.mark.parametrize(
-    'url,engine',
+    'url,engine,scheme',
     [
         ('elasticsearch://127.0.0.1:9200/index',
-         'elasticsearch_backend.ElasticsearchSearchEngine'),
+         'elasticsearch_backend.ElasticsearchSearchEngine',
+         'http',),
+        ('elasticsearchs://127.0.0.1:9200/index',
+         'elasticsearch_backend.ElasticsearchSearchEngine',
+         'https',),
         ('elasticsearch2://127.0.0.1:9200/index',
-         'elasticsearch2_backend.Elasticsearch2SearchEngine'),
+         'elasticsearch2_backend.Elasticsearch2SearchEngine',
+         'http',),
+        ('elasticsearch2s://127.0.0.1:9200/index',
+         'elasticsearch2_backend.Elasticsearch2SearchEngine',
+         'https',),
         ('elasticsearch5://127.0.0.1:9200/index',
-         'elasticsearch5_backend.Elasticsearch5SearchEngine'),
+         'elasticsearch5_backend.Elasticsearch5SearchEngine',
+         'http'),
+        ('elasticsearch5s://127.0.0.1:9200/index',
+         'elasticsearch5_backend.Elasticsearch5SearchEngine',
+         'https'),
         ('elasticsearch7://127.0.0.1:9200/index',
-         'elasticsearch7_backend.Elasticsearch7SearchEngine'),
+         'elasticsearch7_backend.Elasticsearch7SearchEngine',
+         'http'),
+        ('elasticsearch7s://127.0.0.1:9200/index',
+         'elasticsearch7_backend.Elasticsearch7SearchEngine',
+         'https'),
     ],
     ids=[
         'elasticsearch',
+        'elasticsearchs',
         'elasticsearch2',
+        'elasticsearch2s',
         'elasticsearch5',
+        'elasticsearch5s',
         'elasticsearch7',
+        'elasticsearch7s',
     ]
 )
-def test_elasticsearch_parsing(url, engine):
+def test_elasticsearch_parsing(url, engine, scheme):
     """Ensure all supported Elasticsearch engines are recognized."""
     timeout = 360
     url = '{}?TIMEOUT={}'.format(url, timeout)
@@ -63,6 +83,7 @@ def test_elasticsearch_parsing(url, engine):
     assert 'TIMEOUT' in url.keys()
     assert url['TIMEOUT'] == timeout
     assert 'PATH' not in url
+    assert url["URL"].startswith(scheme + ":")
 
 
 @pytest.mark.parametrize('storage', ['file', 'ram'])