extend example

management/refactor/store/sqlite_store.go

@@ -1,6 +1,7 @@
 package store
 
 import (
+	"errors"
 	"path/filepath"
 	"runtime"
 	"sync"
@@ -12,8 +13,17 @@ import (
 	"gorm.io/gorm/clause"
 	"gorm.io/gorm/logger"
 
-	"github.com/netbirdio/netbird/management/refactor/peers"
-	"github.com/netbirdio/netbird/management/refactor/settings"
+	dnsTypes "github.com/netbirdio/netbird/management/refactor/resources/dns/types"
+	groupTypes "github.com/netbirdio/netbird/management/refactor/resources/groups/types"
+	"github.com/netbirdio/netbird/management/refactor/resources/peers"
+	policyTypes "github.com/netbirdio/netbird/management/refactor/resources/policies/types"
+	routeTypes "github.com/netbirdio/netbird/management/refactor/resources/routes/types"
+	"github.com/netbirdio/netbird/management/refactor/resources/settings"
+	setupKeyTypes "github.com/netbirdio/netbird/management/refactor/resources/setup_keys/types"
+	userTypes "github.com/netbirdio/netbird/management/refactor/resources/users/types"
+
+	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/status"
 	"github.com/netbirdio/netbird/management/server/telemetry"
 )
 
@@ -23,37 +33,13 @@ const (
 
 // SqliteStore represents an account storage backed by a Sqlite DB persisted to disk
 type DefaultSqliteStore struct {
-	db                *gorm.DB
+	DB                *gorm.DB
 	storeFile         string
 	accountLocks      sync.Map
 	globalAccountLock sync.Mutex
 	metrics           telemetry.AppMetrics
 	installationPK    int
-}
-
-func (s *DefaultSqliteStore) FindSettings(accountID string) (settings.Settings, error) {
-	// TODO implement me
-	panic("implement me")
-}
-
-func (s *DefaultSqliteStore) FindPeerByPubKey(pubKey string) (peers.Peer, error) {
-	// TODO implement me
-	panic("implement me")
-}
-
-func (s *DefaultSqliteStore) FindPeerByID(id string) (peers.Peer, error) {
-	// TODO implement me
-	panic("implement me")
-}
-
-func (s *DefaultSqliteStore) FindAllPeersInAccount(id string) ([]peers.Peer, error) {
-	// TODO implement me
-	panic("implement me")
-}
-
-func (s *DefaultSqliteStore) UpdatePeer(peer peers.Peer) error {
-	// TODO implement me
-	panic("implement me")
+	accounts          map[string]*DefaultAccount
 }
 
 type installation struct {
@@ -63,10 +49,10 @@ type installation struct {
 
 // NewSqliteStore restores a store from the file located in the datadir
 func NewDefaultSqliteStore(dataDir string, metrics telemetry.AppMetrics) (*DefaultSqliteStore, error) {
-	storeStr := "store.db?cache=shared"
+	storeStr := "store.DB?cache=shared"
 	if runtime.GOOS == "windows" {
 		// Vo avoid `The process cannot access the file because it is being used by another process` on Windows
-		storeStr = "store.db"
+		storeStr = "store.DB"
 	}
 
 	file := filepath.Join(dataDir, storeStr)
@@ -85,7 +71,7 @@ func NewDefaultSqliteStore(dataDir string, metrics telemetry.AppMetrics) (*Defau
 	conns := runtime.NumCPU()
 	sql.SetMaxOpenConns(conns) // TODO: make it configurable
 
-	// err = db.AutoMigrate(
+	// err = DB.AutoMigrate(
 	// 	&SetupKey{}, &Peer{}, &User{}, &PersonalAccessToken{}, &Group{}, &Rule{},
 	// 	&Account{}, &Policy{}, &PolicyRule{}, &route.Route{}, &nbdns.NameServerGroup{},
 	// 	&installation{},
@@ -94,12 +80,7 @@ func NewDefaultSqliteStore(dataDir string, metrics telemetry.AppMetrics) (*Defau
 	// 	return nil, err
 	// }
 
-	return &DefaultSqliteStore{db: db, storeFile: file, metrics: metrics, installationPK: 1}, nil
-}
-
-func (s *DefaultSqliteStore) GetLicense() string {
-	// TODO implement me
-	panic("implement me")
+	return &DefaultSqliteStore{DB: db, storeFile: file, metrics: metrics, installationPK: 1}, nil
 }
 
 // AcquireGlobalLock acquires global lock across all the accounts and returns a function that releases the lock
@@ -138,17 +119,161 @@ func (s *DefaultSqliteStore) AcquireAccountLock(accountID string) (unlock func()
 	return unlock
 }
 
+func (s *DefaultSqliteStore) LoadAccount(accountID string) error {
+	var account DefaultAccount
+	result := s.DB.Model(&account).
+		Preload("UsersG.PATsG"). // have to be specifies as this is nester reference
+		Preload(clause.Associations).
+		First(&account, "id = ?", accountID)
+	if result.Error != nil {
+		log.Errorf("error when getting account from the store: %s", result.Error)
+		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
+			return status.Errorf(status.NotFound, "account not found")
+		}
+		return status.Errorf(status.Internal, "issue getting account from store")
+	}
+
+	// we have to manually preload policy rules as it seems that gorm preloading doesn't do it for us
+	for i, policy := range account.Policies {
+		var rules []*policyTypes.DefaultPolicyRule
+		err := s.DB.Model(&policyTypes.DefaultPolicyRule{}).Find(&rules, "policy_id = ?", policy.ID).Error
+		if err != nil {
+			return status.Errorf(status.NotFound, "rule not found")
+		}
+		account.Policies[i].Rules = rules
+	}
+
+	account.SetupKeys = make(map[string]*setupKeyTypes.DefaultSetupKey, len(account.SetupKeysG))
+	for _, key := range account.SetupKeysG {
+		account.SetupKeys[key.Key] = key.Copy()
+	}
+	account.SetupKeysG = nil
+
+	account.Peers = make(map[string]*nbpeer.Peer, len(account.PeersG))
+	for _, peer := range account.PeersG {
+		account.Peers[peer.ID] = peer.Copy()
+	}
+	account.PeersG = nil
+
+	account.Users = make(map[string]*userTypes.DefaultUser, len(account.UsersG))
+	for _, user := range account.UsersG {
+		user.PATs = make(map[string]*PersonalAccessToken, len(user.PATs))
+		for _, pat := range user.PATsG {
+			user.PATs[pat.ID] = pat.Copy()
+		}
+		account.Users[user.Id] = user.Copy()
+	}
+	account.UsersG = nil
+
+	account.Groups = make(map[string]*groupTypes.DefaultGroup, len(account.GroupsG))
+	for _, group := range account.GroupsG {
+		account.Groups[group.ID] = group.Copy()
+	}
+	account.GroupsG = nil
+
+	account.Routes = make(map[string]*routeTypes.DefaultRoute, len(account.RoutesG))
+	for _, route := range account.RoutesG {
+		account.Routes[route.ID] = route.Copy()
+	}
+	account.RoutesG = nil
+
+	account.NameServerGroups = make(map[string]*dnsTypes.DefaultNameServerGroup, len(account.NameServerGroupsG))
+	for _, ns := range account.NameServerGroupsG {
+		account.NameServerGroups[ns.ID] = ns.Copy()
+	}
+	account.NameServerGroupsG = nil
+
+	s.accounts[account.Id] = &account
+
+	return nil
+}
+
+func (s *DefaultSqliteStore) WriteAccount(accountID string) error {
+	start := time.Now()
+
+	account, ok := s.accounts[accountID]
+	if !ok {
+		return status.Errorf(status.NotFound, "account not found")
+	}
+
+	for _, key := range account.SetupKeys {
+		account.SetupKeysG = append(account.SetupKeysG, *key)
+	}
+
+	for id, peer := range account.Peers {
+		peer.ID = id
+		account.PeersG = append(account.PeersG, *peer)
+	}
+
+	for id, user := range account.Users {
+		user.Id = id
+		for id, pat := range user.PATs {
+			pat.ID = id
+			user.PATsG = append(user.PATsG, *pat)
+		}
+		account.UsersG = append(account.UsersG, *user)
+	}
+
+	for id, group := range account.Groups {
+		group.ID = id
+		account.GroupsG = append(account.GroupsG, *group)
+	}
+
+	for id, route := range account.Routes {
+		route.ID = id
+		account.RoutesG = append(account.RoutesG, *route)
+	}
+
+	for id, ns := range account.NameServerGroups {
+		ns.ID = id
+		account.NameServerGroupsG = append(account.NameServerGroupsG, *ns)
+	}
+
+	err := s.DB.Transaction(func(tx *gorm.DB) error {
+		result := tx.Select(clause.Associations).Delete(account.Policies, "account_id = ?", account.Id)
+		if result.Error != nil {
+			return result.Error
+		}
+
+		result = tx.Select(clause.Associations).Delete(account.UsersG, "account_id = ?", account.Id)
+		if result.Error != nil {
+			return result.Error
+		}
+
+		result = tx.Select(clause.Associations).Delete(account)
+		if result.Error != nil {
+			return result.Error
+		}
+
+		result = tx.
+			Session(&gorm.Session{FullSaveAssociations: true}).
+			Clauses(clause.OnConflict{UpdateAll: true}).Create(account)
+		if result.Error != nil {
+			return result.Error
+		}
+		return nil
+	})
+
+	took := time.Since(start)
+	if s.metrics != nil {
+		s.metrics.StoreMetrics().CountPersistenceDuration(took)
+	}
+	log.Debugf("took %d ms to persist an account to the SQLite", took.Milliseconds())
+
+	return err
+}
+
 func (s *DefaultSqliteStore) SaveInstallationID(ID string) error {
 	installation := installation{InstallationIDValue: ID}
 	installation.ID = uint(s.installationPK)
 
-	return s.db.Clauses(clause.OnConflict{UpdateAll: true}).Create(&installation).Error
+	return s.DB.Clauses(clause.OnConflict{UpdateAll: true}).Create(&installation).Error
 }
 
 func (s *DefaultSqliteStore) GetInstallationID() string {
 	var installation installation
 
-	if result := s.db.First(&installation, "id = ?", s.installationPK); result.Error != nil {
+	if result := s.DB.First(&installation, "id = ?", s.installationPK); result.Error != nil {
 		return ""
 	}
 
@@ -164,3 +289,57 @@ func (s *DefaultSqliteStore) Close() error {
 func (s *DefaultSqliteStore) GetStoreEngine() StoreEngine {
 	return SqliteStoreEngine
 }
+
+func (s *DefaultSqliteStore) GetLicense() string {
+	// TODO implement me
+	panic("implement me")
+}
+
+func (s *DefaultSqliteStore) FindSettings(accountID string) (settings.Settings, error) {
+	account, ok := s.accounts[accountID]
+	if !ok {
+		return nil, status.Errorf(status.NotFound, "account not found")
+	}
+	return account.Settings, nil
+}
+
+func (s *DefaultSqliteStore) FindPeerByPubKey(accountID string, pubKey string) (peers.Peer, error) {
+	a, ok := s.accounts[accountID]
+	if !ok {
+		return nil, status.Errorf(status.NotFound, "account not found")
+	}
+	for _, peer := range a.Peers {
+		if peer.Key == pubKey {
+			return peer.Copy(), nil
+		}
+	}
+
+	return nil, status.Errorf(status.NotFound, "peer with the public key %s not found", pubKey)
+}
+
+func (s *DefaultSqliteStore) FindPeerByID(accountID string, id string) (peers.Peer, error) {
+	a, ok := s.accounts[accountID]
+	if !ok {
+		return nil, status.Errorf(status.NotFound, "account not found")
+	}
+	for _, peer := range a.Peers {
+		if peer.ID == id {
+			return peer.Copy(), nil
+		}
+	}
+
+	return nil, status.Errorf(status.NotFound, "peer with the ID %s not found", id)
+}
+
+func (s *DefaultSqliteStore) FindAllPeersInAccount(accountId string) ([]peers.Peer, error) {
+	a, ok := s.accounts[accountID]
+	if !ok {
+		return nil, status.Errorf(status.NotFound, "account not found")
+	}
+	return a.Peers, nil
+}
+
+func (s *DefaultSqliteStore) UpdatePeer(peer peers.Peer) error {
+	// TODO implement me
+	panic("implement me")
+}