package yanagishima.servlet;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.http.client.fluent.Request;
import yanagishima.config.YanagishimaConfig;
import yanagishima.util.AccessControlUtil;
import yanagishima.util.HttpRequestUtil;

import javax.inject.Inject;
import javax.inject.Singleton;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;

@Singleton
public class QueryStatusServlet extends HttpServlet {

	private static final long serialVersionUID = 1L;

	private YanagishimaConfig yanagishimaConfig;

	private static final int LIMIT = 100;

	@Inject
	public QueryStatusServlet(YanagishimaConfig yanagishimaConfig) {
		this.yanagishimaConfig = yanagishimaConfig;
	}

	@Override
	protected void doGet(HttpServletRequest request,
			HttpServletResponse response) throws ServletException, IOException {

		String datasource = HttpRequestUtil.getParam(request, "datasource");
		if(yanagishimaConfig.isCheckDatasource()) {
			if(!AccessControlUtil.validateDatasource(request, datasource)) {
				try {
					response.sendError(SC_FORBIDDEN);
					return;
				} catch (IOException e) {
					throw new RuntimeException(e);
				}
			}
		}
		String queryid = Optional.ofNullable(request.getParameter("queryid")).get();
		String prestoCoordinatorServer = yanagishimaConfig
				.getPrestoCoordinatorServer(datasource);
		response.setContentType("application/json");
		PrintWriter writer = response.getWriter();
		String json = Request.Get(prestoCoordinatorServer + "/v1/query/" + queryid)
				.execute().returnContent().asString(StandardCharsets.UTF_8);
		ObjectMapper mapper = new ObjectMapper();
		Map map = mapper.readValue(json, Map.class);
		if(map.containsKey("outputStage")) {
			map.remove("outputStage");
		}
		if(map.containsKey("session")) {
			map.remove("session");
		}
		writer.println(mapper.writeValueAsString(map));
	}

}

@@ -87,44 +87,52 @@
 rt o
-rg.apache.http.c
+khttp3.OkHttpC
 lient
-.fluent.Request
+;%0Aimport okhttp3.Response
 ;%0Aim
@@ -561,110 +561,108 @@
 ava.
-nio.charset.StandardCharsets;%0Aimport java.util.List;%0Aimport java.util.Map;%0Aimport java.util.Optional;%0A
+util.Map;%0Aimport java.util.Optional;%0A%0Aimport static com.facebook.presto.client.OkHttpUtil.basicAuth;
 %0Aimp
@@ -903,36 +903,52 @@
 ate 
-static final int LIMIT = 100
+OkHttpClient httpClient = new OkHttpClient()
 ;%0A%0A%09
@@ -1132,16 +1132,187 @@
 equest,%0A
+%09%09%09%09%09%09 HttpServletResponse response) throws ServletException, IOException %7B%0A%09%09doPost(request, response);%0A%09%7D%0A%0A%09@Override%0A%09protected void doPost(HttpServletRequest request,%0A
 %09%09%09HttpS
@@ -1704,16 +1704,17 @@
 %09%09%7D%0A%09%09%7D%0A
+%0A
 %09%09String
@@ -1836,21 +1836,16 @@
 maConfig
-%0A%09%09%09%09
 .getPres
@@ -1989,134 +1989,824 @@
 n = 
-Request.Get(prestoCoordinatorServer + %22/v1/query/%22 + queryid)%0A%09%09%09%09.execute().returnContent().asString(StandardCharsets.UTF_8);
+null;%0A%09%09okhttp3.Request prestoRequest = new okhttp3.Request.Builder().url(prestoCoordinatorServer + %22/v1/query/%22 + queryid).build();%0A%09%09Optional%3CString%3E prestoUser = Optional.ofNullable(request.getParameter(%22presto_user%22));%0A%09%09Optional%3CString%3E prestoPassword = Optional.ofNullable(request.getParameter(%22presto_password%22));%0A%09%09if (prestoUser.isPresent() && prestoPassword.isPresent()) %7B%0A%09%09%09OkHttpClient.Builder clientBuilder = httpClient.newBuilder();%0A%09%09%09clientBuilder.addInterceptor(basicAuth(prestoUser.get(), prestoPassword.get()));%0A%09%09%09try (Response prestoResponse = clientBuilder.build().newCall(prestoRequest).execute()) %7B%0A%09%09%09%09json = prestoResponse.body().string();%0A%09%09%09%7D%0A%09%09%7D else %7B%0A%09%09%09try (Response prestoResponse = httpClient.newCall(prestoRequest).execute()) %7B%0A%09%09%09%09json = prestoResponse.body().string();%0A%09%09%09%7D%0A%09%09%7D%0A
 %0A%09%09O

Add auth logic to /queryStatus API