package yanagishima.servlet;

import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.OkHttpClient;
import okhttp3.Response;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
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.util.HashMap;
import java.util.Map;
import java.util.Optional;

import static com.facebook.presto.client.OkHttpUtil.basicAuth;
import static javax.servlet.http.HttpServletResponse.SC_FORBIDDEN;

@Singleton
public class QueryStatusServlet extends HttpServlet {

	private static Logger LOGGER = LoggerFactory.getLogger(QueryStatusServlet.class);

	private static final long serialVersionUID = 1L;

	private YanagishimaConfig yanagishimaConfig;

	private OkHttpClient httpClient = new OkHttpClient();

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

	@Override
	protected void doPost(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 = null;
		okhttp3.Request prestoRequest = new okhttp3.Request.Builder().url(prestoCoordinatorServer + "/v1/query/" + queryid).build();
		Optional<String> prestoUser = Optional.ofNullable(request.getParameter("user"));
		Optional<String> prestoPassword = Optional.ofNullable(request.getParameter("password"));
		if (prestoUser.isPresent() && prestoPassword.isPresent()) {
			OkHttpClient.Builder clientBuilder = httpClient.newBuilder();
			clientBuilder.addInterceptor(basicAuth(prestoUser.get(), prestoPassword.get()));
			try (Response prestoResponse = clientBuilder.build().newCall(prestoRequest).execute()) {
				json = prestoResponse.body().string();
			}
		} else {
			try (Response prestoResponse = httpClient.newCall(prestoRequest).execute()) {
				json = prestoResponse.body().string();
			}
		}

		Map map = null;
		ObjectMapper mapper = new ObjectMapper();
		try {
			map = mapper.readValue(json, Map.class);
			if(map.containsKey("outputStage")) {
				map.remove("outputStage");
			}
			if(map.containsKey("session")) {
				map.remove("session");
			}
		} catch (IOException e) {
			LOGGER.error(e.getMessage(), e);
			map = new HashMap();
			map.put("state", "FAILED");
			map.put("failureInfo", "");
		}
		writer.println(mapper.writeValueAsString(map));
	}

}

@@ -23,16 +23,75 @@
 rvlet;%0A%0A
+import com.fasterxml.jackson.core.JsonProcessingException;%0A
 import c
@@ -2643,32 +2643,73 @@
 t).execute()) %7B%0A
+%09%09%09%09if(prestoResponse.isSuccessful()) %7B%0A%09
 %09%09%09%09json = prest
@@ -2742,144 +2742,409 @@
 %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();
+%09%7D else %7B%0A%09%09%09%09%09writer.println(getError(prestoResponse.code(), prestoResponse.message()));%0A%09%09%09%09%09return;%0A%09%09%09%09%7D%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%09if(prestoResponse.isSuccessful()) %7B%0A%09%09%09%09%09json = prestoResponse.body().string();%0A%09%09%09%09%7D else %7B%0A%09%09%09%09%09writer.println(getError(prestoResponse.code(), prestoResponse.message()));%0A%09%09%09%09%09return;%0A%09%09%09%09%7D
 %0A%09%09%09
@@ -3616,10 +3616,436 @@
 ));%0A%09%7D%0A%0A
+%09private String getError(int code, String message) %7B%0A%09%09try %7B%0A%09%09%09Map%3CString, Object%3E map = new HashMap%3CString, Object%3E();%0A%09%09%09map.put(%22state%22, %22FAILED%22);%0A%09%09%09map.put(%22failureInfo%22, %22%22);%0A%09%09%09map.put(%22error%22, String.format(%22code=%25d, message=%25s%22, code, message));%0A%09%09%09ObjectMapper mapper = new ObjectMapper();%0A%09%09%09return mapper.writeValueAsString(map);%0A%09%09%7D catch (JsonProcessingException e) %7B%0A%09%09%09throw new RuntimeException(e);%0A%09%09%7D%0A%09%7D%0A%0A
 %7D%0A

Improve error handling