src/main/java/com/googlesource/gerrit/plugins/kafka/rest/KafkaRestClient.java - plugins/events-kafka - Git at Google

 // Copyright (C) 2021 The Android Open Source Project
 //
 // 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.googlesource.gerrit.plugins.kafka.rest;

 import static java.nio.charset.StandardCharsets.UTF_8;

 import com.google.common.base.Function;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.net.MediaType;
 import com.google.common.util.concurrent.AsyncFunction;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.JdkFutureAdapters;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.gerrit.common.Nullable;
 import com.google.inject.Inject;
 import com.google.inject.assistedinject.Assisted;
 import com.googlesource.gerrit.plugins.kafka.config.KafkaProperties;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URLEncoder;
 import java.util.Enumeration;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.stream.Collectors;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpHeaders;
 import org.apache.http.HttpResponse;
 import org.apache.http.HttpStatus;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.config.RequestConfig.Builder;
 import org.apache.http.client.methods.HttpDelete;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.client.methods.HttpRequestBase;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
 import org.apache.http.impl.nio.client.CloseableHttpAsyncClient;
 import org.apache.log4j.Appender;
 import org.apache.log4j.AppenderSkeleton;
 import org.apache.log4j.Level;
 import org.apache.log4j.LogManager;
 import org.apache.log4j.Logger;

 public class KafkaRestClient {
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
   private static final String KAFKA_V2_JSON = "application/vnd.kafka.json.v2+json";
   private static final String KAFKA_V2 = "application/vnd.kafka.v2+json";

   private final HttpHostProxy proxy;
   private final CloseableHttpAsyncClient httpclient;
   private final ExecutorService futureExecutor;
   private final int kafkaRestApiTimeoutMsec;
   private final KafkaProperties configuration;

   private static boolean logConfigured;

   public interface Factory {
     KafkaRestClient create(KafkaProperties configuration);
   }

   @Inject
   public KafkaRestClient(
       HttpHostProxy httpHostProxy,
       @FutureExecutor ExecutorService executor,
       HttpAsyncClientBuilderFactory credentialsFactory,
       @Assisted KafkaProperties configuration)
       throws URISyntaxException {
     proxy = httpHostProxy;
     httpclient = proxy.apply(credentialsFactory.create()).build();
     httpclient.start();
     this.configuration = configuration;
     kafkaRestApiTimeoutMsec = (int) configuration.getRestApiTimeout().toMillis();
     if (configuration.isHttpWireLog()) {
       enableHttpWireLog();
     }
     this.futureExecutor = executor;
   }

   public static void enableHttpWireLog() {
     if (!logConfigured) {
       Logger httpWireLoggger = Logger.getLogger("org.apache.http.wire");
       httpWireLoggger.setLevel(Level.DEBUG);

       @SuppressWarnings("rawtypes")
       Enumeration rootLoggerAppenders = LogManager.getRootLogger().getAllAppenders();
       while (rootLoggerAppenders.hasMoreElements()) {
         Appender logAppender = (Appender) rootLoggerAppenders.nextElement();
         if (logAppender instanceof AppenderSkeleton) {
           ((AppenderSkeleton) logAppender).setThreshold(Level.DEBUG);
         }
         httpWireLoggger.addAppender(logAppender);
       }

       logConfigured = true;
     }
   }

   public ListenableFuture<HttpResponse> execute(HttpRequestBase request, int... expectedStatuses) {
     return Futures.transformAsync(
         listenableFutureOf(httpclient.execute(request, null)),
         (res) -> {
           IOException exc =
               getResponseException(
                   String.format("HTTP %s %s FAILED", request.getMethod(), request.getURI()),
                   res,
                   expectedStatuses);
           if (exc == null) {
             return Futures.immediateFuture(res);
           }
           return Futures.immediateFailedFuture(exc);
         },
         futureExecutor);
   }

   public <I, O> ListenableFuture<O> mapAsync(
       ListenableFuture<I> inputFuture, AsyncFunction<? super I, ? extends O> mapFunction) {
     return Futures.transformAsync(inputFuture, mapFunction, futureExecutor);
   }

   public <I, O> ListenableFuture<O> map(
       ListenableFuture<I> inputFuture, Function<? super I, ? extends O> mapFunction) {
     return Futures.transform(inputFuture, mapFunction, futureExecutor);
   }

   public HttpGet createGetTopic(String topic) {
     HttpGet get = new HttpGet(resolveKafkaRestApiUri("/topics/" + topic));
     get.addHeader(HttpHeaders.ACCEPT, KAFKA_V2);
     get.setConfig(createRequestConfig());
     return get;
   }

   public HttpGet createGetRecords(URI consumerUri) {
     HttpGet get = new HttpGet(consumerUri.resolve(consumerUri.getPath() + "/records"));
     get.addHeader(HttpHeaders.ACCEPT, KAFKA_V2_JSON);
     get.setConfig(createRequestConfig());
     return get;
   }

   public HttpPost createPostToConsumer(String consumerGroup) {
     HttpPost post =
         new HttpPost(
             resolveKafkaRestApiUri("/consumers/" + URLEncoder.encode(consumerGroup, UTF_8)));
     post.addHeader(HttpHeaders.ACCEPT, MediaType.ANY_TYPE.toString());
     post.setConfig(createRequestConfig());
     post.setEntity(
         new StringEntity(
             "{\"format\": \"json\",\"auto.offset.reset\": \"earliest\", \"auto.commit.enable\":\"true\", \"consumer.request.timeout.ms\": \"1000\"}",
             ContentType.create(KAFKA_V2, UTF_8)));
     return post;
   }

   public HttpDelete createDeleteToConsumer(URI consumerUri) {
     HttpDelete delete = new HttpDelete(consumerUri);
     delete.addHeader(HttpHeaders.ACCEPT, "*/*");
     delete.setConfig(createRequestConfig());
     return delete;
   }

   public HttpDelete createDeleteToConsumerSubscriptions(URI consumerUri) {
     URI subscriptionUri = consumerUri.resolve("subscription");
     HttpDelete delete = new HttpDelete(subscriptionUri);
     delete.addHeader(HttpHeaders.ACCEPT, "*/*");
     delete.setConfig(createRequestConfig());
     return delete;
   }

   public HttpPost createPostToSubscribe(URI consumerUri, String topic) {
     HttpPost post = new HttpPost(consumerUri.resolve(consumerUri.getPath() + "/subscription"));
     post.addHeader(HttpHeaders.ACCEPT, "*/*");
     post.setConfig(createRequestConfig());
     post.setEntity(
         new StringEntity(
             String.format("{\"topics\":[\"%s\"]}", topic), ContentType.create(KAFKA_V2, UTF_8)));
     return post;
   }

   public HttpPost createPostToTopic(String topic, HttpEntity postBodyEntity) {
     HttpPost post =
         new HttpPost(resolveKafkaRestApiUri("/topics/" + URLEncoder.encode(topic, UTF_8)));
     post.addHeader(HttpHeaders.ACCEPT, "*/*");
     post.setConfig(createRequestConfig());
     post.setEntity(postBodyEntity);
     return post;
   }

   public HttpPost createPostSeekTopicFromBeginning(
       URI consumerUri, String topic, Set<Integer> partitions) {
     HttpPost post =
         new HttpPost(consumerUri.resolve(consumerUri.getPath() + "/positions/beginning"));
     post.addHeader(HttpHeaders.ACCEPT, "*/*");
     post.setConfig(createRequestConfig());
     post.setEntity(
         new StringEntity(
             String.format(
                 "{\"partitions\":[%s]}",
                 partitions.stream()
                     .map(
                         partition ->
                             String.format("{\"topic\":\"%s\",\"partition\":%d}", topic, partition))
                     .collect(Collectors.joining(","))),
             ContentType.create(KAFKA_V2, UTF_8)));
     return post;
   }

   @Nullable
   public IOException getResponseException(
       String errorMessage, HttpResponse response, int... okHttpStatuses) {
     int responseHttpStatus = response.getStatusLine().getStatusCode();
     if (okHttpStatuses.length == 0) {
       okHttpStatuses =
           new int[] {HttpStatus.SC_OK, HttpStatus.SC_CREATED, HttpStatus.SC_NO_CONTENT};
     }
     for (int httpStatus : okHttpStatuses) {
       if (responseHttpStatus == httpStatus) {
         return null;
       }
     }

     String responseBody = "";
     try {
       responseBody = getStringEntity(response);
     } catch (IOException e) {
       logger.atWarning().withCause(e).log(
           "Unable to extrace the string entity for response %d (%s)",
           response.getStatusLine().getStatusCode(), response.getStatusLine().getReasonPhrase());
     }

     return new IOException(
         String.format(
             "%s\nHTTP status %d (%s)\n%s",
             errorMessage,
             response.getStatusLine().getStatusCode(),
             response.getStatusLine().getReasonPhrase(),
             responseBody));
   }

   protected String getStringEntity(HttpResponse response) throws IOException {
     HttpEntity entity = response.getEntity();
     if (entity != null) {
       try (ByteArrayOutputStream outStream = new ByteArrayOutputStream()) {
         entity.writeTo(outStream);
         outStream.close();
         return outStream.toString(UTF_8);
       }
     }
     return "";
   }

   private <V> ListenableFuture<V> listenableFutureOf(Future<V> future) {
     return JdkFutureAdapters.listenInPoolThread(future, futureExecutor);
   }

   private RequestConfig createRequestConfig() {
     Builder configBuilder =
         RequestConfig.custom()
             .setConnectionRequestTimeout(kafkaRestApiTimeoutMsec)
             .setConnectTimeout(kafkaRestApiTimeoutMsec)
             .setSocketTimeout(kafkaRestApiTimeoutMsec);
     configBuilder = proxy.apply(configBuilder);
     RequestConfig config = configBuilder.build();
     return config;
   }

   public void close() throws IOException {
     httpclient.close();
   }

   public URI resolveKafkaRestApiUri(String path) {
     try {
       URI restApiUri = configuration.getRestApiUri();
       return restApiUri.resolve(path);
     } catch (URISyntaxException e) {
       throw new IllegalArgumentException("Invalid Kafka REST API URI", e);
     }
   }

   public URI resolveKafkaRestApiUri(String kafkaRestId, String path) {
     URI restApiUri;
     try {
       restApiUri = configuration.getRestApiUri(kafkaRestId);
       return restApiUri.resolve(restApiUri.getPath() + path);
     } catch (URISyntaxException e) {
       throw new IllegalArgumentException("Invalid Kafka REST API URI", e);
     }
   }
 }
	// Copyright (C) 2021 The Android Open Source Project
	//
	// 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.googlesource.gerrit.plugins.kafka.rest;

	import static java.nio.charset.StandardCharsets.UTF_8;

	import com.google.common.base.Function;
	import com.google.common.flogger.FluentLogger;
	import com.google.common.net.MediaType;
	import com.google.common.util.concurrent.AsyncFunction;
	import com.google.common.util.concurrent.Futures;
	import com.google.common.util.concurrent.JdkFutureAdapters;
	import com.google.common.util.concurrent.ListenableFuture;
	import com.google.gerrit.common.Nullable;
	import com.google.inject.Inject;
	import com.google.inject.assistedinject.Assisted;
	import com.googlesource.gerrit.plugins.kafka.config.KafkaProperties;
	import java.io.ByteArrayOutputStream;
	import java.io.IOException;
	import java.net.URI;
	import java.net.URISyntaxException;
	import java.net.URLEncoder;
	import java.util.Enumeration;
	import java.util.Set;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Future;
	import java.util.stream.Collectors;
	import org.apache.http.HttpEntity;
	import org.apache.http.HttpHeaders;
	import org.apache.http.HttpResponse;
	import org.apache.http.HttpStatus;
	import org.apache.http.client.config.RequestConfig;
	import org.apache.http.client.config.RequestConfig.Builder;
	import org.apache.http.client.methods.HttpDelete;
	import org.apache.http.client.methods.HttpGet;
	import org.apache.http.client.methods.HttpPost;
	import org.apache.http.client.methods.HttpRequestBase;
	import org.apache.http.entity.ContentType;
	import org.apache.http.entity.StringEntity;
	import org.apache.http.impl.nio.client.CloseableHttpAsyncClient;
	import org.apache.log4j.Appender;
	import org.apache.log4j.AppenderSkeleton;
	import org.apache.log4j.Level;
	import org.apache.log4j.LogManager;
	import org.apache.log4j.Logger;

	public class KafkaRestClient {
	private static final FluentLogger logger = FluentLogger.forEnclosingClass();
	private static final String KAFKA_V2_JSON = "application/vnd.kafka.json.v2+json";
	private static final String KAFKA_V2 = "application/vnd.kafka.v2+json";

	private final HttpHostProxy proxy;
	private final CloseableHttpAsyncClient httpclient;
	private final ExecutorService futureExecutor;
	private final int kafkaRestApiTimeoutMsec;
	private final KafkaProperties configuration;

	private static boolean logConfigured;

	public interface Factory {
	KafkaRestClient create(KafkaProperties configuration);
	}

	@Inject
	public KafkaRestClient(
	HttpHostProxy httpHostProxy,
	@FutureExecutor ExecutorService executor,
	HttpAsyncClientBuilderFactory credentialsFactory,
	@Assisted KafkaProperties configuration)
	throws URISyntaxException {
	proxy = httpHostProxy;
	httpclient = proxy.apply(credentialsFactory.create()).build();
	httpclient.start();
	this.configuration = configuration;
	kafkaRestApiTimeoutMsec = (int) configuration.getRestApiTimeout().toMillis();
	if (configuration.isHttpWireLog()) {
	enableHttpWireLog();
	}
	this.futureExecutor = executor;
	}

	public static void enableHttpWireLog() {
	if (!logConfigured) {
	Logger httpWireLoggger = Logger.getLogger("org.apache.http.wire");
	httpWireLoggger.setLevel(Level.DEBUG);

	@SuppressWarnings("rawtypes")
	Enumeration rootLoggerAppenders = LogManager.getRootLogger().getAllAppenders();
	while (rootLoggerAppenders.hasMoreElements()) {
	Appender logAppender = (Appender) rootLoggerAppenders.nextElement();
	if (logAppender instanceof AppenderSkeleton) {
	((AppenderSkeleton) logAppender).setThreshold(Level.DEBUG);
	}
	httpWireLoggger.addAppender(logAppender);
	}

	logConfigured = true;
	}
	}

	public ListenableFuture<HttpResponse> execute(HttpRequestBase request, int... expectedStatuses) {
	return Futures.transformAsync(
	listenableFutureOf(httpclient.execute(request, null)),
	(res) -> {
	IOException exc =
	getResponseException(
	String.format("HTTP %s %s FAILED", request.getMethod(), request.getURI()),
	res,
	expectedStatuses);
	if (exc == null) {
	return Futures.immediateFuture(res);
	}
	return Futures.immediateFailedFuture(exc);
	},
	futureExecutor);
	}

	public <I, O> ListenableFuture<O> mapAsync(
	ListenableFuture<I> inputFuture, AsyncFunction<? super I, ? extends O> mapFunction) {
	return Futures.transformAsync(inputFuture, mapFunction, futureExecutor);
	}

	public <I, O> ListenableFuture<O> map(
	ListenableFuture<I> inputFuture, Function<? super I, ? extends O> mapFunction) {
	return Futures.transform(inputFuture, mapFunction, futureExecutor);
	}

	public HttpGet createGetTopic(String topic) {
	HttpGet get = new HttpGet(resolveKafkaRestApiUri("/topics/" + topic));
	get.addHeader(HttpHeaders.ACCEPT, KAFKA_V2);
	get.setConfig(createRequestConfig());
	return get;
	}

	public HttpGet createGetRecords(URI consumerUri) {
	HttpGet get = new HttpGet(consumerUri.resolve(consumerUri.getPath() + "/records"));
	get.addHeader(HttpHeaders.ACCEPT, KAFKA_V2_JSON);
	get.setConfig(createRequestConfig());
	return get;
	}

	public HttpPost createPostToConsumer(String consumerGroup) {
	HttpPost post =
	new HttpPost(
	resolveKafkaRestApiUri("/consumers/" + URLEncoder.encode(consumerGroup, UTF_8)));
	post.addHeader(HttpHeaders.ACCEPT, MediaType.ANY_TYPE.toString());
	post.setConfig(createRequestConfig());
	post.setEntity(
	new StringEntity(
	"{\"format\": \"json\",\"auto.offset.reset\": \"earliest\", \"auto.commit.enable\":\"true\", \"consumer.request.timeout.ms\": \"1000\"}",
	ContentType.create(KAFKA_V2, UTF_8)));
	return post;
	}

	public HttpDelete createDeleteToConsumer(URI consumerUri) {
	HttpDelete delete = new HttpDelete(consumerUri);
	delete.addHeader(HttpHeaders.ACCEPT, "/");
	delete.setConfig(createRequestConfig());
	return delete;
	}

	public HttpDelete createDeleteToConsumerSubscriptions(URI consumerUri) {
	URI subscriptionUri = consumerUri.resolve("subscription");
	HttpDelete delete = new HttpDelete(subscriptionUri);
	delete.addHeader(HttpHeaders.ACCEPT, "/");
	delete.setConfig(createRequestConfig());
	return delete;
	}

	public HttpPost createPostToSubscribe(URI consumerUri, String topic) {
	HttpPost post = new HttpPost(consumerUri.resolve(consumerUri.getPath() + "/subscription"));
	post.addHeader(HttpHeaders.ACCEPT, "/");
	post.setConfig(createRequestConfig());
	post.setEntity(
	new StringEntity(
	String.format("{\"topics\":[\"%s\"]}", topic), ContentType.create(KAFKA_V2, UTF_8)));
	return post;
	}

	public HttpPost createPostToTopic(String topic, HttpEntity postBodyEntity) {
	HttpPost post =
	new HttpPost(resolveKafkaRestApiUri("/topics/" + URLEncoder.encode(topic, UTF_8)));
	post.addHeader(HttpHeaders.ACCEPT, "/");
	post.setConfig(createRequestConfig());
	post.setEntity(postBodyEntity);
	return post;
	}

	public HttpPost createPostSeekTopicFromBeginning(
	URI consumerUri, String topic, Set<Integer> partitions) {
	HttpPost post =
	new HttpPost(consumerUri.resolve(consumerUri.getPath() + "/positions/beginning"));
	post.addHeader(HttpHeaders.ACCEPT, "/");
	post.setConfig(createRequestConfig());
	post.setEntity(
	new StringEntity(
	String.format(
	"{\"partitions\":[%s]}",
	partitions.stream()
	.map(
	partition ->
	String.format("{\"topic\":\"%s\",\"partition\":%d}", topic, partition))
	.collect(Collectors.joining(","))),
	ContentType.create(KAFKA_V2, UTF_8)));
	return post;
	}

	@Nullable
	public IOException getResponseException(
	String errorMessage, HttpResponse response, int... okHttpStatuses) {
	int responseHttpStatus = response.getStatusLine().getStatusCode();
	if (okHttpStatuses.length == 0) {
	okHttpStatuses =
	new int[] {HttpStatus.SC_OK, HttpStatus.SC_CREATED, HttpStatus.SC_NO_CONTENT};
	}
	for (int httpStatus : okHttpStatuses) {
	if (responseHttpStatus == httpStatus) {
	return null;
	}
	}

	String responseBody = "";
	try {
	responseBody = getStringEntity(response);
	} catch (IOException e) {
	logger.atWarning().withCause(e).log(
	"Unable to extrace the string entity for response %d (%s)",
	response.getStatusLine().getStatusCode(), response.getStatusLine().getReasonPhrase());
	}

	return new IOException(
	String.format(
	"%s\nHTTP status %d (%s)\n%s",
	errorMessage,
	response.getStatusLine().getStatusCode(),
	response.getStatusLine().getReasonPhrase(),
	responseBody));
	}

	protected String getStringEntity(HttpResponse response) throws IOException {
	HttpEntity entity = response.getEntity();
	if (entity != null) {
	try (ByteArrayOutputStream outStream = new ByteArrayOutputStream()) {
	entity.writeTo(outStream);
	outStream.close();
	return outStream.toString(UTF_8);
	}
	}
	return "";
	}

	private <V> ListenableFuture<V> listenableFutureOf(Future<V> future) {
	return JdkFutureAdapters.listenInPoolThread(future, futureExecutor);
	}

	private RequestConfig createRequestConfig() {
	Builder configBuilder =
	RequestConfig.custom()
	.setConnectionRequestTimeout(kafkaRestApiTimeoutMsec)
	.setConnectTimeout(kafkaRestApiTimeoutMsec)
	.setSocketTimeout(kafkaRestApiTimeoutMsec);
	configBuilder = proxy.apply(configBuilder);
	RequestConfig config = configBuilder.build();
	return config;
	}

	public void close() throws IOException {
	httpclient.close();
	}

	public URI resolveKafkaRestApiUri(String path) {
	try {
	URI restApiUri = configuration.getRestApiUri();
	return restApiUri.resolve(path);
	} catch (URISyntaxException e) {
	throw new IllegalArgumentException("Invalid Kafka REST API URI", e);
	}
	}

	public URI resolveKafkaRestApiUri(String kafkaRestId, String path) {
	URI restApiUri;
	try {
	restApiUri = configuration.getRestApiUri(kafkaRestId);
	return restApiUri.resolve(restApiUri.getPath() + path);
	} catch (URISyntaxException e) {
	throw new IllegalArgumentException("Invalid Kafka REST API URI", e);
	}
	}
	}