fix ws back connection

This commit is contained in:
Khalid Huseynov 2017-03-20 17:41:18 +09:00
parent 7e1ac5d74b
commit 5755a5c462
4 changed files with 23 additions and 17 deletions

View file

@ -194,7 +194,7 @@ public class ZeppelinClient {
}
public void send(Message msg, String noteId) {
Session noteSession = getZeppelinConnection(noteId);
Session noteSession = getZeppelinConnection(noteId, msg.principal, msg.ticket);
if (!isSessionOpen(noteSession)) {
LOG.error("Cannot open websocket connection to Zeppelin note {}", noteId);
return;
@ -202,12 +202,12 @@ public class ZeppelinClient {
noteSession.getRemote().sendStringByFuture(serialize(msg));
}
public Session getZeppelinConnection(String noteId) {
public Session getZeppelinConnection(String noteId, String principal, String ticket) {
if (StringUtils.isBlank(noteId)) {
LOG.warn("Cannot get Websocket session with blanck noteId");
return null;
}
return getNoteSession(noteId);
return getNoteSession(noteId, principal, ticket);
}
/*
@ -220,18 +220,18 @@ public class ZeppelinClient {
}
*/
private Session getNoteSession(String noteId) {
private Session getNoteSession(String noteId, String principal, String ticket) {
LOG.info("Getting Note websocket connection for note {}", noteId);
Session session = notesConnection.get(noteId);
if (!isSessionOpen(session)) {
LOG.info("No open connection for note {}, opening one", noteId);
notesConnection.remove(noteId);
session = openNoteSession(noteId);
session = openNoteSession(noteId, principal, ticket);
}
return session;
}
private Session openNoteSession(String noteId) {
private Session openNoteSession(String noteId, String principal, String ticket) {
ClientUpgradeRequest request = new ClientUpgradeRequest();
ZeppelinWebsocket socket = new ZeppelinWebsocket(noteId);
Future<Session> future = null;
@ -248,7 +248,7 @@ public class ZeppelinClient {
session.close();
session = notesConnection.get(noteId);
} else {
String getNote = serialize(zeppelinGetNoteMsg(noteId));
String getNote = serialize(zeppelinGetNoteMsg(noteId, principal, ticket));
session.getRemote().sendStringByFuture(getNote);
notesConnection.put(noteId, session);
}
@ -259,11 +259,13 @@ public class ZeppelinClient {
return (session != null) && (session.isOpen());
}
private Message zeppelinGetNoteMsg(String noteId) {
private Message zeppelinGetNoteMsg(String noteId, String principal, String ticket) {
Message getNoteMsg = new Message(Message.OP.GET_NOTE);
HashMap<String, Object> data = new HashMap<String, Object>();
data.put("id", noteId);
getNoteMsg.data = data;
getNoteMsg.principal = principal;
getNoteMsg.ticket = ticket;
return getNoteMsg;
}
@ -275,24 +277,22 @@ public class ZeppelinClient {
if (zeppelinMsg == null) {
return;
}
LOG.info("message from {}", zeppelinMsg.principal);
String token;
if (!isActionable(zeppelinMsg.op)) {
return;
}
if (StringUtils.isEmpty(zeppelinMsg.principal) || zeppelinMsg.principal == "anonymous") {
token = "anonymous";
token = StringUtils.EMPTY;
} else {
token = UserTokenContainer.instance.getUserToken(zeppelinMsg.principal);
}
meta.put("token", token);
Client client = Client.getInstance();
if (client == null) {
LOG.warn("Client isn't initialized yet");
return;
}
ZeppelinhubMessage hubMsg = ZeppelinhubMessage.newMessage(zeppelinMsg, meta);
if (token == "anonymous") {
if (StringUtils.isEmpty(token)) {
relayToAllZeppelinHub(hubMsg, noteId);
} else {
client.relayToZeppelinHub(hubMsg.serialize(), zeppelinMsg.ticket);

View file

@ -37,6 +37,7 @@ import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.session.Zeppelinh
import org.apache.zeppelin.notebook.repo.zeppelinhub.websocket.utils.ZeppelinhubUtils;
import org.apache.zeppelin.notebook.socket.Message;
import org.apache.zeppelin.notebook.socket.Message.OP;
import org.apache.zeppelin.ticket.TicketContainer;
import org.eclipse.jetty.util.ssl.SslContextFactory;
import org.eclipse.jetty.websocket.api.Session;
import org.eclipse.jetty.websocket.client.ClientUpgradeRequest;
@ -236,6 +237,8 @@ public class ZeppelinhubClient {
return;
}
zeppelinMsg.data = (Map<String, Object>) hubMsg.data;
zeppelinMsg.principal = hubMsg.meta.get("owner");
zeppelinMsg.ticket = TicketContainer.instance.getTicket(zeppelinMsg.principal);
Client client = Client.getInstance();
if (client == null) {
LOG.warn("Base client isn't initialized, returning");
@ -260,6 +263,7 @@ public class ZeppelinhubClient {
Message zeppelinMsg = new Message(OP.RUN_PARAGRAPH);
JSONArray paragraphs = data.getJSONArray("data");
String principal = data.getJSONObject("meta").getString("owner");
for (int i = 0; i < paragraphs.length(); i++) {
if (!(paragraphs.get(i) instanceof JSONObject)) {
LOG.warn("Wrong \"paragraph\" format for RUN_NOTEBOOK");
@ -267,6 +271,8 @@ public class ZeppelinhubClient {
}
zeppelinMsg.data = gson.fromJson(paragraphs.getString(i),
new TypeToken<Map<String, Object>>(){}.getType());
zeppelinMsg.principal = principal;
zeppelinMsg.ticket = TicketContainer.instance.getTicket(principal);
client.relayToZeppelin(zeppelinMsg, noteId);
LOG.info("\nSending RUN_PARAGRAPH message to Zeppelin ");
}

View file

@ -62,25 +62,25 @@ public class ZeppelinClientTest {
LOG.info("Zeppelin websocket client started");
// Connection to note AAAA
Session connectionA = client.getZeppelinConnection("AAAA");
Session connectionA = client.getZeppelinConnection("AAAA", "anonymous", "anonymous");
assertNotNull(connectionA);
assertTrue(connectionA.isOpen());
assertEquals(client.countConnectedNotes(), 1);
assertEquals(connectionA, client.getZeppelinConnection("AAAA"));
assertEquals(connectionA, client.getZeppelinConnection("AAAA", "anonymous", "anonymous"));
// Connection to note BBBB
Session connectionB = client.getZeppelinConnection("BBBB");
Session connectionB = client.getZeppelinConnection("BBBB", "anonymous", "anonymous");
assertNotNull(connectionB);
assertTrue(connectionB.isOpen());
assertEquals(client.countConnectedNotes(), 2);
assertEquals(connectionB, client.getZeppelinConnection("BBBB"));
assertEquals(connectionB, client.getZeppelinConnection("BBBB", "anonymous", "anonymous"));
// Remove connection to note AAAA
client.removeNoteConnection("AAAA");
assertEquals(client.countConnectedNotes(), 1);
assertNotEquals(connectionA, client.getZeppelinConnection("AAAA"));
assertNotEquals(connectionA, client.getZeppelinConnection("AAAA", "anonymous", "anonymous"));
assertEquals(client.countConnectedNotes(), 2);
client.stop();
}