From c4bac17b830df8eee90a8c7725e1410c5412cfe9 Mon Sep 17 00:00:00 2001
From: mohitdeuex <105265192+mohitdeuex@users.noreply.github.com>
Date: Wed, 29 Jun 2022 23:20:54 +0530
Subject: [PATCH] [Backend][taskNotifications]] Modfied for webscoket only
 transport and added logs (#5730)

---
 .../catalog/socket/WebSocketManager.java      | 36 ++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/catalog-rest-service/src/main/java/org/openmetadata/catalog/socket/WebSocketManager.java b/catalog-rest-service/src/main/java/org/openmetadata/catalog/socket/WebSocketManager.java
index 61c4a35c11d..7bb3dad0b0d 100644
--- a/catalog-rest-service/src/main/java/org/openmetadata/catalog/socket/WebSocketManager.java
+++ b/catalog-rest-service/src/main/java/org/openmetadata/catalog/socket/WebSocketManager.java
@@ -32,7 +32,15 @@ public class WebSocketManager {
         "connection",
         args -> {
           SocketIoSocket socket = (SocketIoSocket) args[0];
-          String userId = socket.getInitialHeaders().get("UserId").get(0);
+          final String userId;
+          String tempId;
+          try {
+            tempId = socket.getInitialHeaders().get("UserId").get(0);
+          } catch (Exception ex) {
+            tempId = socket.getInitialQuery().get("userId");
+          }
+          userId = tempId;
+
           if (userId != null && !userId.equals("")) {
             LOG.info(
                 "Client :"
@@ -57,6 +65,27 @@ public class WebSocketManager {
                   allUserConnection.remove(socket.getId());
                   activityFeedEndpoints.put(id, allUserConnection);
                 });
+            socket.on(
+                "connect_error",
+                args1 -> {
+                  LOG.error(
+                      "Connection ERROR for user:"
+                          + userId
+                          + "with Remote Address :"
+                          + socket.getInitialHeaders().get("RemoteAddress")
+                          + " disconnected.");
+                });
+            socket.on(
+                "connect_failed",
+                args1 -> {
+                  LOG.error(
+                      "Connection failed ERROR for user:"
+                          + userId
+                          + "with Remote Address :"
+                          + socket.getInitialHeaders().get("RemoteAddress")
+                          + " disconnected.");
+                });
+
             UUID id = UUID.fromString(userId);
             Map<String, SocketIoSocket> userSocketConnections;
             userSocketConnections =
@@ -65,6 +94,11 @@ public class WebSocketManager {
             activityFeedEndpoints.put(id, userSocketConnections);
           }
         });
+    ns.on(
+        "error",
+        args -> {
+          LOG.error("Connection error on the server");
+        });
   }
 
   public static WebSocketManager getInstance() {