diff --git a/esrally/mechanic/supplier.py b/esrally/mechanic/supplier.py index 22f8db430..33b8eafd5 100644 --- a/esrally/mechanic/supplier.py +++ b/esrally/mechanic/supplier.py @@ -509,7 +509,7 @@ def __init__(self, repo, version, distributions_root): def fetch(self): io.ensure_dir(self.distributions_root) - download_url = self.repo.download_url + download_url = net.add_url_param_elastic_no_kpi(self.repo.download_url) distribution_path = os.path.join(self.distributions_root, self.repo.file_name) self.logger.info("Resolved download URL [%s] for version [%s]", download_url, self.version) if not os.path.isfile(distribution_path) or not self.repo.cache: diff --git a/esrally/utils/net.py b/esrally/utils/net.py index 447b82f33..3099006aa 100644 --- a/esrally/utils/net.py +++ b/esrally/utils/net.py @@ -19,7 +19,7 @@ import os import socket import urllib.error -from urllib.parse import quote +from urllib.parse import quote, parse_qs, urlencode, urlparse, urlunparse import certifi import urllib3 @@ -181,6 +181,17 @@ def download_http(url, local_path, expected_size_in_bytes=None, progress_indicat progress_indicator(bytes_read, size_from_content_header) return expected_size_in_bytes +def add_url_param_elastic_no_kpi(url): + return _add_url_param(url, {"x-elastic-no-kpi": "true"}) + + +def _add_url_param(url, params): + url_parsed = urlparse(url) + query = parse_qs(url_parsed.query) + query.update(params) + return urlunparse((url_parsed.scheme, url_parsed.netloc, url_parsed.path, url_parsed.params, + urlencode(query, doseq=True), url_parsed.fragment)) + def download(url, local_path, expected_size_in_bytes=None, progress_indicator=None): """ diff --git a/tests/utils/net_test.py b/tests/utils/net_test.py index a1730a5ff..998c818f9 100644 --- a/tests/utils/net_test.py +++ b/tests/utils/net_test.py @@ -60,6 +60,18 @@ def test_gcs_object_url(self, seed): assert net._build_gcs_object_url(bucket_name, bucket_path) == \ "https://storage.googleapis.com/storage/v1/b/unittest-bucket.test.me/o/path%2Fto%2Fobject?alt=media" + def test_add_url_param_elastic_no_kpi(self): + url = "https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.2.0.tar.gz" + assert net.add_url_param_elastic_no_kpi(url) == \ + "https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.2.0.tar.gz?x-elastic-no-kpi=true" + + def test_add_url_param_encoding_and_update(self): + url = "https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.2.0.tar.gz?flag1=true" + params = {"flag1": "test me", "flag2": "test@me"} + # pylint: disable=protected-access + assert net._add_url_param(url, params) == \ + "https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.2.0.tar.gz?flag1=test+me&flag2=test%40me" + def test_progress(self): progress = net.Progress("test") mock_progress = mock.Mock()