feat: ngql websocket

app/config/service.ts

@@ -1,11 +1,18 @@
 import { _delete, get, post, put } from '../utils/http';
+import ngqlRunner from '@app/utils/websocket';
 
 const service = {
+  // execNGQL: (params, config?) => {
+  //   return post('/api-nebula/db/exec')(params, config);
+  // },
   execNGQL: (params, config?) => {
-    return post('/api-nebula/db/exec')(params, config);
+    return ngqlRunner.runNgql(params, config);
   },
+  // batchExecNGQL: (params, config?) => {
+  //   return post('/api-nebula/db/batchExec')(params, config);
+  // },
   batchExecNGQL: (params, config?) => {
-    return post('/api-nebula/db/batchExec')(params, config);
+    return ngqlRunner.runBatchNgql(params, config);
   },
   execSeqNGQL: (params, config?) => {
     return post('/api-nebula/db/exec_seq')(params, config);

app/stores/global.ts

@@ -3,8 +3,9 @@ import cookies from 'js-cookie';
 import { message } from 'antd';
 import { Base64 } from 'js-base64';
 import { getI18n } from '@vesoft-inc/i18n';
-import service from '@app/config/service';
 import { BrowserHistory } from 'history';
+import service from '@app/config/service';
+import ngqlRunner from '@app/utils/websocket';
 import { getRootStore, resetStore } from '.';
 
 const { intl } = getI18n();
@@ -21,10 +22,14 @@ export class GlobalStore {
   _username = cookies.get('nu');
   _host = cookies.get('nh');
   version = process.env.VERSION;
+
+  ngqlRunner = ngqlRunner;
+
   constructor() {
     makeObservable(this, {
       _username: observable,
       _host: observable,
+      ngqlRunner: observable.ref,
       update: action,
     });
   }
@@ -100,7 +105,9 @@ export class GlobalStore {
       cookies.set('nh', _host);
       cookies.set('nu', username);
       this.update({ _host, _username: username });
-      return true;
+      const socketConncted = await this.ngqlRunner.connect(`ws://${location.host}/nebula_ws`);
+      return socketConncted;
+      // return true;
     }
 
     this.update({ _host: '', _username: '' });

app/stores/graph.ts

@@ -81,7 +81,7 @@ export class GraphStore {
     return edgeMap;
   }
 
-  constructor(data?) {
+  constructor() {
     makeAutoObservable(this, {
       nodeHovering: observable.ref,
       nodeDragging: observable.ref,

app/utils/websocket.ts

@@ -0,0 +1,223 @@
+import { getRootStore } from '@app/stores';
+import { message } from 'antd';
+import { v4 as uuidv4 } from 'uuid';
+import { safeParse } from './function';
+import { HttpResCode } from './http';
+
+interface MessageReceive {
+  header: {
+    msgId: string;
+    sendTime: number;
+  };
+  body: {
+    msgType: string;
+    content: Record<string, unknown>;
+  };
+}
+
+export const WsHeartbeatReq = '1';
+export const WsHeartbeatRes = '2';
+
+export class NgqlRunner {
+  socket: WebSocket | undefined = undefined;
+
+  socketUrl: string | URL | undefined;
+  socketProtocols: string | string[] | undefined;
+
+  socketMessageListeners: Array<(e: MessageEvent) => void> = [];
+
+  socketIsConnecting = false;
+  socketConnectingPromise: Promise<boolean> | undefined;
+  socketPingTimeInterval: number | undefined;
+
+  constructor() {
+    const urlItem = localStorage.getItem('socketUrl');
+    const protocolsItem = localStorage.getItem('socketProtocols');
+
+    urlItem && (this.socketUrl = safeParse<string>(urlItem));
+    protocolsItem && (this.socketProtocols = safeParse<string>(protocolsItem));
+  }
+
+  addSocketMessageListener = (listener: (e: MessageEvent) => void) => {
+    this.socket?.addEventListener('message', listener);
+    this.socketMessageListeners.push(listener);
+  };
+
+  rmSocketMessageListener = (listener: (e: MessageEvent) => void) => {
+    this.socket?.removeEventListener('message', listener);
+    this.socketMessageListeners = this.socketMessageListeners.filter(l => l !== listener);
+  };
+
+  clearSocketMessageListener = () => {
+    this.socketMessageListeners.forEach((l) => {
+      this.socket?.removeEventListener('message', l);
+    });
+    this.socketMessageListeners = [];
+  };
+
+  connect = (url: string | URL, protocols?: string | string[]) => {
+    if (this.socketIsConnecting) {
+      return this.socketConnectingPromise;
+    } else if (this.socket?.readyState === WebSocket.OPEN) {
+      return Promise.resolve(true);
+    }
+    this.socketIsConnecting = true;
+    this.socketConnectingPromise = new Promise<boolean>((resolve) => {
+      const socket = new WebSocket(url, protocols);
+      socket.onopen = () => {
+        console.log('=====ngqlSocket open');
+        this.socket = socket;
+        this.socketUrl = url;
+        this.socketProtocols = protocols;
+
+        localStorage.setItem('socketUrl', JSON.stringify(url));
+        protocols && localStorage.setItem('socketProtocols', JSON.stringify(protocols));
+
+        if (this.socketPingTimeInterval) {
+          clearTimeout(this.socketPingTimeInterval);
+        }
+        this.socketPingTimeInterval = window.setInterval(this.ping, 1000 * 30);
+        this.socketIsConnecting = false;
+        this.socketConnectingPromise = undefined;
+
+        socket.onerror = undefined;
+        socket.onclose = undefined;
+
+        // reconnect
+        this.socket.addEventListener('close', this.onDisconnect);
+        this.socket.addEventListener('error', this.onError);
+
+        resolve(true);
+      };
+      socket.onerror = (e) => {
+        console.error('=====ngqlSocket error', e);
+        this.socketIsConnecting = false;
+        this.socketConnectingPromise = undefined;
+        resolve(false);
+      };
+      socket.onclose = () => {
+        console.log('=====ngqlSocket close');
+        this.socket = undefined;
+      };
+    })
+    return this.socketConnectingPromise;
+  };
+
+  reConnect = () => {
+    return this.connect(this.socketUrl, this.socketProtocols);
+  };
+
+  onError = (e: Event) => {
+    console.error('=====ngqlSocket error', e);
+    message.error('WebSocket error, try to reconnect...');
+    this.onDisconnect();
+  };
+
+  onDisconnect = () => {
+    this.socket?.removeEventListener('close', this.onDisconnect);
+    this.socket?.removeEventListener('error', this.onError);
+    this.clearSocketMessageListener();
+    this.socket?.close();
+
+    clearTimeout(this.socketPingTimeInterval);
+    this.socketPingTimeInterval = undefined;
+    this.socket = undefined;
+
+    this.socketUrl && setTimeout(this.reConnect, 1000);
+  }
+
+  desctory = () => {
+    clearTimeout(this.socketPingTimeInterval);
+    this.clearSocketMessageListener();
+    this.socket?.close();
+    this.socket = undefined;
+    this.socketUrl = undefined;
+    this.socketProtocols = undefined;
+  };
+
+  ping = () => {
+    this.socket?.readyState === WebSocket.OPEN && this.socket.send(WsHeartbeatReq);
+  };
+
+  runNgql = async ({ gql, paramList }: { gql: string; paramList?: string[] }, _config: any) => {
+    const message = {
+      header: {
+        msgId: uuidv4(),
+        version: '1.0',
+      },
+      body: {
+        product: 'Studio',
+        msgType: 'ngql',
+        content: { gql, paramList },
+      },
+    };
+
+    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
+      await this.reConnect();
+    }
+
+    return new Promise((resolve) => {
+      const receiveMsg = (e: MessageEvent<string>) => {
+        if (e.data === WsHeartbeatRes) {
+          return;
+        }
+        const msgReceive = safeParse<MessageReceive>(e.data);
+        if (msgReceive?.body?.content?.code === HttpResCode.ErrSession) {
+          this.desctory();
+          getRootStore().global.logout();
+          return;
+        }
+        if (msgReceive?.header?.msgId === message.header.msgId) {
+          resolve(msgReceive.body.content);
+          this.rmSocketMessageListener(receiveMsg);
+        }
+      };
+
+      this.socket?.send(JSON.stringify(message));
+      this.addSocketMessageListener(receiveMsg);
+    });
+  };
+
+  runBatchNgql = async ({ gqls, paramList }: { gqls: string[]; paramList?: string[] }, _config: any) => {
+    const message = {
+      header: {
+        msgId: uuidv4(),
+        version: '1.0',
+      },
+      body: {
+        product: 'Studio',
+        msgType: 'batch_ngql',
+        content: { gqls, paramList },
+      },
+    };
+
+    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
+      await this.reConnect();
+    }
+
+    return new Promise((resolve) => {
+      const receiveMsg = (e: MessageEvent<string>) => {
+        if (e.data === WsHeartbeatRes) {
+          return;
+        }
+        const msgReceive = safeParse<MessageReceive>(e.data);
+        if (msgReceive?.body?.content?.code === HttpResCode.ErrSession) {
+          this.desctory();
+          getRootStore().global.logout();
+          return;
+        }
+        if (msgReceive?.header?.msgId === message.header.msgId) {
+          resolve(msgReceive.body.content);
+          this.rmSocketMessageListener(receiveMsg);
+        }
+      };
+
+      this.socket?.send(JSON.stringify(message));
+      this.addSocketMessageListener(receiveMsg);
+    });
+  };
+}
+
+const ngqlRunner = new NgqlRunner();
+
+export default ngqlRunner;

config/webpack.dev.js

@@ -62,6 +62,13 @@ const devConfig = {
         target: 'http://127.0.0.1:9000',
         changeOrigin: true,
       },
+      {
+        path: '/nebula_ws',
+        target: 'ws://127.0.0.1:9000',
+        changeOrigin: true,
+        secure: false,
+        ws: true,
+      },
     ]
   },
 };

package.json

@@ -83,6 +83,7 @@
     "@types/react-redux": "^7.1.22",
     "@types/react-router-dom": "^5.1.0",
     "@types/supertest": "^2.0.0",
+    "@types/uuid": "^8.3.4",
     "@types/webpack": "^5.28.0",
     "@typescript-eslint/eslint-plugin": "^5.0.0",
     "@typescript-eslint/parser": "^5.0.0",

server/api/studio/internal/service/gateway.go

@@ -95,29 +95,12 @@ func (s *gatewayService) DisconnectDB() (*types.AnyResponse, error) {
 	return &types.AnyResponse{Data: response.StandardHandlerDataFieldAny(nil)}, nil
 }
 
-func isSessionError(err error) bool {
-	subErrMsgStr := []string{
-		"session expired",
-		"connection refused",
-		"broken pipe",
-		"an existing connection was forcibly closed",
-		"Token is expired",
-		"Session not existed",
-	}
-	for _, subErrMsg := range subErrMsgStr {
-		if strings.Contains(err.Error(), subErrMsg) {
-			return true
-		}
-	}
-	return false
-}
-
 func (s *gatewayService) ExecNGQL(request *types.ExecNGQLParams) (*types.AnyResponse, error) {
 	authData := s.ctx.Value(auth.CtxKeyUserInfo{}).(*auth.AuthData)
 
 	execute, _, err := dao.Execute(authData.NSID, request.Gql, request.ParamList)
 	if err != nil {
-		if isSessionError(err) {
+		if auth.IsSessionError(err) {
 			return nil, ecode.WithSessionMessage(err)
 		}
 		return nil, ecode.WithErrorMessage(ecode.ErrInternalServer, err, "execute failed")
@@ -149,7 +132,7 @@ func (s *gatewayService) BatchExecNGQL(request *types.BatchExecNGQLParams) (*typ
 		execute, _, err := dao.Execute(NSID, gql, make([]string, 0))
 		gqlRes := map[string]interface{}{"gql": gql, "data": execute}
 		if err != nil {
-			if isSessionError(err) {
+			if auth.IsSessionError(err) {
 				return nil, ecode.WithSessionMessage(err)
 			}
 			gqlRes["message"] = err.Error()