/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You 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 org.apache.zeppelin.livy; import com.google.gson.Gson; import com.google.gson.GsonBuilder; import com.google.gson.annotations.SerializedName; import org.apache.commons.lang.StringUtils; import org.apache.http.client.HttpClient; import org.apache.http.conn.ssl.SSLConnectionSocketFactory; import org.apache.http.conn.ssl.SSLContexts; import org.apache.http.impl.client.HttpClients; import org.apache.commons.lang.exception.ExceptionUtils; import org.apache.zeppelin.interpreter.*; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.http.HttpEntity; import org.springframework.http.HttpHeaders; import org.springframework.http.HttpMethod; import org.springframework.http.ResponseEntity; import org.springframework.http.client.HttpComponentsClientHttpRequestFactory; import org.springframework.security.kerberos.client.KerberosRestTemplate; import org.springframework.web.client.HttpClientErrorException; import org.springframework.web.client.RestClientException; import org.springframework.web.client.RestTemplate; import javax.net.ssl.SSLContext; import java.io.FileInputStream; import java.io.IOException; import java.security.KeyStore; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; /** * Base class for livy interpreters. */ public abstract class BaseLivyInterpreter extends Interpreter { protected static final Logger LOGGER = LoggerFactory.getLogger(BaseLivyInterpreter.class); private static Gson gson = new GsonBuilder().setPrettyPrinting().disableHtmlEscaping().create(); private static String SESSION_NOT_FOUND_PATTERN = "\"Session '\\d+' not found.\""; protected volatile SessionInfo sessionInfo; private String livyURL; private int sessionCreationTimeout; private int pullStatusInterval; protected boolean displayAppInfo; protected LivyVersion livyVersion; private RestTemplate restTemplate; Set<Object> paragraphsToCancel = Collections.newSetFromMap( new ConcurrentHashMap<Object, Boolean>()); private ConcurrentHashMap<String, Integer> paragraphId2StmtProgressMap = new ConcurrentHashMap<>(); public BaseLivyInterpreter(Properties property) { super(property); this.livyURL = property.getProperty("zeppelin.livy.url"); this.displayAppInfo = Boolean.parseBoolean( property.getProperty("zeppelin.livy.displayAppInfo", "false")); this.sessionCreationTimeout = Integer.parseInt( property.getProperty("zeppelin.livy.session.create_timeout", 120 + "")); this.pullStatusInterval = Integer.parseInt( property.getProperty("zeppelin.livy.pull_status.interval.millis", 1000 + "")); this.restTemplate = createRestTemplate(); } public abstract String getSessionKind(); @Override public void open() { try { initLivySession(); } catch (LivyException e) { String msg = "Fail to create session, please check livy interpreter log and " + "livy server log"; throw new RuntimeException(msg, e); } } @Override public void close() { if (sessionInfo != null) { closeSession(sessionInfo.id); // reset sessionInfo to null so that we won't close it twice. sessionInfo = null; } } protected void initLivySession() throws LivyException { this.sessionInfo = createSession(getUserName(), getSessionKind()); if (displayAppInfo) { if (sessionInfo.appId == null) { // livy 0.2 don't return appId and sparkUiUrl in response so that we need to get it // explicitly by ourselves. sessionInfo.appId = extractAppId(); } if (sessionInfo.appInfo == null || StringUtils.isEmpty(sessionInfo.appInfo.get("sparkUiUrl"))) { sessionInfo.webUIAddress = extractWebUIAddress(); } else { sessionInfo.webUIAddress = sessionInfo.appInfo.get("sparkUiUrl"); } LOGGER.info("Create livy session successfully with sessionId: {}, appId: {}, webUI: {}", sessionInfo.id, sessionInfo.appId, sessionInfo.webUIAddress); } else { LOGGER.info("Create livy session successfully with sessionId: {}", this.sessionInfo.id); } // check livy version try { this.livyVersion = getLivyVersion(); LOGGER.info("Use livy " + livyVersion); } catch (APINotFoundException e) { this.livyVersion = new LivyVersion("0.2.0"); LOGGER.info("Use livy 0.2.0"); } } protected abstract String extractAppId() throws LivyException; protected abstract String extractWebUIAddress() throws LivyException; public SessionInfo getSessionInfo() { return sessionInfo; } @Override public InterpreterResult interpret(String st, InterpreterContext context) { if (StringUtils.isEmpty(st)) { return new InterpreterResult(InterpreterResult.Code.SUCCESS, ""); } try { return interpret(st, context.getParagraphId(), this.displayAppInfo, true); } catch (LivyException e) { LOGGER.error("Fail to interpret:" + st, e); return new InterpreterResult(InterpreterResult.Code.ERROR, InterpreterUtils.getMostRelevantMessage(e)); } } @Override public void cancel(InterpreterContext context) { paragraphsToCancel.add(context.getParagraphId()); LOGGER.info("Added paragraph " + context.getParagraphId() + " for cancellation."); } @Override public FormType getFormType() { return FormType.SIMPLE; } @Override public int getProgress(InterpreterContext context) { if (livyVersion.isGetProgressSupported()) { String paraId = context.getParagraphId(); Integer progress = paragraphId2StmtProgressMap.get(paraId); return progress == null ? 0 : progress; } return 0; } private SessionInfo createSession(String user, String kind) throws LivyException { try { Map<String, String> conf = new HashMap<>(); for (Map.Entry<Object, Object> entry : property.entrySet()) { if (entry.getKey().toString().startsWith("livy.spark.") && !entry.getValue().toString().isEmpty()) conf.put(entry.getKey().toString().substring(5), entry.getValue().toString()); } CreateSessionRequest request = new CreateSessionRequest(kind, user == null || user.equals("anonymous") ? null : user, conf); SessionInfo sessionInfo = SessionInfo.fromJson( callRestAPI("/sessions", "POST", request.toJson())); long start = System.currentTimeMillis(); // pull the session status until it is idle or timeout while (!sessionInfo.isReady()) { if ((System.currentTimeMillis() - start) / 1000 > sessionCreationTimeout) { String msg = "The creation of session " + sessionInfo.id + " is timeout within " + sessionCreationTimeout + " seconds, appId: " + sessionInfo.appId + ", log: " + sessionInfo.log; throw new LivyException(msg); } Thread.sleep(pullStatusInterval); sessionInfo = getSessionInfo(sessionInfo.id); LOGGER.info("Session {} is in state {}, appId {}", sessionInfo.id, sessionInfo.state, sessionInfo.appId); if (sessionInfo.isFinished()) { String msg = "Session " + sessionInfo.id + " is finished, appId: " + sessionInfo.appId + ", log: " + sessionInfo.log; throw new LivyException(msg); } } return sessionInfo; } catch (Exception e) { LOGGER.error("Error when creating livy session for user " + user, e); throw new LivyException(e); } } private SessionInfo getSessionInfo(int sessionId) throws LivyException { return SessionInfo.fromJson(callRestAPI("/sessions/" + sessionId, "GET")); } public InterpreterResult interpret(String code, String paragraphId, boolean displayAppInfo, boolean appendSessionExpired) throws LivyException { StatementInfo stmtInfo = null; boolean sessionExpired = false; try { try { stmtInfo = executeStatement(new ExecuteRequest(code)); } catch (SessionNotFoundException e) { LOGGER.warn("Livy session {} is expired, new session will be created.", sessionInfo.id); sessionExpired = true; // we don't want to create multiple sessions because it is possible to have multiple thread // to call this method, like LivySparkSQLInterpreter which use ParallelScheduler. So we need // to check session status again in this sync block synchronized (this) { if (isSessionExpired()) { initLivySession(); } } stmtInfo = executeStatement(new ExecuteRequest(code)); } // pull the statement status while (!stmtInfo.isAvailable()) { if (paragraphId != null && paragraphsToCancel.contains(paragraphId)) { cancel(stmtInfo.id, paragraphId); return new InterpreterResult(InterpreterResult.Code.ERROR, "Job is cancelled"); } try { Thread.sleep(pullStatusInterval); } catch (InterruptedException e) { LOGGER.error("InterruptedException when pulling statement status.", e); throw new LivyException(e); } stmtInfo = getStatementInfo(stmtInfo.id); if (paragraphId != null) { paragraphId2StmtProgressMap.put(paragraphId, (int) (stmtInfo.progress * 100)); } } if (appendSessionExpired) { return appendSessionExpire(getResultFromStatementInfo(stmtInfo, displayAppInfo), sessionExpired); } else { return getResultFromStatementInfo(stmtInfo, displayAppInfo); } } finally { if (paragraphId != null) { paragraphId2StmtProgressMap.remove(paragraphId); paragraphsToCancel.remove(paragraphId); } } } private void cancel(int id, String paragraphId) { if (livyVersion.isCancelSupported()) { try { LOGGER.info("Cancelling statement " + id); cancelStatement(id); } catch (LivyException e) { LOGGER.error("Fail to cancel statement " + id + " for paragraph " + paragraphId, e); } finally { paragraphsToCancel.remove(paragraphId); } } else { LOGGER.warn("cancel is not supported for this version of livy: " + livyVersion); paragraphsToCancel.clear(); } } protected LivyVersion getLivyVersion() throws LivyException { return new LivyVersion((LivyVersionResponse.fromJson(callRestAPI("/version", "GET")).version)); } private boolean isSessionExpired() throws LivyException { try { getSessionInfo(sessionInfo.id); return false; } catch (SessionNotFoundException e) { return true; } catch (LivyException e) { throw e; } } private InterpreterResult appendSessionExpire(InterpreterResult result, boolean sessionExpired) { if (sessionExpired) { InterpreterResult result2 = new InterpreterResult(result.code()); result2.add(InterpreterResult.Type.HTML, "<font color=\"red\">Previous livy session is expired, new livy session is created. " + "Paragraphs that depend on this paragraph need to be re-executed!" + "</font>"); for (InterpreterResultMessage message : result.message()) { result2.add(message.getType(), message.getData()); } return result2; } else { return result; } } private InterpreterResult getResultFromStatementInfo(StatementInfo stmtInfo, boolean displayAppInfo) { if (stmtInfo.output != null && stmtInfo.output.isError()) { return new InterpreterResult(InterpreterResult.Code.ERROR, stmtInfo.output.evalue); } else if (stmtInfo.isCancelled()) { // corner case, output might be null if it is cancelled. return new InterpreterResult(InterpreterResult.Code.ERROR, "Job is cancelled"); } else if (stmtInfo.output == null) { // This case should never happen, just in case return new InterpreterResult(InterpreterResult.Code.ERROR, "Empty output"); } else { //TODO(zjffdu) support other types of data (like json, image and etc) String result = stmtInfo.output.data.plain_text; // check table magic result first if (stmtInfo.output.data.application_livy_table_json != null) { StringBuilder outputBuilder = new StringBuilder(); boolean notFirstColumn = false; for (Map header : stmtInfo.output.data.application_livy_table_json.headers) { if (notFirstColumn) { outputBuilder.append("\t"); } outputBuilder.append(header.get("name")); notFirstColumn = true; } outputBuilder.append("\n"); for (List<Object> row : stmtInfo.output.data.application_livy_table_json.records) { outputBuilder.append(StringUtils.join(row, "\t")); outputBuilder.append("\n"); } return new InterpreterResult(InterpreterResult.Code.SUCCESS, InterpreterResult.Type.TABLE, outputBuilder.toString()); } else if (stmtInfo.output.data.image_png != null) { return new InterpreterResult(InterpreterResult.Code.SUCCESS, InterpreterResult.Type.IMG, (String) stmtInfo.output.data.image_png); } else if (result != null) { result = result.trim(); if (result.startsWith("<link") || result.startsWith("<script") || result.startsWith("<style") || result.startsWith("<div")) { result = "%html " + result; } } if (displayAppInfo) { InterpreterResult interpreterResult = new InterpreterResult(InterpreterResult.Code.SUCCESS); interpreterResult.add(result); String appInfoHtml = "<hr/>Spark Application Id: " + sessionInfo.appId + "<br/>" + "Spark WebUI: <a href=\"" + sessionInfo.webUIAddress + "\">" + sessionInfo.webUIAddress + "</a>"; interpreterResult.add(InterpreterResult.Type.HTML, appInfoHtml); return interpreterResult; } else { return new InterpreterResult(InterpreterResult.Code.SUCCESS, result); } } } private StatementInfo executeStatement(ExecuteRequest executeRequest) throws LivyException { return StatementInfo.fromJson(callRestAPI("/sessions/" + sessionInfo.id + "/statements", "POST", executeRequest.toJson())); } private StatementInfo getStatementInfo(int statementId) throws LivyException { return StatementInfo.fromJson( callRestAPI("/sessions/" + sessionInfo.id + "/statements/" + statementId, "GET")); } private void cancelStatement(int statementId) throws LivyException { callRestAPI("/sessions/" + sessionInfo.id + "/statements/" + statementId + "/cancel", "POST"); } private RestTemplate createRestTemplate() { HttpClient httpClient = null; if (livyURL.startsWith("https:")) { String keystoreFile = property.getProperty("zeppelin.livy.ssl.trustStore"); String password = property.getProperty("zeppelin.livy.ssl.trustStorePassword"); if (StringUtils.isBlank(keystoreFile)) { throw new RuntimeException("No zeppelin.livy.ssl.trustStore specified for livy ssl"); } if (StringUtils.isBlank(password)) { throw new RuntimeException("No zeppelin.livy.ssl.trustStorePassword specified " + "for livy ssl"); } FileInputStream inputStream = null; try { inputStream = new FileInputStream(keystoreFile); KeyStore trustStore = KeyStore.getInstance(KeyStore.getDefaultType()); trustStore.load(new FileInputStream(keystoreFile), password.toCharArray()); SSLContext sslContext = SSLContexts.custom() .loadTrustMaterial(trustStore) .build(); SSLConnectionSocketFactory csf = new SSLConnectionSocketFactory(sslContext); httpClient = HttpClients.custom().setSSLSocketFactory(csf).build(); } catch (Exception e) { throw new RuntimeException("Failed to create SSL HttpClient", e); } finally { if (inputStream != null) { try { inputStream.close(); } catch (IOException e) { LOGGER.error("Failed to close keystore file", e); } } } } String keytabLocation = property.getProperty("zeppelin.livy.keytab"); String principal = property.getProperty("zeppelin.livy.principal"); if (StringUtils.isNotEmpty(keytabLocation) && StringUtils.isNotEmpty(principal)) { if (httpClient == null) { return new KerberosRestTemplate(keytabLocation, principal); } else { return new KerberosRestTemplate(keytabLocation, principal, httpClient); } } if (httpClient == null) { return new RestTemplate(); } else { return new RestTemplate(new HttpComponentsClientHttpRequestFactory(httpClient)); } } private String callRestAPI(String targetURL, String method) throws LivyException { return callRestAPI(targetURL, method, ""); } private String callRestAPI(String targetURL, String method, String jsonData) throws LivyException { targetURL = livyURL + targetURL; LOGGER.debug("Call rest api in {}, method: {}, jsonData: {}", targetURL, method, jsonData); HttpHeaders headers = new HttpHeaders(); headers.add("Content-Type", "application/json"); headers.add("X-Requested-By", "zeppelin"); ResponseEntity<String> response = null; try { if (method.equals("POST")) { HttpEntity<String> entity = new HttpEntity<>(jsonData, headers); response = restTemplate.exchange(targetURL, HttpMethod.POST, entity, String.class); } else if (method.equals("GET")) { HttpEntity<String> entity = new HttpEntity<>(headers); response = restTemplate.exchange(targetURL, HttpMethod.GET, entity, String.class); } else if (method.equals("DELETE")) { HttpEntity<String> entity = new HttpEntity<>(headers); response = restTemplate.exchange(targetURL, HttpMethod.DELETE, entity, String.class); } } catch (HttpClientErrorException e) { response = new ResponseEntity(e.getResponseBodyAsString(), e.getStatusCode()); LOGGER.error(String.format("Error with %s StatusCode: %s", response.getStatusCode().value(), e.getResponseBodyAsString())); } catch (RestClientException e) { // Exception happens when kerberos is enabled. if (e.getCause() instanceof HttpClientErrorException) { HttpClientErrorException cause = (HttpClientErrorException) e.getCause(); if (cause.getResponseBodyAsString().matches(SESSION_NOT_FOUND_PATTERN)) { throw new SessionNotFoundException(cause.getResponseBodyAsString()); } throw new LivyException(cause.getResponseBodyAsString() + "\n" + ExceptionUtils.getFullStackTrace(ExceptionUtils.getRootCause(e))); } throw new LivyException(e); } if (response == null) { throw new LivyException("No http response returned"); } LOGGER.debug("Get response, StatusCode: {}, responseBody: {}", response.getStatusCode(), response.getBody()); if (response.getStatusCode().value() == 200 || response.getStatusCode().value() == 201) { return response.getBody(); } else if (response.getStatusCode().value() == 404) { if (response.getBody().matches(SESSION_NOT_FOUND_PATTERN)) { throw new SessionNotFoundException(response.getBody()); } else { throw new APINotFoundException("No rest api found for " + targetURL + ", " + response.getStatusCode()); } } else { String responseString = response.getBody(); if (responseString.contains("CreateInteractiveRequest[\\\"master\\\"]")) { return responseString; } throw new LivyException(String.format("Error with %s StatusCode: %s", response.getStatusCode().value(), responseString)); } } private void closeSession(int sessionId) { try { callRestAPI("/sessions/" + sessionId, "DELETE"); } catch (Exception e) { LOGGER.error(String.format("Error closing session for user with session ID: %s", sessionId), e); } } /* * We create these POJO here to accommodate livy 0.3 which is not released yet. livy rest api has * some changes from version to version. So we create these POJO in zeppelin side to accommodate * incompatibility between versions. Later, when livy become more stable, we could just depend on * livy client jar. */ private static class CreateSessionRequest { public final String kind; @SerializedName("proxyUser") public final String user; public final Map<String, String> conf; public CreateSessionRequest(String kind, String user, Map<String, String> conf) { this.kind = kind; this.user = user; this.conf = conf; } public String toJson() { return gson.toJson(this); } } /** * */ public static class SessionInfo { public final int id; public String appId; public String webUIAddress; public final String owner; public final String proxyUser; public final String state; public final String kind; public final Map<String, String> appInfo; public final List<String> log; public SessionInfo(int id, String appId, String owner, String proxyUser, String state, String kind, Map<String, String> appInfo, List<String> log) { this.id = id; this.appId = appId; this.owner = owner; this.proxyUser = proxyUser; this.state = state; this.kind = kind; this.appInfo = appInfo; this.log = log; } public boolean isReady() { return state.equals("idle"); } public boolean isFinished() { return state.equals("error") || state.equals("dead") || state.equals("success"); } public static SessionInfo fromJson(String json) { return gson.fromJson(json, SessionInfo.class); } } private static class ExecuteRequest { public final String code; public ExecuteRequest(String code) { this.code = code; } public String toJson() { return gson.toJson(this); } } private static class StatementInfo { public Integer id; public String state; public double progress; public StatementOutput output; public StatementInfo() { } public static StatementInfo fromJson(String json) { return gson.fromJson(json, StatementInfo.class); } public boolean isAvailable() { return state.equals("available") || state.equals("cancelled"); } public boolean isCancelled() { return state.equals("cancelled"); } private static class StatementOutput { public String status; public String execution_count; public Data data; public String ename; public String evalue; public Object traceback; public TableMagic tableMagic; public boolean isError() { return status.equals("error"); } public String toJson() { return gson.toJson(this); } private static class Data { @SerializedName("text/plain") public String plain_text; @SerializedName("image/png") public String image_png; @SerializedName("application/json") public String application_json; @SerializedName("application/vnd.livy.table.v1+json") public TableMagic application_livy_table_json; } private static class TableMagic { @SerializedName("headers") List<Map> headers; @SerializedName("data") List<List> records; } } } private static class LivyVersionResponse { public String url; public String branch; public String revision; public String version; public String date; public String user; public static LivyVersionResponse fromJson(String json) { return gson.fromJson(json, LivyVersionResponse.class); } } }