diff --git a/locode_generate.go b/locode_generate.go index 7235675..ff7ff4f 100644 --- a/locode_generate.go +++ b/locode_generate.go @@ -69,8 +69,10 @@ var ( Table: locodeDB, } - err = locodedb.FillDatabase(locodeDB, airportDB, continentsDB, names, targetDB) + res, err := locodedb.FillDatabase(locodeDB, airportDB, continentsDB, names, targetDB) ExitOnErr(cmd, "", err) + cmd.Printf("Records added to locode db: %d\n", res.AddedRecordCount) + cmd.Printf("Records ignored: %d\n", res.IgnoredRecordCount) }, } ) diff --git a/pkg/locode/db/db.go b/pkg/locode/db/db.go index 90b19ed..82c02cb 100644 --- a/pkg/locode/db/db.go +++ b/pkg/locode/db/db.go @@ -4,6 +4,7 @@ import ( "errors" "fmt" "runtime" + "sync/atomic" "git.frostfs.info/TrueCloudLab/frostfs-locode-db/pkg/locode" "golang.org/x/sync/errgroup" @@ -75,38 +76,58 @@ type NamesDB interface { 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) error { +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 { - return processTableRecord(tableRecord, airports, continents, names, db) + 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 errG.Wait() + + 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) error { +func processTableRecord(tableRecord locode.Record, airports AirportDB, continents ContinentsDB, names NamesDB, db DB) (bool, error) { if tableRecord.LOCODE.LocationCode() == "" { - return nil + return false, nil } dbKey, err := NewKey(tableRecord.LOCODE) if err != nil { - return err + return false, err } dbRecord, err := NewRecord(tableRecord) if err != nil { if errors.Is(err, errParseCoordinates) { - return nil + return false, nil } - return err + return false, err } geoPoint := dbRecord.GeoPoint() @@ -116,10 +137,10 @@ func processTableRecord(tableRecord locode.Record, airports AirportDB, continent airportRecord, err := airports.Get(tableRecord) if err != nil { if errors.Is(err, ErrAirportNotFound) { - return nil + return false, nil } - return err + return false, err } geoPoint = airportRecord.Point @@ -132,10 +153,10 @@ func processTableRecord(tableRecord locode.Record, airports AirportDB, continent countryName, err = names.CountryName(dbKey.CountryCode()) if err != nil { if errors.Is(err, ErrCountryNotFound) { - return nil + return false, nil } - return err + return false, err } } @@ -145,10 +166,10 @@ func processTableRecord(tableRecord locode.Record, airports AirportDB, continent subDivName, err := names.SubDivName(dbKey.CountryCode(), subDivCode) if err != nil { if errors.Is(err, ErrSubDivNotFound) { - return nil + return false, nil } - return err + return false, err } dbRecord.SetSubDivName(subDivName) @@ -156,14 +177,14 @@ func processTableRecord(tableRecord locode.Record, airports AirportDB, continent continent, err := continents.PointContinent(geoPoint) if err != nil { - return fmt.Errorf("could not calculate continent geo point: %w", err) + return false, fmt.Errorf("could not calculate continent geo point: %w", err) } else if continent.Is(ContinentUnknown) { - return nil + return false, nil } dbRecord.SetContinent(continent) - return db.Put(*dbKey, *dbRecord) + return true, db.Put(*dbKey, *dbRecord) } // LocodeRecord returns the record from the FrostFS location database