[ZEPPELIN-1437, 1438] Multi-user note management - user aware reload broadcast

zeppelin-server/src/main/java/org/apache/zeppelin/rest/NotebookRestApi.java

@@ -162,6 +162,7 @@ public Response putNotePermissions(@PathParam("noteId") String noteId, String re
     AuthenticationInfo subject = new AuthenticationInfo(SecurityUtils.getPrincipal());
     note.persist(subject);
     notebookServer.broadcastNote(note);
+    notebookServer.broadcastNoteList(subject);
     return new JsonResponse<>(Status.OK).build();
   }
 

zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java

@@ -92,7 +92,7 @@ public ZeppelinServer() throws Exception {
         notebookWsServer, heliumApplicationFactory, depResolver);
     this.notebookRepo = new NotebookRepoSync(conf);
     this.notebookIndex = new LuceneSearch();
-    this.notebookAuthorization = new NotebookAuthorization(conf);
+    this.notebookAuthorization = NotebookAuthorization.init(conf);
     this.credentials = new Credentials(conf.credentialsPersist(), conf.getCredentialsPath());
     notebook = new Notebook(conf,
         notebookRepo, schedulerFactory, replFactory, notebookWsServer,

zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java

@@ -59,6 +59,7 @@
 import java.net.URISyntaxException;
 import java.net.UnknownHostException;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 /**
@@ -85,6 +86,8 @@ String getKey() {
   Gson gson = new GsonBuilder().setDateFormat("yyyy-MM-dd'T'HH:mm:ssZ").create();
   final Map<String, List<NotebookSocket>> noteSocketMap = new HashMap<>();
   final Queue<NotebookSocket> connectedSockets = new ConcurrentLinkedQueue<>();
+  final Map<String, Queue<NotebookSocket>> userConnectedSockets = 
+    new ConcurrentHashMap<String, Queue<NotebookSocket>>();
 
   private Notebook notebook() {
     return ZeppelinServer.notebook;
@@ -160,6 +163,9 @@ public void onMessage(NotebookSocket conn, String msg) {
           userAndRoles.addAll(roles);
         }
       }
+      if (StringUtils.isEmpty(conn.getUser())) {
+        addUserConnection(messagereceived.principal, conn);
+      }
       AuthenticationInfo subject = new AuthenticationInfo(messagereceived.principal);
 
       /** Lets be elegant here */
@@ -264,6 +270,26 @@ public void onClose(NotebookSocket conn, int code, String reason) {
         .getRemoteAddr(), conn.getRequest().getRemotePort(), code, reason);
     removeConnectionFromAllNote(conn);
     connectedSockets.remove(conn);
+    removeUserConnection(conn.getUser(), conn);
+  }
+
+  private void removeUserConnection(String user, NotebookSocket conn) {
+    if (userConnectedSockets.containsKey(user)) {
+      userConnectedSockets.get(user).remove(conn);
+    } else {
+      LOG.warn("Closing connection that is absent in user connections");
+    }
+  }
+
+  private void addUserConnection(String user, NotebookSocket conn) {
+    conn.setUser(user);
+    if (userConnectedSockets.containsKey(user)) {
+      userConnectedSockets.get(user).add(conn);
+    } else {
+      Queue<NotebookSocket> socketQueue = new ConcurrentLinkedQueue<>();
+      socketQueue.add(conn);
+      userConnectedSockets.put(user, socketQueue);
+    }
   }
 
   protected Message deserializeMessage(String msg) {
@@ -379,8 +405,12 @@ private void broadcastExcept(String noteId, Message m, NotebookSocket exclude) {
     }
   }
 
-  private void broadcastAll(Message m) {
-    for (NotebookSocket conn : connectedSockets) {
+  private void multicastToUser(String user, Message m) {
+    if (!userConnectedSockets.containsKey(user)) {
+      LOG.warn("Broadcasting to user that is not in connections map");
+      return;
+    }
+    for (NotebookSocket conn: userConnectedSockets.get(user)) {
       try {
         conn.send(serializeMessage(m));
       } catch (IOException e) {
@@ -472,6 +502,7 @@ public List<Map<String, String>> generateNotebooksInfo(boolean needsReload,
         LOG.error("Fail to reload notes from repository", e);
       }
     }
+
     List<Note> notes = notebook.getAllNotes(subject);
     List<Map<String, String>> notesInfo = new LinkedList<>();
     for (Note note : notes) {
@@ -500,8 +531,20 @@ public void broadcastInterpreterBindings(String noteId,
   }
 
   public void broadcastNoteList(AuthenticationInfo subject) {
+    if (subject == null) {
+      subject = new AuthenticationInfo(StringUtils.EMPTY);
+    }
+    //send first to requesting user
     List<Map<String, String>> notesInfo = generateNotebooksInfo(false, subject);
-    broadcastAll(new Message(OP.NOTES_INFO).put("notes", notesInfo));
+    multicastToUser(subject.getUser(), new Message(OP.NOTES_INFO).put("notes", notesInfo));
+    //to others afterwards
+    for (String user: userConnectedSockets.keySet()) {
+      if (subject.getUser() == user) {
+        continue;
+      }
+      notesInfo = generateNotebooksInfo(false, new AuthenticationInfo(user));
+      multicastToUser(user, new Message(OP.NOTES_INFO).put("notes", notesInfo));
+    }
   }
 
   public void unicastNoteList(NotebookSocket conn, AuthenticationInfo subject) {
@@ -510,8 +553,21 @@ public void unicastNoteList(NotebookSocket conn, AuthenticationInfo subject) {
   }
 
   public void broadcastReloadedNoteList(AuthenticationInfo subject) {
+    if (subject == null) {
+      subject = new AuthenticationInfo(StringUtils.EMPTY);
+    }
+    //reload and reply first to requesting user
     List<Map<String, String>> notesInfo = generateNotebooksInfo(true, subject);
-    broadcastAll(new Message(OP.NOTES_INFO).put("notes", notesInfo));
+    multicastToUser(subject.getUser(), new Message(OP.NOTES_INFO).put("notes", notesInfo));
+    //to others afterwards
+    for (String user: userConnectedSockets.keySet()) {
+      if (subject.getUser() == user) {
+        continue;
+      }
+      //reloaded already above; parameter - false
+      notesInfo = generateNotebooksInfo(false, new AuthenticationInfo(user));
+      multicastToUser(user, new Message(OP.NOTES_INFO).put("notes", notesInfo));
+    }
   }
 
   void permissionError(NotebookSocket conn, String op,

zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookSocket.java

@@ -20,6 +20,7 @@
 
 import javax.servlet.http.HttpServletRequest;
 
+import org.apache.commons.lang.StringUtils;
 import org.eclipse.jetty.websocket.api.Session;
 import org.eclipse.jetty.websocket.api.WebSocketAdapter;
 
@@ -32,12 +33,14 @@ public class NotebookSocket extends WebSocketAdapter {
   private NotebookSocketListener listener;
   private HttpServletRequest request;
   private String protocol;
+  private String user;
 
   public NotebookSocket(HttpServletRequest req, String protocol,
       NotebookSocketListener listener) {
     this.listener = listener;
     this.request = req;
     this.protocol = protocol;
+    this.user = StringUtils.EMPTY;
   }
 
   @Override
@@ -69,4 +72,11 @@ public void send(String serializeMessage) throws IOException {
     connection.getRemote().sendString(serializeMessage);
   }
 
+  public String getUser() {
+    return user;
+  }
+
+  public void setUser(String user) {
+    this.user = user;
+  }
 }

zeppelin-server/src/test/java/org/apache/zeppelin/rest/ZeppelinRestApiTest.java

@@ -32,6 +32,7 @@
 import org.apache.zeppelin.notebook.Paragraph;
 import org.apache.zeppelin.scheduler.Job.Status;
 import org.apache.zeppelin.server.ZeppelinServer;
+import org.apache.zeppelin.user.AuthenticationInfo;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.FixMethodOrder;
@@ -341,7 +342,9 @@ public void testListNotebooks() throws IOException {
     Map<String, Object> resp = gson.fromJson(get.getResponseBodyAsString(), new TypeToken<Map<String, Object>>() {
     }.getType());
     List<Map<String, String>> body = (List<Map<String, String>>) resp.get("body");
-    assertEquals("List notebooks are equal", ZeppelinServer.notebook.getAllNotes().size(), body.size());
+    //TODO(khalid): anonymous or specific user notes?
+    AuthenticationInfo subject = new AuthenticationInfo("anonymous");
+    assertEquals("List notebooks are equal", ZeppelinServer.notebook.getAllNotes(subject).size(), body.size());
     get.releaseConnection();
   }
 

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/Notebook.java

@@ -31,14 +31,12 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
-
 import com.google.common.base.Predicate;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.Sets;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.google.gson.stream.JsonReader;
-import org.apache.commons.codec.binary.StringUtils;
 import org.quartz.CronScheduleBuilder;
 import org.quartz.CronTrigger;
 import org.quartz.JobBuilder;
@@ -484,6 +482,7 @@ public void reloadAllNotes(AuthenticationInfo subject) throws IOException {
     }
 
     List<NoteInfo> noteInfos = notebookRepo.list(subject);
+
     for (NoteInfo info : noteInfos) {
       loadNoteFromRepo(info.getId(), subject);
     }
@@ -534,7 +533,7 @@ public int compare(Note note1, Note note2) {
       return noteList;
     }
   }
-
+  
   public List<Note> getAllNotes(AuthenticationInfo subject) {
     final Set<String> entities = Sets.newHashSet();
     if (subject != null) {

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/NotebookAuthorization.java

@@ -17,9 +17,13 @@
 
 package org.apache.zeppelin.notebook;
 
+import com.google.common.base.Predicate;
+import com.google.common.collect.FluentIterable;
+import com.google.common.collect.Sets;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import org.apache.zeppelin.conf.ZeppelinConfiguration;
+import org.apache.zeppelin.user.AuthenticationInfo;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -31,29 +35,44 @@
  */
 public class NotebookAuthorization {
   private static final Logger LOG = LoggerFactory.getLogger(NotebookAuthorization.class);
-
+  private static NotebookAuthorization instance = null;
   /*
    * { "note1": { "owners": ["u1"], "readers": ["u1", "u2"], "writers": ["u1"] },  "note2": ... } }
    */
-  private Map<String, Map<String, Set<String>>> authInfo = new HashMap<>();
-  private ZeppelinConfiguration conf;
-  private Gson gson;
-  private String filePath;
-
-  public NotebookAuthorization(ZeppelinConfiguration conf) {
-    this.conf = conf;
-    filePath = conf.getNotebookAuthorizationPath();
-    GsonBuilder builder = new GsonBuilder();
-    builder.setPrettyPrinting();
-    gson = builder.create();
-    try {
-      loadFromFile();
-    } catch (IOException e) {
-      LOG.error("Error loading NotebookAuthorization", e);
+  private static Map<String, Map<String, Set<String>>> authInfo = new HashMap<>();
+  private static ZeppelinConfiguration conf;
+  private static Gson gson;
+  private static String filePath;
+
+  private NotebookAuthorization() {}
+
+  public static NotebookAuthorization init(ZeppelinConfiguration config) {
+    if (instance == null) {
+      instance = new NotebookAuthorization();
+      conf = config;
+      filePath = conf.getNotebookAuthorizationPath();
+      GsonBuilder builder = new GsonBuilder();
+      builder.setPrettyPrinting();
+      gson = builder.create();
+      try {
+        loadFromFile();
+      } catch (IOException e) {
+        LOG.error("Error loading NotebookAuthorization", e);
+      }
     }
+    return instance;
   }
 
-  private void loadFromFile() throws IOException {
+  public static NotebookAuthorization getInstance() {
+    if (instance == null) {
+      LOG.warn("Notebook authorization module was called without initialization,"
+          + " initializing with default configuration");
+      init(ZeppelinConfiguration.create());
+    }
+    return instance;
+  }
+
+  private static void loadFromFile() throws IOException {
     File settingFile = new File(filePath);
     LOG.info(settingFile.getAbsolutePath());
     if (!settingFile.exists()) {
@@ -74,7 +93,7 @@ private void loadFromFile() throws IOException {
     String json = sb.toString();
     NotebookAuthorizationInfoSaving info = gson.fromJson(json,
             NotebookAuthorizationInfoSaving.class);
-    this.authInfo = info.authInfo;
+    authInfo = info.authInfo;
   }
 
   private void saveToFile() {
@@ -225,4 +244,16 @@ public void removeNote(String noteId) {
     saveToFile();
   }
 
+  public List<NoteInfo> filterByUser(List<NoteInfo> notes, AuthenticationInfo subject) {
+    final Set<String> entities = Sets.newHashSet();
+    if (subject != null) {
+      entities.add(subject.getUser());
+    }
+    return FluentIterable.from(notes).filter(new Predicate<NoteInfo>() {
+      @Override
+      public boolean apply(NoteInfo input) {
+        return input != null && isReader(input.getId(), entities);
+      }
+    }).toList();
+  }
 }

zeppelin-zengine/src/main/java/org/apache/zeppelin/notebook/repo/NotebookRepoSync.java

@@ -31,6 +31,7 @@
 import org.apache.zeppelin.conf.ZeppelinConfiguration.ConfVars;
 import org.apache.zeppelin.notebook.Note;
 import org.apache.zeppelin.notebook.NoteInfo;
+import org.apache.zeppelin.notebook.NotebookAuthorization;
 import org.apache.zeppelin.notebook.Paragraph;
 import org.apache.zeppelin.user.AuthenticationInfo;
 import org.slf4j.Logger;
@@ -178,9 +179,11 @@ public void remove(String noteId, AuthenticationInfo subject) throws IOException
    */
   void sync(int sourceRepoIndex, int destRepoIndex, AuthenticationInfo subject) throws IOException {
     LOG.info("Sync started");
+    NotebookAuthorization auth = NotebookAuthorization.getInstance();
     NotebookRepo srcRepo = getRepo(sourceRepoIndex);
     NotebookRepo dstRepo = getRepo(destRepoIndex);
-    List <NoteInfo> srcNotes = srcRepo.list(subject);
+    List <NoteInfo> allSrcNotes = srcRepo.list(subject);
+    List <NoteInfo> srcNotes = auth.filterByUser(allSrcNotes, subject);
     List <NoteInfo> dstNotes = dstRepo.list(subject);
 
     Map<String, List<String>> noteIDs = notesCheckDiff(srcNotes, srcRepo, dstNotes, dstRepo);

zeppelin-zengine/src/test/java/org/apache/zeppelin/helium/HeliumApplicationFactoryTest.java

@@ -89,7 +89,7 @@ public void setUp() throws Exception {
 
     SearchService search = mock(SearchService.class);
     notebookRepo = new VFSNotebookRepo(conf);
-    NotebookAuthorization notebookAuthorization = new NotebookAuthorization(conf);
+    NotebookAuthorization notebookAuthorization = NotebookAuthorization.init(conf);
     notebook = new Notebook(
         conf,
         notebookRepo,