package locodedb import ( "errors" "fmt" "runtime" "sync/atomic" "git.frostfs.info/TrueCloudLab/frostfs-locode-db/pkg/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) } type FillDatabaseResult struct { AddedRecordCount int IgnoredRecordCount int } // FillDatabase generates the FrostFS location database based on the UN/LOCODE table. func FillDatabase(table SourceTable, airports AirportDB, continents ContinentsDB, names NamesDB, db DB) (FillDatabaseResult, error) { var errG errgroup.Group var added, ignored atomic.Int32 // Pick some sane default, after this the performance stopped increasing. errG.SetLimit(runtime.NumCPU() * 16) _ = table.IterateAll(func(tableRecord locode.Record) error { errG.Go(func() error { wasAdded, err := processTableRecord(tableRecord, airports, continents, names, db) if err != nil { return err } if wasAdded { added.Add(1) } else { ignored.Add(1) } return nil }) return nil }) return FillDatabaseResult{ AddedRecordCount: int(added.Load()), IgnoredRecordCount: int(ignored.Load()), }, errG.Wait() } func processTableRecord(tableRecord locode.Record, airports AirportDB, continents ContinentsDB, names NamesDB, db DB) (bool, error) { if tableRecord.LOCODE.LocationCode() == "" { return false, nil } dbKey, err := NewKey(tableRecord.LOCODE) if err != nil { return false, err } dbRecord, err := NewRecord(tableRecord) if err != nil { if errors.Is(err, errParseCoordinates) { return false, nil } return false, err } geoPoint := dbRecord.GeoPoint() countryName := "" if geoPoint == nil { airportRecord, err := airports.Get(tableRecord) if err != nil { if errors.Is(err, ErrAirportNotFound) { return false, nil } return false, 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 false, nil } return false, 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 false, nil } return false, err } dbRecord.SetSubDivName(subDivName) } continent, err := continents.PointContinent(geoPoint) if err != nil { return false, fmt.Errorf("could not calculate continent geo point: %w", err) } else if continent.Is(ContinentUnknown) { return false, nil } dbRecord.SetContinent(continent) return true, 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) }