frostfs-locode-db/pkg/locode/db/boltdb/calls.go

package locodebolt

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"

	locodedb "git.frostfs.info/TrueCloudLab/frostfs-locode-db/pkg/locode/db"
	"go.etcd.io/bbolt"
)

// Open opens an underlying BoltDB instance.
//
// Timeout of BoltDB opening is 3s (only for Linux or Darwin).
//
// Opens BoltDB in read-only mode if DB is read-only.
func (db *DB) Open() error {
	// copy-paste from metabase:
	// consider universal Open/Close for BoltDB wrappers

	err := os.MkdirAll(filepath.Dir(db.path), db.mode|0o110)
	if err != nil {
		return fmt.Errorf("could not create dir for BoltDB: %w", err)
	}

	db.bolt, err = bbolt.Open(db.path, db.mode, db.boltOpts)
	if err != nil {
		return fmt.Errorf("could not open BoltDB: %w", err)
	}

	return nil
}

// Close closes an underlying BoltDB instance.
//
// Must not be called before successful Open call.
func (db *DB) Close() error {
	return db.bolt.Close()
}

func countryBucketKey(cc *locodedb.CountryCode) ([]byte, error) {
	return []byte(cc.String()), nil
}

func locationBucketKey(lc *locodedb.LocationCode) ([]byte, error) {
	return []byte(lc.String()), nil
}

type recordJSON struct {
	CountryName  string
	LocationName string
	SubDivName   string
	SubDivCode   string
	Latitude     float64
	Longitude    float64
	Continent    string
}

func recordValue(r locodedb.Record) ([]byte, error) {
	p := r.GeoPoint()

	rj := &recordJSON{
		CountryName:  r.CountryName(),
		LocationName: r.LocationName(),
		SubDivName:   r.SubDivName(),
		SubDivCode:   r.SubDivCode(),
		Latitude:     p.Latitude(),
		Longitude:    p.Longitude(),
		Continent:    r.Continent().String(),
	}

	return json.Marshal(rj)
}

func recordFromValue(data []byte) (*locodedb.Record, error) {
	rj := new(recordJSON)

	if err := json.Unmarshal(data, rj); err != nil {
		return nil, err
	}

	r := new(locodedb.Record)
	r.SetCountryName(rj.CountryName)
	r.SetLocationName(rj.LocationName)
	r.SetSubDivName(rj.SubDivName)
	r.SetSubDivCode(rj.SubDivCode)
	r.SetGeoPoint(locodedb.NewPoint(rj.Latitude, rj.Longitude))

	cont := locodedb.ContinentFromString(rj.Continent)
	r.SetContinent(&cont)

	return r, nil
}

// Put saves the record by key in an underlying BoltDB instance.
//
// Country code from the key is used for allocating the 1st level buckets.
// Records are stored in country buckets by the location code from the key.
// The records are stored in internal binary JSON format.
//
// Must not be called before successful Open call.
// Must not be called in read-only mode: behavior is undefined.
func (db *DB) Put(key locodedb.Key, rec locodedb.Record) error {
	return db.bolt.Batch(func(tx *bbolt.Tx) error {
		countryKey, err := countryBucketKey(key.CountryCode())
		if err != nil {
			return err
		}

		bktCountry, err := tx.CreateBucketIfNotExists(countryKey)
		if err != nil {
			return fmt.Errorf("could not create country bucket: %w", err)
		}

		locationKey, err := locationBucketKey(key.LocationCode())
		if err != nil {
			return err
		}

		cont, err := recordValue(rec)
		if err != nil {
			return err
		}

		return bktCountry.Put(locationKey, cont)
	})
}

var errRecordNotFound = errors.New("record not found")

// Get reads the record by key from underlying BoltDB instance.
//
// Returns an error if no record is presented by key in DB.
//
// Must not be called before successful Open call.
func (db *DB) Get(key locodedb.Key) (rec *locodedb.Record, err error) {
	err = db.bolt.View(func(tx *bbolt.Tx) error {
		countryKey, err := countryBucketKey(key.CountryCode())
		if err != nil {
			return err
		}

		bktCountry := tx.Bucket(countryKey)
		if bktCountry == nil {
			return errRecordNotFound
		}

		locationKey, err := locationBucketKey(key.LocationCode())
		if err != nil {
			return err
		}

		data := bktCountry.Get(locationKey)
		if data == nil {
			return errRecordNotFound
		}

		rec, err = recordFromValue(data)

		return err
	})

	return
}