package yanagishima.servlet;

import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.OkHttpClient;
import okhttp3.Response;
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 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) {
			map = new HashMap();
			map.put("state", "FAILED");
			map.put("failureInfo", "");
		}
		writer.println(mapper.writeValueAsString(map));
	}

}

@@ -125,24 +125,81 @@
 3.Response;%0A
+import org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A
 import yanag
@@ -875,16 +875,100 @@
 vlet %7B%0A%0A
+%09private static Logger LOGGER = LoggerFactory.getLogger(QueryStatusServlet.class);%0A%0A
 %09private
@@ -3071,16 +3071,52 @@
 on e) %7B%0A
+%09%09%09LOGGER.error(e.getMessage(), e);%0A
 %09%09%09map =

Add logging