Upgrade Elasticsearch client to 6.2.0

Closes #395 Relates #462
elastic · Apr 10, 2018 · 101a2e6 · 101a2e6
1 parent 2b6b012
commit 101a2e6
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 23 deletions.
diff --git a/esrally/client.py b/esrally/client.py
@@ -20,19 +20,39 @@ def __init__(self, hosts, client_options):
         logger.info("Creating ES client connected to %s with options [%s]", hosts, masked_client_options)
         self.hosts = hosts
         self.client_options = client_options
+        self.ssl_context = None
+
+        # we're using an SSL context now and it is not allowed to have use_ssl present in client options anymore
+        if client_options.pop("use_ssl", False):
+            import ssl
+            from elasticsearch.connection import create_ssl_context
+            logger.info("SSL support: on")
+            client_options["scheme"] = "https"
+
+            self.ssl_context = create_ssl_context(cafile=client_options.pop("ca_certs", certifi.where()))
+
+            if not client_options.pop("verify_certs", True):
+                logger.info("SSL certificate verification: off")
+                self.ssl_context.check_hostname = False
+                self.ssl_context.verify_mode = ssl.CERT_NONE
+
+                logger.warning("User has enabled SSL but disabled certificate verification. This is dangerous but may be ok for a "
+                               "benchmark. Disabling urllib warnings now to avoid a logging storm. "
+                               "See https://urllib3.readthedocs.io/en/latest/advanced-usage.html#ssl-warnings for details.")
+                # disable:  "InsecureRequestWarning: Unverified HTTPS request is being made. Adding certificate verification is strongly \
+                # advised. See: https://urllib3.readthedocs.io/en/latest/advanced-usage.html#ssl-warnings"
+                urllib3.disable_warnings()
+            else:
+                logger.info("SSL certificate verification: on")
+        else:
+            logger.info("SSL support: off")
+            client_options["scheme"] = "http"
 
-        if self._is_set(client_options, "use_ssl") and self._is_set(client_options, "verify_certs") and "ca_certs" not in client_options:
-            self.client_options["ca_certs"] = certifi.where()
-        elif self._is_set(client_options, "use_ssl") and not self._is_set(client_options, "verify_certs"):
-            logger.warning("User has enabled SSL but disabled certificate verification. This is dangerous but may be ok for a benchmark. "
-                           "Disabling urllib warnings now to avoid a logging storm. "
-                           "See https://urllib3.readthedocs.io/en/latest/advanced-usage.html#ssl-warnings for details.")
-            # disable:  "InsecureRequestWarning: Unverified HTTPS request is being made. Adding certificate verification is strongly \
-            # advised. See: https://urllib3.readthedocs.io/en/latest/advanced-usage.html#ssl-warnings"
-            urllib3.disable_warnings()
         if self._is_set(client_options, "basic_auth_user") and self._is_set(client_options, "basic_auth_password"):
-            # Maybe we should remove these keys from the dict?
-            self.client_options["http_auth"] = (client_options["basic_auth_user"], client_options["basic_auth_password"])
+            logger.info("HTTP basic authentication: on")
+            self.client_options["http_auth"] = (client_options.pop("basic_auth_user"), client_options.pop("basic_auth_password"))
+        else:
+            logger.info("HTTP basic authentication: off")
 
     def _is_set(self, client_opts, k):
         try:
@@ -60,8 +80,12 @@ class ConfigurableHttpConnection(elasticsearch.Urllib3HttpConnection):
             def __init__(self, compressed=False, **kwargs):
                 super(ConfigurableHttpConnection, self).__init__(**kwargs)
                 if compressed:
+                    logger.info("HTTP compression: on")
                     self.headers.update(urllib3.make_headers(accept_encoding=True))
                     self.headers.update({"Content-Encoding": "gzip"})
+                else:
+                    logger.info("HTTP compression: off")
                 self.pool = PoolWrap(self.pool, **kwargs)
 
-        return elasticsearch.Elasticsearch(hosts=self.hosts, connection_class=ConfigurableHttpConnection, **self.client_options)
+        return elasticsearch.Elasticsearch(hosts=self.hosts, connection_class=ConfigurableHttpConnection,
+                                           ssl_context=self.ssl_context, **self.client_options)
diff --git a/esrally/metrics.py b/esrally/metrics.py
@@ -132,18 +132,24 @@ def __init__(self, cfg):
         password = self._config.opts("reporting", "datastore.password")
         verify = self._config.opts("reporting", "datastore.ssl.verification_mode", default_value="full", mandatory=False) != "none"
         ca_path = self._config.opts("reporting", "datastore.ssl.certificate_authorities", default_value=None, mandatory=False)
-        if ca_path is None and verify:
-            ca_path = certifi.where()
 
+        from esrally import client
+
+        # Instead of duplicating code, we're just adapting the metrics store specific properties to match the regular client options.
+        client_options = {
+            "use_ssl": secure,
+            "verify_certs": verify,
+            "timeout": 120
+        }
+        if ca_path:
+            client_options["ca_certs"] = ca_path
         if user and password:
-            auth = (user, password)
-        else:
-            auth = None
-        logger.info("Creating connection to metrics store at %s:%s" % (host, port))
-        import elasticsearch
-        self._client = elasticsearch.Elasticsearch(hosts=[{"host": host, "port": port}],
-                                                   use_ssl=secure, http_auth=auth, verify_certs=verify, ca_certs=ca_path,
-                                                   timeout=120, request_timeout=120)
+            client_options["basic_auth_user"] = user
+            client_options["basic_auth_password"] = password
+
+        logger.info("Creating connection to metrics store at %s:%s", host, port)
+        factory = client.EsClientFactory(hosts=[{"host": host, "port": port}], client_options=client_options)
+        self._client = factory.create()
 
     def create(self):
         return EsClient(self._client)

diff --git a/setup.py b/setup.py
@@ -27,7 +27,7 @@ def str_from_file(name):
 #
 ################################################################################################
 install_requires = [
-    "elasticsearch==6.0.0",
+    "elasticsearch==6.2.0",
     "psutil==5.4.0",
     "py-cpuinfo==3.2.0",
     "tabulate==0.8.1",