package locodedb import ( "errors" "fmt" "runtime" "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) } // 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) }