frostfs-node/pkg/util/locode/db/db.go

184 lines
4.5 KiB
Go
Raw Normal View History

package locodedb
import (
"errors"
"fmt"
"runtime"
"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/util/locode"
"golang.org/x/sync/errgroup"
)
// SourceTable is an interface of the UN/LOCODE table.
type SourceTable interface {
// Must iterate over all entries of the table
// and pass next entry to the handler.
//
// Must return handler's errors directly.
IterateAll(func(locode.Record) error) error
}
// DB is an interface of FrostFS location database.
type DB interface {
// Must save the record by key in the database.
Put(Key, Record) error
// Must return the record by key from the database.
Get(Key) (*Record, error)
}
// AirportRecord represents the entry in FrostFS airport database.
type AirportRecord struct {
// Name of the country where airport is located.
CountryName string
// Geo point where airport is located.
Point *Point
}
// ErrAirportNotFound is returned by AirportRecord readers
// when the required airport is not found.
var ErrAirportNotFound = errors.New("airport not found")
// AirportDB is an interface of FrostFS airport database.
type AirportDB interface {
// Must return the record by UN/LOCODE table record.
//
// Must return ErrAirportNotFound if there is no
// related airport in the database.
Get(locode.Record) (*AirportRecord, error)
}
// ContinentsDB is an interface of FrostFS continent database.
type ContinentsDB interface {
// Must return continent of the geo point.
PointContinent(*Point) (*Continent, error)
}
var ErrSubDivNotFound = errors.New("subdivision not found")
var ErrCountryNotFound = errors.New("country not found")
// NamesDB is an interface of the FrostFS location namespace.
type NamesDB interface {
// Must resolve a country code to a country name.
//
// Must return ErrCountryNotFound if there is no
// country with the provided code.
CountryName(*CountryCode) (string, error)
// Must resolve (country code, subdivision code) to
// a subdivision name.
//
// Must return ErrSubDivNotFound if either country or
// subdivision is not presented in database.
SubDivName(*CountryCode, string) (string, error)
}
// FillDatabase generates the FrostFS location database based on the UN/LOCODE table.
func FillDatabase(table SourceTable, airports AirportDB, continents ContinentsDB, names NamesDB, db DB) error {
var errG errgroup.Group
// Pick some sane default, after this the performance stopped increasing.
errG.SetLimit(runtime.NumCPU() * 4)
_ = table.IterateAll(func(tableRecord locode.Record) error {
errG.Go(func() error {
return processTableRecord(tableRecord, airports, continents, names, db)
})
return nil
})
return errG.Wait()
}
func processTableRecord(tableRecord locode.Record, airports AirportDB, continents ContinentsDB, names NamesDB, db DB) error {
if tableRecord.LOCODE.LocationCode() == "" {
return nil
}
dbKey, err := NewKey(tableRecord.LOCODE)
if err != nil {
return err
}
dbRecord, err := NewRecord(tableRecord)
if err != nil {
if errors.Is(err, errParseCoordinates) {
return nil
}
return err
}
geoPoint := dbRecord.GeoPoint()
countryName := ""
if geoPoint == nil {
airportRecord, err := airports.Get(tableRecord)
if err != nil {
if errors.Is(err, ErrAirportNotFound) {
return nil
}
return err
}
geoPoint = airportRecord.Point
countryName = airportRecord.CountryName
}
dbRecord.SetGeoPoint(geoPoint)
if countryName == "" {
countryName, err = names.CountryName(dbKey.CountryCode())
if err != nil {
if errors.Is(err, ErrCountryNotFound) {
return nil
}
return err
}
}
dbRecord.SetCountryName(countryName)
if subDivCode := dbRecord.SubDivCode(); subDivCode != "" {
subDivName, err := names.SubDivName(dbKey.CountryCode(), subDivCode)
if err != nil {
if errors.Is(err, ErrSubDivNotFound) {
return nil
}
return err
}
dbRecord.SetSubDivName(subDivName)
}
continent, err := continents.PointContinent(geoPoint)
if err != nil {
return fmt.Errorf("could not calculate continent geo point: %w", err)
} else if continent.Is(ContinentUnknown) {
return nil
}
dbRecord.SetContinent(continent)
return db.Put(*dbKey, *dbRecord)
}
// LocodeRecord returns the record from the FrostFS location database
// corresponding to the string representation of UN/LOCODE.
func LocodeRecord(db DB, sLocode string) (*Record, error) {
lc, err := locode.FromString(sLocode)
if err != nil {
return nil, fmt.Errorf("could not parse locode: %w", err)
}
key, err := NewKey(*lc)
if err != nil {
return nil, err
}
return db.Get(*key)
}