Add SaveOnSet

Author: rlucus

gnmi_server/server.go

@@ -5,16 +5,17 @@ import (
 	"errors"
 	"fmt"
 	"github.com/Azure/sonic-mgmt-common/translib"
-	"github.com/sonic-net/sonic-gnmi/common_utils"
-	spb "github.com/sonic-net/sonic-gnmi/proto"
-	spb_gnoi "github.com/sonic-net/sonic-gnmi/proto/gnoi"
-	spb_jwt_gnoi "github.com/sonic-net/sonic-gnmi/proto/gnoi/jwt"
-	sdc "github.com/sonic-net/sonic-gnmi/sonic_data_client"
+	"github.com/Azure/sonic-mgmt-common/translib/transformer"
 	log "github.com/golang/glog"
 	"github.com/golang/protobuf/proto"
 	gnmipb "github.com/openconfig/gnmi/proto/gnmi"
 	gnmi_extpb "github.com/openconfig/gnmi/proto/gnmi_ext"
 	gnoi_system_pb "github.com/openconfig/gnoi/system"
+	"github.com/sonic-net/sonic-gnmi/common_utils"
+	spb "github.com/sonic-net/sonic-gnmi/proto"
+	spb_gnoi "github.com/sonic-net/sonic-gnmi/proto/gnoi"
+	spb_jwt_gnoi "github.com/sonic-net/sonic-gnmi/proto/gnoi/jwt"
+	sdc "github.com/sonic-net/sonic-gnmi/sonic_data_client"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -39,6 +40,10 @@ type Server struct {
 	config  *Config
 	cMu     sync.Mutex
 	clients map[string]*Client
+	// SaveStartupConfig points to a function that is called to save changes of
+	// configuration to a file. By default it points to an empty function -
+	// the configuration is not saved to a file.
+	SaveStartupConfig func()
 }
 type AuthTypes map[string]bool
 
@@ -133,9 +138,10 @@ func NewServer(config *Config, opts []grpc.ServerOption) (*Server, error) {
 	reflection.Register(s)
 
 	srv := &Server{
-		s:       s,
-		config:  config,
-		clients: map[string]*Client{},
+		s:                 s,
+		config:            config,
+		clients:           map[string]*Client{},
+		SaveStartupConfig: SaveOnSetDisabled,
 	}
 	var err error
 	if srv.config.Port < 0 {
@@ -150,7 +156,7 @@ func NewServer(config *Config, opts []grpc.ServerOption) (*Server, error) {
 	if srv.config.EnableTranslibWrite || srv.config.EnableNativeWrite {
 		gnoi_system_pb.RegisterSystemServer(srv.s, srv)
 	}
-	if srv.config.EnableTranslibWrite {		
+	if srv.config.EnableTranslibWrite {
 		spb_gnoi.RegisterSonicServiceServer(srv.s, srv)
 	}
 	log.V(1).Infof("Created Server on %s, read-only: %t", srv.Address(), !srv.config.EnableTranslibWrite)
@@ -376,6 +382,18 @@ func (s *Server) Get(ctx context.Context, req *gnmipb.GetRequest) (*gnmipb.GetRe
 	return &gnmipb.GetResponse{Notification: notifications}, nil
 }
 
+// SaveOnSetEnabled saves configuration to a file
+func SaveOnSetEnabled() {
+	if err := transformer.SaveStartupConfig(); err != nil {
+		log.Errorf("Saving startup config failed: %v", err)
+	} else {
+		log.Errorf("Success! Startup config has been saved!")
+	}
+}
+
+// SaveOnSetDisabeld does nothing.
+func SaveOnSetDisabled() {}
+
 func (s *Server) Set(ctx context.Context, req *gnmipb.SetRequest) (*gnmipb.SetResponse, error) {
 	common_utils.IncCounter(common_utils.GNMI_SET)
 	if s.config.EnableTranslibWrite == false && s.config.EnableNativeWrite == false {
@@ -467,6 +485,7 @@ func (s *Server) Set(ctx context.Context, req *gnmipb.SetRequest) (*gnmipb.SetRe
 		common_utils.IncCounter(common_utils.GNMI_SET_FAIL)
 	}
 
+	s.SaveStartupConfig()
 	return &gnmipb.SetResponse{
 		Prefix:   req.GetPrefix(),
 		Response: results,

gnmi_server/server_test.go

@@ -116,6 +116,7 @@ func createServer(t *testing.T, port int64) *Server {
 	if err != nil {
 		t.Errorf("Failed to create gNMI server: %v", err)
 	}
+	s.SaveStartupConfig = SaveOnSetEnabled
 	return s
 }
 
@@ -3095,7 +3096,7 @@ func TestConnectionsKeepAlive(t *testing.T) {
 }
 
 func TestClient(t *testing.T) {
-    // sonic-host:device-test-event is a test event. 
+    // sonic-host:device-test-event is a test event.
     // Events client will drop it on floor.
     events := [] sdc.Evt_rcvd {
         { "test0", 7, 777 },
@@ -3229,7 +3230,7 @@ func TestClient(t *testing.T) {
 }
 
 func TestGnmiSetBatch(t *testing.T) {
-	mockCode := 
+	mockCode :=
 `
 print('No Yang validation for test mode...')
 print('%s')
@@ -3314,7 +3315,7 @@ func TestGNMINative(t *testing.T) {
 		return &dbus.Call{}
 	})
 	defer mock2.Reset()
-	mockCode := 
+	mockCode :=
 `
 print('No Yang validation for test mode...')
 print('%s')

telemetry/telemetry.go

@@ -20,20 +20,21 @@ import (
 
 var (
 	userAuth = gnmi.AuthTypes{"password": false, "cert": false, "jwt": false}
-	port = flag.Int("port", -1, "port to listen on")
+	port     = flag.Int("port", -1, "port to listen on")
 	// Certificate files.
-	caCert            = flag.String("ca_crt", "", "CA certificate for client certificate validation. Optional.")
-	serverCert        = flag.String("server_crt", "", "TLS server certificate")
-	serverKey         = flag.String("server_key", "", "TLS server private key")
-	insecure          = flag.Bool("insecure", false, "Skip providing TLS cert and key, for testing only!")
-	noTLS             = flag.Bool("noTLS", false, "disable TLS, for testing only!")
-	allowNoClientCert = flag.Bool("allow_no_client_auth", false, "When set, telemetry server will request but not require a client certificate.")
-	jwtRefInt         = flag.Uint64("jwt_refresh_int", 900, "Seconds before JWT expiry the token can be refreshed.")
-	jwtValInt         = flag.Uint64("jwt_valid_int", 3600, "Seconds that JWT token is valid for.")
+	caCert              = flag.String("ca_crt", "", "CA certificate for client certificate validation. Optional.")
+	serverCert          = flag.String("server_crt", "", "TLS server certificate")
+	serverKey           = flag.String("server_key", "", "TLS server private key")
+	insecure            = flag.Bool("insecure", false, "Skip providing TLS cert and key, for testing only!")
+	noTLS               = flag.Bool("noTLS", false, "disable TLS, for testing only!")
+	allowNoClientCert   = flag.Bool("allow_no_client_auth", false, "When set, telemetry server will request but not require a client certificate.")
+	jwtRefInt           = flag.Uint64("jwt_refresh_int", 900, "Seconds before JWT expiry the token can be refreshed.")
+	jwtValInt           = flag.Uint64("jwt_valid_int", 3600, "Seconds that JWT token is valid for.")
 	gnmi_translib_write = flag.Bool("gnmi_translib_write", gnmi.ENABLE_TRANSLIB_WRITE, "Enable gNMI translib write for management framework")
 	gnmi_native_write   = flag.Bool("gnmi_native_write", gnmi.ENABLE_NATIVE_WRITE, "Enable gNMI native write")
-	threshold         = flag.Int("threshold", 100, "max number of client connections")
-	idle_conn_duration = flag.Int("idle_conn_duration", 5, "Seconds before server closes idle connections")
+	threshold           = flag.Int("threshold", 100, "max number of client connections")
+	idle_conn_duration  = flag.Int("idle_conn_duration", 5, "Seconds before server closes idle connections")
+	withSaveOnSet       = flag.Bool("with-save-on-set", false, "Enables save-on-set.")
 )
 
 func main() {
@@ -176,6 +177,9 @@ func main() {
 		log.Errorf("Failed to create gNMI server: %v", err)
 		return
 	}
+	if *withSaveOnSet {
+		s.SaveStartupConfig = gnmi.SaveOnSetEnabled
+	}
 
 	log.V(1).Infof("Auth Modes: ", userAuth)
 	log.V(1).Infof("Starting RPC server on address: %s", s.Address())