Badger store: avoid code duplication by using generic methods (#369)

* Badger store: avoid code duplication by using generic methods

* No need to return PT, can return just *T
This commit is contained in:
Nikolay Edigaryev 2025-11-17 18:34:59 +04:00 committed by GitHub
parent dbb180befb
commit 9cdfd75f79
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 146 additions and 287 deletions

View File

@ -3,14 +3,15 @@ package controller
import (
"crypto/tls"
"fmt"
"os"
"strings"
"github.com/cirruslabs/orchard/internal/certificatefingerprint"
"github.com/cirruslabs/orchard/internal/controller"
v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/pterm/pterm"
"github.com/samber/lo"
"github.com/sethvargo/go-password/password"
"os"
"strings"
)
const BootstrapContextName = "bootstrap-context"
@ -33,7 +34,7 @@ func Bootstrap(controllerInstance *controller.Controller, controllerCert tls.Cer
// However, if the BootstrapAdminName service account still exists,
// return its credentials. We'll use them for updating the
// BootstrapContextName context.
if serviceAccount, ok := lo.Find(serviceAccounts, func(serviceAccount *v1.ServiceAccount) bool {
if serviceAccount, ok := lo.Find(serviceAccounts, func(serviceAccount v1.ServiceAccount) bool {
return serviceAccount.Name == BootstrapAdminName
}); ok {
return serviceAccount.Name, serviceAccount.Token, nil

View File

@ -196,8 +196,8 @@ func New(opts ...Option) (*Controller, error) {
return controller, nil
}
func (controller *Controller) ServiceAccounts() ([]*v1.ServiceAccount, error) {
var serviceAccounts []*v1.ServiceAccount
func (controller *Controller) ServiceAccounts() ([]v1.ServiceAccount, error) {
var serviceAccounts []v1.ServiceAccount
var err error
if err := controller.store.View(func(txn storepkg.Transaction) error {

View File

@ -1,46 +1,15 @@
package badger
import (
"encoding/json"
"github.com/cirruslabs/orchard/pkg/resource/v1"
)
var ClusterSettingsKey = []byte("/cluster-settings")
func (txn *Transaction) GetClusterSettings() (_ *v1.ClusterSettings, err error) {
defer func() {
err = mapErr(err)
}()
item, err := txn.badgerTxn.Get(ClusterSettingsKey)
if err != nil {
return nil, err
}
valueBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var clusterSettings v1.ClusterSettings
err = json.Unmarshal(valueBytes, &clusterSettings)
if err != nil {
return nil, err
}
return &clusterSettings, nil
func (txn *Transaction) GetClusterSettings() (*v1.ClusterSettings, error) {
return genericGet[v1.ClusterSettings](txn, ClusterSettingsKey)
}
func (txn *Transaction) SetClusterSettings(clusterSettings v1.ClusterSettings) (err error) {
defer func() {
err = mapErr(err)
}()
valueBytes, err := json.Marshal(clusterSettings)
if err != nil {
return err
}
return txn.badgerTxn.Set(ClusterSettingsKey, valueBytes)
func (txn *Transaction) SetClusterSettings(clusterSettings v1.ClusterSettings) error {
return genericSet[v1.ClusterSettings](txn, ClusterSettingsKey, clusterSettings)
}

View File

@ -0,0 +1,97 @@
package badger
import (
"encoding/json"
"github.com/dgraph-io/badger/v3"
)
func genericSet[T any](txn *Transaction, key []byte, obj T) (err error) {
defer func() {
err = mapErr(err)
}()
valueBytes, err := json.Marshal(obj)
if err != nil {
return err
}
return txn.badgerTxn.Set(key, valueBytes)
}
func genericGet[T any, PT interface {
SetVersion(uint64)
*T
}](txn *Transaction, key []byte) (_ *T, err error) {
defer func() {
err = mapErr(err)
}()
item, err := txn.badgerTxn.Get(key)
if err != nil {
return nil, err
}
valueBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var obj T
err = json.Unmarshal(valueBytes, &obj)
if err != nil {
return nil, err
}
PT(&obj).SetVersion(item.Version())
return &obj, nil
}
func genericList[T any, PT interface {
SetVersion(uint64)
*T
}](txn *Transaction, prefix []byte) (_ []T, err error) {
defer func() {
err = mapErr(err)
}()
// Declare an empty, non-nil slice to
// return [] when no objects are found
result := []T{}
it := txn.badgerTxn.NewIterator(badger.IteratorOptions{
Prefix: prefix,
})
defer it.Close()
for it.Rewind(); it.Valid(); it.Next() {
item := it.Item()
valueBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var obj T
if err := json.Unmarshal(valueBytes, &obj); err != nil {
return nil, err
}
PT(&obj).SetVersion(item.Version())
result = append(result, obj)
}
return result, nil
}
func genericDelete(txn *Transaction, key []byte) (err error) {
defer func() {
err = mapErr(err)
}()
return txn.badgerTxn.Delete(key)
}

View File

@ -1,10 +1,9 @@
package badger
import (
"encoding/json"
"github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/dgraph-io/badger/v3"
"path"
"github.com/cirruslabs/orchard/pkg/resource/v1"
)
const SpaceServiceAccounts = "/service-accounts"
@ -13,88 +12,18 @@ func ServiceAccountKey(name string) []byte {
return []byte(path.Join(SpaceServiceAccounts, name))
}
func (txn *Transaction) GetServiceAccount(name string) (_ *v1.ServiceAccount, err error) {
defer func() {
err = mapErr(err)
}()
key := ServiceAccountKey(name)
item, err := txn.badgerTxn.Get(key)
if err != nil {
return nil, err
}
valueBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var serviceAccount v1.ServiceAccount
err = json.Unmarshal(valueBytes, &serviceAccount)
if err != nil {
return nil, err
}
return &serviceAccount, nil
func (txn *Transaction) GetServiceAccount(name string) (*v1.ServiceAccount, error) {
return genericGet[v1.ServiceAccount](txn, ServiceAccountKey(name))
}
func (txn *Transaction) SetServiceAccount(serviceAccount *v1.ServiceAccount) (err error) {
defer func() {
err = mapErr(err)
}()
key := ServiceAccountKey(serviceAccount.Name)
valueBytes, err := json.Marshal(serviceAccount)
if err != nil {
return err
}
return txn.badgerTxn.Set(key, valueBytes)
func (txn *Transaction) SetServiceAccount(serviceAccount *v1.ServiceAccount) error {
return genericSet[v1.ServiceAccount](txn, ServiceAccountKey(serviceAccount.Name), *serviceAccount)
}
func (txn *Transaction) DeleteServiceAccount(name string) (err error) {
defer func() {
err = mapErr(err)
}()
key := ServiceAccountKey(name)
return txn.badgerTxn.Delete(key)
func (txn *Transaction) DeleteServiceAccount(name string) error {
return genericDelete(txn, ServiceAccountKey(name))
}
func (txn *Transaction) ListServiceAccounts() (_ []*v1.ServiceAccount, err error) {
defer func() {
err = mapErr(err)
}()
// Declare an empty, non-nil slice to return
// [] when no service accounts are found
result := []*v1.ServiceAccount{}
it := txn.badgerTxn.NewIterator(badger.IteratorOptions{
Prefix: []byte(SpaceServiceAccounts),
})
defer it.Close()
for it.Rewind(); it.Valid(); it.Next() {
item := it.Item()
serviceAccountBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var serviceAccount v1.ServiceAccount
if err := json.Unmarshal(serviceAccountBytes, &serviceAccount); err != nil {
return nil, err
}
result = append(result, &serviceAccount)
}
return result, nil
func (txn *Transaction) ListServiceAccounts() ([]v1.ServiceAccount, error) {
return genericList[v1.ServiceAccount](txn, []byte(SpaceServiceAccounts))
}

View File

@ -2,11 +2,9 @@
package badger
import (
"encoding/json"
"path"
"github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/dgraph-io/badger/v3"
)
const SpaceVMs = "/vms"
@ -15,92 +13,18 @@ func VMKey(name string) []byte {
return []byte(path.Join(SpaceVMs, name))
}
func (txn *Transaction) GetVM(name string) (_ *v1.VM, err error) {
defer func() {
err = mapErr(err)
}()
key := VMKey(name)
item, err := txn.badgerTxn.Get(key)
if err != nil {
return nil, err
}
valueBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var vm v1.VM
err = json.Unmarshal(valueBytes, &vm)
if err != nil {
return nil, err
}
vm.Version = item.Version()
return &vm, nil
func (txn *Transaction) GetVM(name string) (*v1.VM, error) {
return genericGet[v1.VM](txn, VMKey(name))
}
func (txn *Transaction) SetVM(vm v1.VM) (err error) {
defer func() {
err = mapErr(err)
}()
key := VMKey(vm.Name)
valueBytes, err := json.Marshal(vm)
if err != nil {
return err
}
return txn.badgerTxn.Set(key, valueBytes)
func (txn *Transaction) SetVM(vm v1.VM) error {
return genericSet[v1.VM](txn, VMKey(vm.Name), vm)
}
func (txn *Transaction) DeleteVM(name string) (err error) {
defer func() {
err = mapErr(err)
}()
key := VMKey(name)
return txn.badgerTxn.Delete(key)
func (txn *Transaction) DeleteVM(name string) error {
return genericDelete(txn, VMKey(name))
}
func (txn *Transaction) ListVMs() (_ []v1.VM, err error) {
defer func() {
err = mapErr(err)
}()
// Declare an empty, non-nil slice to
// return [] when no VMs are found
result := []v1.VM{}
it := txn.badgerTxn.NewIterator(badger.IteratorOptions{
Prefix: []byte(SpaceVMs),
})
defer it.Close()
for it.Rewind(); it.Valid(); it.Next() {
item := it.Item()
vmBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var vm v1.VM
if err := json.Unmarshal(vmBytes, &vm); err != nil {
return nil, err
}
vm.Version = item.Version()
result = append(result, vm)
}
return result, nil
func (txn *Transaction) ListVMs() ([]v1.VM, error) {
return genericList[v1.VM](txn, []byte(SpaceVMs))
}

View File

@ -2,10 +2,9 @@
package badger
import (
"encoding/json"
"github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/dgraph-io/badger/v3"
"path"
"github.com/cirruslabs/orchard/pkg/resource/v1"
)
const SpaceWorkers = "/workers"
@ -14,88 +13,18 @@ func WorkerKey(name string) []byte {
return []byte(path.Join(SpaceWorkers, name))
}
func (txn *Transaction) GetWorker(name string) (_ *v1.Worker, err error) {
defer func() {
err = mapErr(err)
}()
key := WorkerKey(name)
item, err := txn.badgerTxn.Get(key)
if err != nil {
return nil, err
}
valueBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var worker v1.Worker
err = json.Unmarshal(valueBytes, &worker)
if err != nil {
return nil, err
}
return &worker, nil
func (txn *Transaction) GetWorker(name string) (*v1.Worker, error) {
return genericGet[v1.Worker](txn, WorkerKey(name))
}
func (txn *Transaction) SetWorker(worker v1.Worker) (err error) {
defer func() {
err = mapErr(err)
}()
key := WorkerKey(worker.Name)
valueBytes, err := json.Marshal(worker)
if err != nil {
return err
}
return txn.badgerTxn.Set(key, valueBytes)
func (txn *Transaction) SetWorker(worker v1.Worker) error {
return genericSet[v1.Worker](txn, WorkerKey(worker.Name), worker)
}
func (txn *Transaction) DeleteWorker(name string) (err error) {
defer func() {
err = mapErr(err)
}()
key := WorkerKey(name)
return txn.badgerTxn.Delete(key)
func (txn *Transaction) DeleteWorker(name string) error {
return genericDelete(txn, WorkerKey(name))
}
func (txn *Transaction) ListWorkers() (_ []v1.Worker, err error) {
defer func() {
err = mapErr(err)
}()
// Declare an empty, non-nil slice to
// return [] when no workers are found
result := []v1.Worker{}
it := txn.badgerTxn.NewIterator(badger.IteratorOptions{
Prefix: []byte(SpaceWorkers),
})
defer it.Close()
for it.Rewind(); it.Valid(); it.Next() {
item := it.Item()
vmBytes, err := item.ValueCopy(nil)
if err != nil {
return nil, err
}
var worker v1.Worker
if err := json.Unmarshal(vmBytes, &worker); err != nil {
return nil, err
}
result = append(result, worker)
}
return result, nil
func (txn *Transaction) ListWorkers() ([]v1.Worker, error) {
return genericList[v1.Worker](txn, []byte(SpaceWorkers))
}

View File

@ -39,7 +39,7 @@ type Transaction interface {
GetServiceAccount(name string) (result *v1.ServiceAccount, err error)
SetServiceAccount(serviceAccount *v1.ServiceAccount) (err error)
DeleteServiceAccount(name string) (err error)
ListServiceAccounts() (result []*v1.ServiceAccount, err error)
ListServiceAccounts() (result []v1.ServiceAccount, err error)
AppendEvents(event []v1.Event, scope ...string) (err error)
ListEvents(scope ...string) (result []v1.Event, err error)

View File

@ -14,6 +14,8 @@ type ClusterSettings struct {
SchedulerProfile SchedulerProfile `json:"schedulerProfile,omitempty"`
}
func (clusterSettings *ClusterSettings) SetVersion(_ uint64) {}
func NewSchedulerProfile(value string) (SchedulerProfile, error) {
switch value {
case string(SchedulerProfileOptimizeUtilization):

View File

@ -6,3 +6,5 @@ type ServiceAccount struct {
Meta
}
func (serviceAccount *ServiceAccount) SetVersion(_ uint64) {}

View File

@ -97,6 +97,10 @@ type VM struct {
Meta
}
func (vm *VM) SetVersion(version uint64) {
vm.Version = version
}
type VMSpec struct {
NetSoftnetDeprecated bool `json:"net-softnet,omitempty"`
NetSoftnet bool `json:"netSoftnet,omitempty"`

View File

@ -30,3 +30,5 @@ type Worker struct {
func (worker Worker) Offline(workerOfflineTimeout time.Duration) bool {
return time.Since(worker.LastSeen) > workerOfflineTimeout
}
func (worker *Worker) SetVersion(_ uint64) {}