From d7e1df7166f633e5ba09178dd42fd88580d57893 Mon Sep 17 00:00:00 2001
From: Zhenxiao Luo <zluo@twitter.com>
Date: Tue, 9 Jun 2020 12:26:02 -0700
Subject: [PATCH] Add option to ignore Elasticsearch publish_address

The address is used to address Elasticsearch nodes. When
running in a container environment, the published address
may not match the public address of the container.

This option makes the connector ignore the published address
and use the configured address, instead.

Cherry-pick of https://github.com/prestosql/presto/commit/d04b31c4d91b0365419f175698a4dcb1176b4455

Co-authored-by: Martin Traverso mtraverso@gmail.com
---
 .../presto/elasticsearch/ElasticsearchClient.java   |  4 +++-
 .../presto/elasticsearch/ElasticsearchConfig.java   | 13 +++++++++++++
 .../elasticsearch/TestElasticsearchConfig.java      |  3 +++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchClient.java b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchClient.java
index dd2210964b048..bccd76a55cdb9 100644
--- a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchClient.java
+++ b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchClient.java
@@ -112,12 +112,14 @@ public class ElasticsearchClient
     private final AtomicBoolean started = new AtomicBoolean();
     private final Duration refreshInterval;
     private final boolean tlsEnabled;
+    private final boolean ignorePublishAddress;
 
     @Inject
     public ElasticsearchClient(ElasticsearchConfig config, Optional<AwsSecurityConfig> awsSecurityConfig)
     {
         client = createClient(config, awsSecurityConfig);
 
+        this.ignorePublishAddress = config.isIgnorePublishAddress();
         this.scrollSize = config.getScrollSize();
         this.scrollTimeout = config.getScrollTimeout();
         this.refreshInterval = config.getNodeRefreshInterval();
@@ -155,7 +157,7 @@ private void refreshNodes()
                     .map(address -> HttpHost.create(format("%s://%s", tlsEnabled ? "https" : "http", address)))
                     .toArray(HttpHost[]::new);
 
-            if (hosts.length > 0) {
+            if (hosts.length > 0 && !ignorePublishAddress) {
                 client.getLowLevelClient().setHosts(hosts);
             }
             this.nodes.set(nodes);
diff --git a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchConfig.java b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchConfig.java
index 929f39aa70081..a86fa259eb55c 100644
--- a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchConfig.java
+++ b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchConfig.java
@@ -51,6 +51,7 @@ public enum Security
     private File trustStorePath;
     private String keystorePassword;
     private String truststorePassword;
+    private boolean ignorePublishAddress;
     private boolean verifyHostnames = true;
     private Security security;
 
@@ -268,6 +269,18 @@ public ElasticsearchConfig setVerifyHostnames(boolean verify)
         return this;
     }
 
+    public boolean isIgnorePublishAddress()
+    {
+        return ignorePublishAddress;
+    }
+
+    @Config("elasticsearch.ignore-publish-address")
+    public ElasticsearchConfig setIgnorePublishAddress(boolean ignorePublishAddress)
+    {
+        this.ignorePublishAddress = ignorePublishAddress;
+        return this;
+    }
+
     @NotNull
     public Optional<Security> getSecurity()
     {
diff --git a/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/TestElasticsearchConfig.java b/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/TestElasticsearchConfig.java
index 593cc7db82a0f..9f9349bf70d02 100644
--- a/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/TestElasticsearchConfig.java
+++ b/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/TestElasticsearchConfig.java
@@ -49,6 +49,7 @@ public void testDefaults()
                 .setTrustStorePath(null)
                 .setTruststorePassword(null)
                 .setVerifyHostnames(true)
+                .setIgnorePublishAddress(false)
                 .setSecurity(null));
     }
 
@@ -72,6 +73,7 @@ public void testExplicitPropertyMappings()
                 .put("elasticsearch.tls.truststore-path", "/tmp/truststore")
                 .put("elasticsearch.tls.truststore-password", "truststore-password")
                 .put("elasticsearch.tls.verify-hostnames", "false")
+                .put("elasticsearch.ignore-publish-address", "true")
                 .put("elasticsearch.security", "AWS")
                 .build();
 
@@ -92,6 +94,7 @@ public void testExplicitPropertyMappings()
                 .setTrustStorePath(new File("/tmp/truststore"))
                 .setTruststorePassword("truststore-password")
                 .setVerifyHostnames(false)
+                .setIgnorePublishAddress(true)
                 .setSecurity(AWS);
 
         assertFullMapping(properties, expected);