From de8c6a8b639d50d06284a8bb1153690ba1270b49 Mon Sep 17 00:00:00 2001
From: Stephan Preeker <stephan@preeker.net>
Date: Fri, 5 Jan 2018 10:56:41 +0100
Subject: [PATCH] improve elastic server settings

---
 pillar/api/search/queries.py | 23 ++++++++++++++++++-----
 pillar/config.py             | 25 ++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/pillar/api/search/queries.py b/pillar/api/search/queries.py
index 0e223c79..cb81e2ed 100644
--- a/pillar/api/search/queries.py
+++ b/pillar/api/search/queries.py
@@ -4,15 +4,28 @@ from elasticsearch_dsl import Search, Q
 import logging
 
 from pillar import current_app
-
-client = Elasticsearch()
-
 log = logging.getLogger(__name__)
 
 node_agg_terms = ['node_type', 'media', 'tags', 'is_free']
 user_agg_terms = ['roles', ]
 
 
+class TheELKClient():
+    """
+    current_app is not available when on import
+    """
+    client: Elasticsearch = None
+
+    def get_client(self):
+        if not self.client:
+            self.client = Elasticsearch(
+                current_app.config['ELASTIC_SEARCH_HOSTS'])
+        else:
+            return self.client
+
+elk = TheELKClient()
+
+
 def add_aggs_to_search(search, agg_terms):
     """
     Add facets / aggregations to the search result
@@ -42,7 +55,7 @@ def nested_bool(must: list, should: list, terms: dict) -> Search:
     must.append(bool_query)
     bool_query = Q('bool', must=must)
 
-    search = Search(using=client)
+    search = Search(using=elk.get_client())
     search.query = bool_query
 
     return search
@@ -129,7 +142,7 @@ def do_user_search_admin(query: str) -> dict:
         Q('match', full_name=query),
     ]
     bool_query = Q('bool', should=should)
-    search = Search(using=client)
+    search = Search(using=elk.get_client())
     search.query = bool_query
 
     if log.isEnabledFor(logging.DEBUG):
diff --git a/pillar/config.py b/pillar/config.py
index 48d9e44e..557d2cc0 100644
--- a/pillar/config.py
+++ b/pillar/config.py
@@ -80,7 +80,7 @@ ALGOLIA_INDEX_NODES = 'dev_Nodes'
 
 SEARCH_BACKENDS = ('algolia', 'elastic')    # search backend we use
 
-ELASTIC_SEARCH_HOSTS = ['elasticsearch']    # elasticsearch hosts
+
 ELASTIC_INDICES = {
     # elasticsearch indexes
     'NODE': 'nodes',
@@ -88,6 +88,29 @@ ELASTIC_INDICES = {
 }
 
 
+def get_docker_host():
+    """
+    Looks for the DOCKER_HOST environment variable to find the VM
+    running docker-machine.
+
+    If the environment variable is not found, it is assumed that
+    you're running docker on localhost.
+    """
+    d_host = os.getenv('DOCKER_HOST', None)
+    if d_host:
+        if re.match(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$', d_host):
+            return d_host
+
+        return re.match(r'tcp://(.*?):\d+', d_host).group(1)
+    return 'localhost'
+
+
+# Elasticsearch hosts
+ELASTIC_SEARCH_HOSTS = ["{}:{}".format(
+    os.getenv('ELASTICSEARCH_PORT_9200_TCP_ADDR', get_docker_host()),
+    os.getenv('ELASTICSEARCH_PORT_9200_TCP_PORT', 9200))]
+
+
 ZENCODER_API_KEY = '-SECRET-'
 ZENCODER_NOTIFICATIONS_SECRET = '-SECRET-'
 ZENCODER_NOTIFICATIONS_URL = 'http://zencoderfetcher/'