gNOI Reboot changes

sonic-net · Oct 9, 2024 · 5cc436a · 5cc436a
1 parent 0b3232f
commit 5cc436a
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 7 deletions.
diff --git a/common_utils/component_state_helper.go b/common_utils/component_state_helper.go
@@ -5,19 +5,29 @@ import (
 
 	sdcfg "github.com/sonic-net/sonic-gnmi/sonic_db_config"
 	"github.com/go-redis/redis"
-	log "github.com/golang/glog"
 )
 
 const (
 	dbName              = "STATE_DB"
 )
 
 func getRedisDBClient() (*redis.Client, error) {
+	ns, _ := sdcfg.GetDbDefaultNamespace()
+	addr, err := sdcfg.GetDbTcpAddr(dbName, ns)
+	if err != nil {
+		log.Errorf("Addr err: %v", err)
+		return
+	}
+	db, err := sdcfg.GetDbId("STATE_DB", ns)
+	if err != nil {
+		log.Errorf("DB err: %v", err)
+		return
+	}
 	rclient := redis.NewClient(&redis.Options{
 		Network:     "tcp",
-		Addr:        sdcfg.GetDbTcpAddr(dbName),
+		Addr:        addr,
 		Password:    "", // no password set
-		DB:          sdcfg.GetDbId(dbName),
+		DB:          db,
 		DialTimeout: 0,
 	})
 	if rclient == nil {

diff --git a/gnmi_server/gnoi_system.go b/gnmi_server/gnoi_system.go
@@ -244,7 +244,7 @@ func (srv *Server) Reboot(ctx context.Context, req *syspb.RebootRequest) (*syspb
 		return nil, status.Errorf(codes.InvalidArgument, err.Error())
 	}
 	// Initialize State DB.
-	rclient, err := getRedisDBClient(stateDB)
+	rclient, err := getRedisDBClient()
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, err.Error())
 	}
@@ -276,7 +276,7 @@ func (srv *Server) RebootStatus(ctx context.Context, req *syspb.RebootStatusRequ
 	log.V(1).Info("gNOI: RebootStatus")
 	resp := &syspb.RebootStatusResponse{}
 	// Initialize State DB.
-	rclient, err := getRedisDBClient(stateDB)
+	rclient, err := getRedisDBClient()
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, err.Error())
 	}
@@ -311,7 +311,7 @@ func (srv *Server) CancelReboot(ctx context.Context, req *syspb.CancelRebootRequ
 		return nil, status.Errorf(codes.Internal, "Invalid CancelReboot request: message is empty.")
 	}
 	// Initialize State DB.
-	rclient, err := getRedisDBClient(stateDB)
+	rclient, err := getRedisDBClient()
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, err.Error())
 	}

diff --git a/gnmi_server/gnoi_system_test.go b/gnmi_server/gnoi_system_test.go
@@ -117,7 +117,7 @@ func TestSystem(t *testing.T) {
 	defer cancel()
 
 	sc := syspb.NewSystemClient(conn)
-	rclient, err := getRedisDBClient(stateDB)
+	rclient, err := getRedisDBClient()
 	if err != nil {
 		t.Fatalf("Cannot connect to the redis server: %v.", err.Error())
 	}