/*- * #%L * athena-elasticsearch * %% * Copyright (C) 2019 - 2020 Amazon Web Services * %% * 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. * #L% */ package com.amazonaws.connectors.athena.elasticsearch; import com.amazonaws.auth.DefaultAWSCredentialsProviderChain; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.util.regex.Matcher; import java.util.regex.Pattern; /** * This class creates a REST client injected with an endpoint and credentials. */ public class AwsRestHighLevelClientFactory { private static final Logger logger = LoggerFactory.getLogger(AwsRestHighLevelClientFactory.class); /** * credentialsPattern is used for parsing out the username/password credentials out of the endpoint with the * format: "username@password:" (e.g. if the endpoint is http://myusername@mypassword:www.google.com, then this * pattern will be used to extract the following credentials string: "myusername@mypassword:"). */ private static final Pattern credentialsPattern = Pattern.compile("[^/@]+@[^:]+:"); /** * A flag indicating how the created clients will authenticate to an Elasticsearch instance. * If useAwsCredentials = true, the client will be injected with AWS credentials. * If useAwsCredentials = false and username/password are extracted using the credentialsPattern, the client * will be injected with username/password credentials. * If useAwsCredentials = false and username/password are not present, the client will be created with no * credentials. */ private final boolean useAwsCredentials; /** * Stored cache of clients accessible via the endpoint. */ private final CacheableAwsRestHighLevelClient clientCache = new CacheableAwsRestHighLevelClient(); /** * Constructs a new client factory (using a builder) that will create clients injected with credentials. * @param useAwsCredentials if true the factory create clients injected with AWC credentials. * If false, it will create clients injected with username/password credentials. */ public AwsRestHighLevelClientFactory(boolean useAwsCredentials) { this.useAwsCredentials = useAwsCredentials; } /** * Attempts to get a REST client from the client cache synchronously using the endpoint as a key. If the client * does not exist in the cache, createClient() is called to create a new client. The newly created client is * then inserted into the cache and stored for later use. * @param endpoint is the Elasticsearch instance endpoint. * @return an Elasticsearch REST client. */ public synchronized AwsRestHighLevelClient getOrCreateClient(String endpoint) { AwsRestHighLevelClient client = clientCache.get(endpoint); if (client == null) { client = createClient(endpoint); clientCache.put(endpoint, client); } return client; } /** * Creates a new Elasticsearch REST client. If useAwsCredentials = true, the client is injected with AWS * credentials. If useAwsCredentials = false and username/password are extracted using the credentialsPattern, * the client is injected with username/password credentials. Otherwise a default client with no credentials is * returned. * @param endpoint is the Elasticsearch instance endpoint. The latter may contain username/password credentials * for Elasticsearch services that are external to Amazon. * Examples: * 1) https://search-movies-ne3yqu.us-east-1.es.amazonaws.com * 2) http://myusername@mypassword:www.google.com * @return an Elasticsearch REST client. */ private AwsRestHighLevelClient createClient(String endpoint) { if (useAwsCredentials) { return new AwsRestHighLevelClient.Builder(endpoint) .withCredentials(new DefaultAWSCredentialsProviderChain()).build(); } else { Matcher credentials = credentialsPattern.matcher(endpoint); if (credentials.find()) { String usernameAndPassword = credentials.group(); String username = usernameAndPassword.substring(0, usernameAndPassword.indexOf("@")); String password = usernameAndPassword.substring(usernameAndPassword.indexOf("@") + 1, usernameAndPassword.lastIndexOf(":")); String finalEndpoint = endpoint.replace(usernameAndPassword, ""); return new AwsRestHighLevelClient.Builder(finalEndpoint).withCredentials(username, password).build(); } } logger.debug("Default client w/o credentials"); // Default client w/o credentials. return new AwsRestHighLevelClient.Builder(endpoint).build(); } }