From 0084e4e590890e6cb44aa1bb713c63a1ecde7d0e Mon Sep 17 00:00:00 2001
From: Anqi <anqi.wang@vesoft.com>
Date: Mon, 4 Mar 2024 11:07:31 +0800
Subject: [PATCH] support auth for scan

---
 nebula3/sclient/GraphStorageClient.py | 19 +++++++++++++++++++
 tests/test_graph_storage_client.py    |  1 +
 2 files changed, 20 insertions(+)

diff --git a/nebula3/sclient/GraphStorageClient.py b/nebula3/sclient/GraphStorageClient.py
index dbf6b888..702a7073 100644
--- a/nebula3/sclient/GraphStorageClient.py
+++ b/nebula3/sclient/GraphStorageClient.py
@@ -34,6 +34,8 @@ class GraphStorageClient(object):
     DEFAULT_START_TIME = 0
     DEFAULT_END_TIME = sys.maxsize
     DEFAULT_LIMIT = 1000
+    user = ""
+    passwd = ""
 
     def __init__(self, meta_cache, storage_addrs=None, time_out=60000):
         self._meta_cache = meta_cache
@@ -42,6 +44,15 @@ def __init__(self, meta_cache, storage_addrs=None, time_out=60000):
         self._connections = []
         self._create_connection()
 
+    def set_userPasswd(self, user, passwd):
+        """set user and password for scan. only useful for enterprise
+
+        :return:
+        """
+        self.user = user
+        self.passwd = passwd
+
+
     def get_conns(self):
         """get all connections which connect to storaged, the ScanResult use it
 
@@ -225,6 +236,9 @@ def _scan_vertex(
         req.filter = where
         req.only_latest_version = only_latest_version
         req.enable_read_from_follower = enable_read_from_follower
+        req.username = self.user.encode('utf-8')
+        req.password = self.passwd.encode('utf-8')
+        req.need_authenticate = True
         return ScanResult(
             self,
             req=req,
@@ -368,6 +382,11 @@ def _scan_edge(
         req.filter = where
         req.only_latest_version = only_latest_version
         req.enable_read_from_follower = enable_read_from_follower
+        if self.user is not "":
+            req.username = self.user.encode('utf-8')
+        if self.passwd is not "":
+            req.password = self.passwd.encode('utf-8')
+        req.need_authenticate = True
         return ScanResult(
             self,
             req=req,
diff --git a/tests/test_graph_storage_client.py b/tests/test_graph_storage_client.py
index 99d7a0bb..5721459f 100644
--- a/tests/test_graph_storage_client.py
+++ b/tests/test_graph_storage_client.py
@@ -87,6 +87,7 @@ def setup_class(cls):
                 50000,
             )
             cls.graph_storage_client = GraphStorageClient(meta_cache)
+            cls.graph_storage_client.set_userPasswd("root", "nebula")
 
         except Exception:
             import traceback