Refactor locking function to reduce code duplication

internal/service/policy/service.go

@@ -13,8 +13,7 @@ import (
 
 type Storage interface {
 	Policy(ctx context.Context, name, group, version string) (*storage.Policy, error)
-	LockPolicy(ctx context.Context, name, group, version string) error
-	UnlockPolicy(ctx context.Context, name, group, version string) error
+	SetPolicyLock(ctx context.Context, name, group, version string, lock bool) error
 }
 
 type Service struct {
@@ -100,7 +99,7 @@ func (s *Service) Lock(ctx context.Context, req *policy.LockRequest) error {
 		return errors.New(errors.Forbidden, "policy is already locked")
 	}
 
-	if err := s.storage.LockPolicy(ctx, req.PolicyName, req.Group, req.Version); err != nil {
+	if err := s.storage.SetPolicyLock(ctx, req.PolicyName, req.Group, req.Version, true); err != nil {
 		logger.Error("error locking policy", zap.Error(err))
 		return errors.New("error locking policy", err)
 	}
@@ -131,7 +130,7 @@ func (s *Service) Unlock(ctx context.Context, req *policy.UnlockRequest) error {
 		return errors.New(errors.Forbidden, "policy is unlocked")
 	}
 
-	if err := s.storage.UnlockPolicy(ctx, req.PolicyName, req.Group, req.Version); err != nil {
+	if err := s.storage.SetPolicyLock(ctx, req.PolicyName, req.Group, req.Version, false); err != nil {
 		logger.Error("error unlocking policy", zap.Error(err))
 		return errors.New("error unlocking policy", err)
 	}

internal/storage/policies_tmp_store.go

@@ -11,7 +11,7 @@ package storage
 //		Group:       "example",
 //		Version:     "1.0",
 //		Locked:      false,
-//		LastUpdated: time.Now(),
+//		LastUpdate: time.Now(),
 //		Rego: `
 //			package gaiax
 //

internal/storage/storage.go

@@ -12,31 +12,27 @@ import (
 )
 
 type Policy struct {
-	Filename    string
-	Name        string
-	Group       string
-	Version     string
-	Rego        string
-	Locked      bool
-	LastUpdated time.Time
+	Filename   string
+	Name       string
+	Group      string
+	Version    string
+	Rego       string
+	Locked     bool
+	LastUpdate time.Time
 }
 
 type Storage struct {
-	db         *mongo.Client
-	dbname     string
-	collection string
+	policy *mongo.Collection
 }
 
 func New(db *mongo.Client, dbname, collection string) *Storage {
 	return &Storage{
-		db:         db,
-		dbname:     dbname,
-		collection: collection,
+		policy: db.Database(dbname).Collection(collection),
 	}
 }
 
 func (s *Storage) Policy(ctx context.Context, name, group, version string) (*Policy, error) {
-	result := s.db.Database(s.dbname).Collection(s.collection).FindOne(ctx, bson.M{
+	result := s.policy.FindOne(ctx, bson.M{
 		"name":    name,
 		"group":   group,
 		"version": version,
@@ -57,8 +53,8 @@ func (s *Storage) Policy(ctx context.Context, name, group, version string) (*Pol
 	return &policy, nil
 }
 
-func (s *Storage) LockPolicy(ctx context.Context, name, group, version string) error {
-	_, err := s.db.Database(s.dbname).Collection(s.collection).UpdateOne(
+func (s *Storage) SetPolicyLock(ctx context.Context, name, group, version string, lock bool) error {
+	_, err := s.policy.UpdateOne(
 		ctx,
 		bson.M{
 			"name":    name,
@@ -67,24 +63,8 @@ func (s *Storage) LockPolicy(ctx context.Context, name, group, version string) e
 		},
 		bson.M{
 			"$set": bson.M{
-				"locked": true,
-			},
-		},
-	)
-	return err
-}
-
-func (s *Storage) UnlockPolicy(ctx context.Context, name, group, version string) error {
-	_, err := s.db.Database(s.dbname).Collection(s.collection).UpdateOne(
-		ctx,
-		bson.M{
-			"name":    name,
-			"group":   group,
-			"version": version,
-		},
-		bson.M{
-			"$set": bson.M{
-				"locked": false,
+				"locked":     lock,
+				"lastUpdate": time.Now(),
 			},
 		},
 	)