Add support for user/password auth

Cherry-pick of trinodb/trino#4165

Co-authored-by: Martin Traverso <mtraverso@gmail.com>

presto-elasticsearch/pom.xml

@@ -243,6 +243,30 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>net.java.dev.jna</groupId>
+            <artifactId>jna</artifactId>
+            <version>5.5.0</version>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>org.testcontainers</groupId>
+            <artifactId>testcontainers</artifactId>
+            <version>1.15.2</version>
+            <scope>test</scope>
+            <exclusions>
+                <exclusion>
+                    <groupId>org.slf4j</groupId>
+                    <artifactId>slf4j-api</artifactId>
+                </exclusion>
+                <exclusion>
+                    <groupId>com.fasterxml.jackson.core</groupId>
+                    <artifactId>jackson-annotations</artifactId>
+                </exclusion>
+            </exclusions>
+        </dependency>
+
         <dependency>
             <groupId>org.testcontainers</groupId>
             <artifactId>elasticsearch</artifactId>
@@ -256,10 +280,6 @@
                     <groupId>org.apache.commons</groupId>
                     <artifactId>commons-compress</artifactId>
                 </exclusion>
-                <exclusion>
-                    <groupId>net.java.dev.jna</groupId>
-                    <artifactId>jna</artifactId>
-                </exclusion>
             </exclusions>
         </dependency>
 

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchConfig.java

@@ -32,7 +32,8 @@ public class ElasticsearchConfig
 {
     public enum Security
     {
-        AWS
+        AWS,
+        PASSWORD
     }
 
     private String host;

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchConnectorModule.java

@@ -30,6 +30,7 @@
 import static com.facebook.airlift.json.JsonBinder.jsonBinder;
 import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
 import static com.facebook.presto.elasticsearch.ElasticsearchConfig.Security.AWS;
+import static com.facebook.presto.elasticsearch.ElasticsearchConfig.Security.PASSWORD;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static java.util.Objects.requireNonNull;
@@ -55,13 +56,21 @@ protected void setup(Binder binder)
         binder.install(new DecoderModule());
 
         newOptionalBinder(binder, AwsSecurityConfig.class);
+        newOptionalBinder(binder, PasswordConfig.class);
 
         install(installModuleIf(
                 ElasticsearchConfig.class,
                 config -> config.getSecurity()
                         .filter(isEqual(AWS))
                         .isPresent(),
                 conditionalBinder -> configBinder(conditionalBinder).bindConfig(AwsSecurityConfig.class)));
+
+        install(installModuleIf(
+                ElasticsearchConfig.class,
+                config -> config.getSecurity()
+                        .filter(isEqual(PASSWORD))
+                        .isPresent(),
+                conditionalBinder -> configBinder(conditionalBinder).bindConfig(PasswordConfig.class)));
     }
 
     private static final class TypeDeserializer

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/PasswordConfig.java

@@ -0,0 +1,52 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.elasticsearch;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigSecuritySensitive;
+
+import javax.validation.constraints.NotNull;
+
+public class PasswordConfig
+{
+    private String user;
+    private String password;
+
+    @NotNull
+    public String getUser()
+    {
+        return user;
+    }
+
+    @Config("elasticsearch.auth.user")
+    public PasswordConfig setUser(String user)
+    {
+        this.user = user;
+        return this;
+    }
+
+    @NotNull
+    public String getPassword()
+    {
+        return password;
+    }
+
+    @Config("elasticsearch.auth.password")
+    @ConfigSecuritySensitive
+    public PasswordConfig setPassword(String password)
+    {
+        this.password = password;
+        return this;
+    }
+}

presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/client/ElasticsearchClient.java

@@ -24,6 +24,7 @@
 import com.facebook.airlift.security.pem.PemReader;
 import com.facebook.presto.elasticsearch.AwsSecurityConfig;
 import com.facebook.presto.elasticsearch.ElasticsearchConfig;
+import com.facebook.presto.elasticsearch.PasswordConfig;
 import com.facebook.presto.spi.PrestoException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -35,10 +36,14 @@
 import io.airlift.units.Duration;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpHost;
+import org.apache.http.auth.AuthScope;
+import org.apache.http.auth.UsernamePasswordCredentials;
+import org.apache.http.client.CredentialsProvider;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.conn.ssl.NoopHostnameVerifier;
 import org.apache.http.entity.ByteArrayEntity;
 import org.apache.http.entity.StringEntity;
+import org.apache.http.impl.client.BasicCredentialsProvider;
 import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
 import org.apache.http.impl.nio.reactor.IOReactorConfig;
 import org.apache.http.message.BasicHeader;
@@ -129,9 +134,12 @@ public class ElasticsearchClient
     private final boolean ignorePublishAddress;
 
     @Inject
-    public ElasticsearchClient(ElasticsearchConfig config, Optional<AwsSecurityConfig> awsSecurityConfig)
+    public ElasticsearchClient(
+            ElasticsearchConfig config,
+            Optional<AwsSecurityConfig> awsSecurityConfig,
+            Optional<PasswordConfig> passwordConfig)
     {
-        client = createClient(config, awsSecurityConfig);
+        client = createClient(config, awsSecurityConfig, passwordConfig);
 
         this.ignorePublishAddress = config.isIgnorePublishAddress();
         this.scrollSize = config.getScrollSize();
@@ -183,7 +191,10 @@ private void refreshNodes()
         }
     }
 
-    private static RestHighLevelClient createClient(ElasticsearchConfig config, Optional<AwsSecurityConfig> awsSecurityConfig)
+    private static RestHighLevelClient createClient(
+            ElasticsearchConfig config,
+            Optional<AwsSecurityConfig> awsSecurityConfig,
+            Optional<PasswordConfig> passwordConfig)
     {
         RestClientBuilder builder = RestClient.builder(
                 new HttpHost(config.getHost(), config.getPort(), config.isTlsEnabled() ? "https" : "http"))
@@ -216,6 +227,12 @@ private static RestHighLevelClient createClient(ElasticsearchConfig config, Opti
                 }
             }
 
+            passwordConfig.ifPresent(securityConfig -> {
+                CredentialsProvider credentials = new BasicCredentialsProvider();
+                credentials.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(securityConfig.getUser(), securityConfig.getPassword()));
+                clientBuilder.setDefaultCredentialsProvider(credentials);
+            });
+
             awsSecurityConfig.ifPresent(securityConfig -> clientBuilder.addInterceptorLast(new AwsRequestSigner(
                     securityConfig.getRegion(),
                     getAwsCredentialsProvider(securityConfig))));

presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchConnectorTest.java

@@ -13,25 +13,47 @@
  */
 package com.facebook.presto.elasticsearch;
 
+import com.facebook.presto.testing.MaterializedResult;
+import com.facebook.presto.testing.MaterializedRow;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestIntegrationSmokeTest;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.net.HostAndPort;
 import io.airlift.tpch.TpchTable;
-import org.elasticsearch.action.admin.indices.alias.IndicesAliasesRequest;
+import org.apache.http.HttpHost;
+import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.RestHighLevelClient;
 import org.intellij.lang.annotations.Language;
+import org.testng.annotations.AfterClass;
 import org.testng.annotations.Test;
 
+import java.io.IOException;
+
+import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.elasticsearch.ElasticsearchQueryRunner.createElasticsearchQueryRunner;
-import static com.facebook.presto.elasticsearch.EmbeddedElasticsearchNode.createEmbeddedElasticsearchNode;
-import static org.elasticsearch.client.Requests.indexAliasesRequest;
-import static org.elasticsearch.client.Requests.refreshRequest;
+import static com.facebook.presto.testing.MaterializedResult.resultBuilder;
+import static com.facebook.presto.testing.assertions.Assert.assertEquals;
+import static java.lang.String.format;
 
+@Test(singleThreaded = true)
 public class ElasticsearchConnectorTest
         extends AbstractTestIntegrationSmokeTest
 {
-    private EmbeddedElasticsearchNode embeddedElasticsearchNode;
+    private final String elasticsearchServer = "docker.elastic.co/elasticsearch/elasticsearch-oss:6.0.0";
+    private ElasticsearchServer elasticsearch;
+    private RestHighLevelClient client;
+
+    @AfterClass(alwaysRun = true)
+    public final void destroy()
+            throws IOException
+    {
+        elasticsearch.stop();
+        client.close();
+    }
 
     @Test
     public void testSelectInformationSchemaForMultiIndexAlias()
+            throws IOException
     {
         addAlias("nation", "multi_alias");
         addAlias("region", "multi_alias");
@@ -132,25 +154,71 @@ public void testSelectInformationSchemaColumns()
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        embeddedElasticsearchNode = createEmbeddedElasticsearchNode();
-        return createElasticsearchQueryRunner(embeddedElasticsearchNode, TpchTable.getTables());
+        elasticsearch = new ElasticsearchServer(elasticsearchServer, ImmutableMap.of());
+        HostAndPort address = elasticsearch.getAddress();
+        client = new RestHighLevelClient(RestClient.builder(new HttpHost(address.getHost(), address.getPort())));
+
+        return createElasticsearchQueryRunner(elasticsearch.getAddress(),
+                TpchTable.getTables(),
+                ImmutableMap.of(),
+                ImmutableMap.of());
+    }
+
+    @Test
+    @Override
+    public void testDescribeTable()
+    {
+        MaterializedResult actualColumns = computeActual("DESC orders").toTestTypes();
+        MaterializedResult.Builder builder = resultBuilder(getQueryRunner().getDefaultSession(), VARCHAR, VARCHAR, VARCHAR, VARCHAR);
+        for (MaterializedRow row : actualColumns.getMaterializedRows()) {
+            builder.row(row.getField(0), row.getField(1), "", "");
+        }
+        MaterializedResult actualResult = builder.build();
+        builder = resultBuilder(getQueryRunner().getDefaultSession(), VARCHAR, VARCHAR, VARCHAR, VARCHAR);
+        MaterializedResult expectedColumns = builder
+                .row("clerk", "varchar", "", "")
+                .row("comment", "varchar", "", "")
+                .row("custkey", "bigint", "", "")
+                .row("orderdate", "timestamp", "", "")
+                .row("orderkey", "bigint", "", "")
+                .row("orderpriority", "varchar", "", "")
+                .row("orderstatus", "varchar", "", "")
+                .row("shippriority", "bigint", "", "")
+                .row("totalprice", "real", "", "")
+                .build();
+        assertEquals(actualResult, expectedColumns, format("%s != %s", actualResult, expectedColumns));
+    }
+
+    @Test
+    public void testMultipleRangesPredicate()
+    {
+        assertQuery("" +
+                "SELECT orderkey, custkey, orderstatus, totalprice, orderdate, orderpriority, clerk, shippriority, comment " +
+                "FROM orders " +
+                "WHERE orderkey BETWEEN 10 AND 50 OR orderkey BETWEEN 100 AND 150");
+    }
+
+    @Test
+    public void testRangePredicate()
+    {
+        // List columns explicitly, as there's no defined order in Elasticsearch
+        assertQuery("" +
+                "SELECT orderkey, custkey, orderstatus, totalprice, orderdate, orderpriority, clerk, shippriority, comment " +
+                "FROM orders " +
+                "WHERE orderkey BETWEEN 10 AND 50");
+    }
+
+    @Test
+    public void testSelectAll()
+    {
+        // List columns explicitly, as there's no defined order in Elasticsearch
+        assertQuery("SELECT orderkey, custkey, orderstatus, totalprice, orderdate, orderpriority, clerk, shippriority, comment  FROM orders");
     }
 
     private void addAlias(String index, String alias)
+            throws IOException
     {
-        embeddedElasticsearchNode.getClient()
-                .admin()
-                .indices()
-                .aliases(indexAliasesRequest()
-                        .addAliasAction(IndicesAliasesRequest.AliasActions.add()
-                                .index(index)
-                                .alias(alias)))
-                .actionGet();
-
-        embeddedElasticsearchNode.getClient()
-                .admin()
-                .indices()
-                .refresh(refreshRequest(alias))
-                .actionGet();
+        client.getLowLevelClient()
+                .performRequest("PUT", format("/%s/_alias/%s", index, alias));
     }
 }

presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java

@@ -24,7 +24,6 @@
 import com.facebook.presto.tests.ResultsSession;
 import org.elasticsearch.action.bulk.BulkRequest;
 import org.elasticsearch.action.index.IndexRequest;
-import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.client.RestHighLevelClient;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 
@@ -43,6 +42,7 @@
 import static com.facebook.presto.common.type.Varchars.isVarcharType;
 import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
+import static org.elasticsearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE;
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
 
 public class ElasticsearchLoader
@@ -108,7 +108,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data)
                     throw new UncheckedIOException("Error loading data into Elasticsearch index: " + tableName, e);
                 }
             }
-            request.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+            request.setRefreshPolicy(IMMEDIATE);
             try {
                 restClient.bulk(request);
             }