diff --git a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/NotebookRepoSync.java b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/NotebookRepoSync.java index 8553349ed6..6bbd5bca99 100644 --- a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/NotebookRepoSync.java +++ b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/NotebookRepoSync.java @@ -75,7 +75,6 @@ public class NotebookRepoSync implements NotebookRepo { } for (int i = 0; i < Math.min(storageClassNames.length, getMaxRepoNum()); i++) { - @SuppressWarnings("static-access") Class notebookStorageClass; try { notebookStorageClass = getClass().forName(storageClassNames[i].trim()); diff --git a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/ZeppelinHubRepo.java b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/ZeppelinHubRepo.java index fb57c0c38c..2f33f6f9fc 100644 --- a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/ZeppelinHubRepo.java +++ b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/ZeppelinHubRepo.java @@ -34,6 +34,7 @@ import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserTokenContainer; import org.apache.zeppelin.notebook.repo.zeppelinhub.model.UserSessionContainer; import org.apache.zeppelin.notebook.repo.zeppelinhub.rest.ZeppelinhubRestApiHandler; import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.Client; +import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.utils.ZeppelinhubUtils; import org.apache.zeppelin.user.AuthenticationInfo; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -76,7 +77,6 @@ public class ZeppelinHubRepo implements NotebookRepo { websocketClient = Client.initialize(getZeppelinWebsocketUri(conf), getZeppelinhubWebsocketUri(conf), token, conf); websocketClient.start(); - } private String getZeppelinHubWsUri(URI api) throws URISyntaxException { @@ -241,7 +241,6 @@ public class ZeppelinHubRepo implements NotebookRepo { return EMPTY_NOTE; } String endpoint = Joiner.on("/").join(noteId, "checkpoint", revId); - String token = getUserToken(subject.getUser()); String response = restApiClient.get(token, endpoint); @@ -321,19 +320,24 @@ public class ZeppelinHubRepo implements NotebookRepo { LOG.info("User {} will switch instance", user); String ticket = UserSessionContainer.instance.getSession(user); List instances; + String currentToken = StringUtils.EMPTY, targetToken = StringUtils.EMPTY; try { instances = tokenManager.getUserInstances(ticket); if (instances.isEmpty()) { return; } - + currentToken = tokenManager.getExistingUserToken(user); for (Instance instance : instances) { if (instance.id == instanceId) { - LOG.info("User {} switched to instance {}", user, instances.get(0).name); + LOG.info("User {} switched to instance {}", user, instance.name); tokenManager.setUserToken(user, instance.token); + targetToken = instance.token; break; } } + if (!StringUtils.isBlank(currentToken) && !StringUtils.isBlank(targetToken)) { + ZeppelinhubUtils.userSwitchTokenRoutine(user, currentToken, targetToken); + } } catch (IOException e) { LOG.error("Cannot switch instance for user {}", user, e); } diff --git a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/model/UserTokenContainer.java b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/model/UserTokenContainer.java index c1d7fad739..3c24decdf3 100644 --- a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/model/UserTokenContainer.java +++ b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/model/UserTokenContainer.java @@ -86,6 +86,17 @@ public class UserTokenContainer { return token; } + public String getExistingUserToken(String principal) { + if (StringUtils.isBlank(principal) || "anonymous".equals(principal)) { + return StringUtils.EMPTY; + } + String token = userTokens.get(principal); + if (token == null) { + return StringUtils.EMPTY; + } + return token; + } + public String removeUserToken(String username) { return userTokens.remove(username); } diff --git a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/utils/ZeppelinhubUtils.java b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/utils/ZeppelinhubUtils.java index e55c6afc37..f0a9db450a 100644 --- a/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/utils/ZeppelinhubUtils.java +++ b/zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/zeppelinhub/websocket/utils/ZeppelinhubUtils.java @@ -118,4 +118,14 @@ public class ZeppelinhubUtils { .send(msg, token); ZeppelinhubClient.getInstance().removeSession(token); } + + public static void userSwitchTokenRoutine(String username, String originToken, + String targetToken) { + String offMsg = ZeppelinhubUtils.deadMessage(originToken); + ZeppelinhubClient.getInstance().send(offMsg, originToken); + ZeppelinhubClient.getInstance().removeSession(originToken); + + String onMsg = ZeppelinhubUtils.liveMessage(targetToken); + ZeppelinhubClient.getInstance().send(onMsg, targetToken); + } }