discordBots/cmd/domaincheckbot/repository/repository.go

package repository

import (
	"encoding/json"
	"github.com/dgraph-io/badger/v4"
	"log"
)

type Domain struct {
	Name string
	NS   []string
}

type Repository interface {
	GetAll() ([]Domain, error)
	GetValue(domainName string) Domain
	SetValue(domain Domain) error
	DeleteValue(domainName string) error
	Close() error
}

type DomainRepository struct {
	db *badger.DB
}

func InitDb() *DomainRepository {
	opts := badger.DefaultOptions("./badger")
	opts.Logger = nil
	db, err := badger.Open(opts)
	if err != nil {
		log.Fatal(err)
	}
	return &DomainRepository{db}
}

func (d *DomainRepository) Close() error {
	return d.db.Close()
}

func (d *DomainRepository) RunGC() error {
	return d.db.RunValueLogGC(0.7)
}

func (d *DomainRepository) GetAll() ([]Domain, error) {
	var domains []Domain
	err := d.db.View(func(txn *badger.Txn) error {
		opts := badger.DefaultIteratorOptions
		opts.PrefetchSize = 10
		it := txn.NewIterator(opts)
		defer it.Close()
		for it.Rewind(); it.Valid(); it.Next() {
			item := it.Item()
			err := item.Value(func(val []byte) error {
				retrievedDomain := Domain{}
				err := json.Unmarshal(val, &retrievedDomain)
				domains = append(domains, retrievedDomain)
				return err
			})
			if err != nil {
				return err
			}
		}
		return nil
	})
	return domains, err
}

func (d *DomainRepository) GetValue(domainName string) (Domain, error) {
	retrievedDomain := Domain{}
	err := d.db.View(func(txn *badger.Txn) error {
		item, err := txn.Get([]byte(domainName))
		if err != nil {
			return err
		}
		err = item.Value(func(val []byte) error {
			err = json.Unmarshal(val, &retrievedDomain)
			return err
		})
		return err
	})
	if err != nil {
		return Domain{}, err
	}
	return retrievedDomain, nil
}

func (d *DomainRepository) SetValue(domain Domain) error {
	jsonBytes, err := json.Marshal(domain)
	if err != nil {
		return err
	}
	err = d.db.Update(func(txn *badger.Txn) error {
		err := txn.Set([]byte(domain.Name), jsonBytes)
		return err
	})
	if err != nil {
		return err
	}
	return nil
}

func (d *DomainRepository) DeleteValue(domainName string) error {
	err := d.db.Update(func(txn *badger.Txn) error {
		err := txn.Delete([]byte(domainName))
		return err
	})
	return err
}