Add milvus-client-id in the headers (#682)

* connect with clientside id Signed-off-by: ryjiang <jiangruiyi@gmail.com> * Add milvus-client-id in headers * chagne debug => info Signed-off-by: ryjiang <jiangruiyi@gmail.com> --------- Signed-off-by: ryjiang <jiangruiyi@gmail.com>
zilliztech · Nov 8, 2024 · 5fc260a · 5fc260a
1 parent abb0d09
commit 5fc260a
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 16 deletions.
diff --git a/client/src/context/Auth.tsx b/client/src/context/Auth.tsx
@@ -64,7 +64,10 @@ export const AuthProvider = (props: { children: React.ReactNode }) => {
   // login API
   const login = async (params: AuthReq) => {
     // create a new client id
-    params.clientId = Math.random().toString(36).substring(16);
+    params.clientId = Math.random().toString(36).substring(7);
+    // save clientId to local storage
+    // console.log('params.clientId', params.clientId);
+    window.localStorage.setItem(MILVUS_CLIENT_ID, params.clientId);
     // connect to Milvus
     const res = await MilvusService.connect(params);
     // update auth request

diff --git a/client/src/context/Data.tsx b/client/src/context/Data.tsx
@@ -410,10 +410,14 @@ export const DataProvider = (props: { children: React.ReactNode }) => {
       // update database get from auth
       setDatabase(authReq.database);
 
-      // create socket
-      socket.current = isElectron ? io(url as string) : io();
-      // register client
-      socket.current.emit(WS_EVENTS.REGISTER, clientId);
+      const extraHeaders = {
+        'milvus-client-id': clientId,
+        'x-attu-database': authReq.database,
+      };
+
+      socket.current = isElectron
+        ? io(url as string, { extraHeaders })
+        : io({ extraHeaders });
 
       socket.current.on('connect', async () => {
         // console.info('--- ws connected ---', clientId);

diff --git a/server/src/milvus/milvus.service.ts b/server/src/milvus/milvus.service.ts
@@ -38,6 +38,15 @@ export class MilvusService {
     // Format the address to remove the http prefix
     const milvusAddress = MilvusService.formatAddress(address);
 
+    // if client exists, return the client
+    if (clientCache.has(clientId)) {
+      const cache = clientCache.get(clientId);
+      return {
+        clientId: cache.milvusClient.clientId,
+        database: cache.database,
+      };
+    }
+
     try {
       // Create a new Milvus client with the provided connection details
       const clientConfig: ClientConfig = {
@@ -147,15 +156,18 @@ export class MilvusService {
   }
 
   async closeConnection(clientId: string): Promise<CONNECT_STATUS> {
-    const { milvusClient } = clientCache.get(clientId);
+    if (clientCache.has(clientId)) {
+      const { milvusClient } = clientCache.get(clientId);
 
-    const res = await milvusClient.closeConnection();
-    // clear cache on disconnect
-    clientCache.delete(milvusClient.clientId);
-    // clear crons
-    cronsManager.deleteCronJob(clientId);
+      console.info('Deleting client', clientId);
 
-    return res;
+      const res = await milvusClient.closeConnection();
+      // clear cache on disconnect
+      clientCache.delete(milvusClient.clientId);
+      // clear crons
+      cronsManager.deleteCronJob(clientId);
+      return res;
+    }
   }
 
   async useDatabase(clientId: string, db: string) {

diff --git a/server/src/socket.ts b/server/src/socket.ts
@@ -2,7 +2,7 @@
 import { Server, Socket } from 'socket.io';
 import * as http from 'http';
 import chalk from 'chalk';
-import { WS_EVENTS, isElectron } from './utils';
+import { isElectron } from './utils';
 
 export let io: Server;
 export let clients = new Map<string, Socket>();
@@ -16,8 +16,9 @@ export function initWebSocket(server: http.Server) {
   });
 
   io.on('connection', (socket: Socket) => {
-    // register client
-    socket.on(WS_EVENTS.REGISTER, (clientId: string) => {
+    const clientId = socket.handshake.headers['milvus-client-id'] as string;
+
+    if (clientId) {
       clients.set(clientId, socket);
       if (!isElectron()) {
         console.info(chalk.green(`ws client connected ${clientId}`));
@@ -37,7 +38,7 @@ export function initWebSocket(server: http.Server) {
           );
         }
       });
-    });
+    }
   });
 
   // Handle server-level errors