mirror of
https://github.com/apache/zeppelin
synced 2026-05-24 09:38:26 +00:00
fix ws back connection
This commit is contained in:
parent
7e1ac5d74b
commit
5755a5c462
4 changed files with 23 additions and 17 deletions
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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 ");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue