From f2113a9c021a59f115c9b64c176d5d909a6d7361 Mon Sep 17 00:00:00 2001
From: Airat Arifullin <a.arifullin@yadro.com>
Date: Wed, 22 Jan 2025 22:03:36 +0300
Subject: [PATCH] [#21] locode: Introduce locode overriding

Signed-off-by: Airat Arifullin <a.arifullin@yadro.com>
---
 Makefile                      |  1 +
 locode_generate.go            | 10 +++++---
 pkg/locode/table/csv/calls.go | 47 +++++++++++++++++++++++++++++++++++
 pkg/locode/table/csv/table.go | 17 ++++++++++++-
 4 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/Makefile b/Makefile
index 0e6b28a..27428ad 100755
--- a/Makefile
+++ b/Makefile
@@ -51,6 +51,7 @@ locode_db: in/unlocode-CodeList.csv in/unlocode-SubdivisionCodes.csv in/continen
 	--continents in/continents.geojson \
 	--countries in/countries.dat \
 	--in in/unlocode-CodeList.csv \
+	--overrides in/unlocode-CodeList_overrides.csv \
 	--subdiv in/unlocode-SubdivisionCodes.csv \
 	--out locode_db
 	chmod 644 locode_db
diff --git a/locode_generate.go b/locode_generate.go
index 57ec4f8..7235675 100644
--- a/locode_generate.go
+++ b/locode_generate.go
@@ -16,6 +16,7 @@ type namesDB struct {
 
 const (
 	locodeGenerateInputFlag      = "in"
+	locodeGenerateOverridesFlag  = "overrides"
 	locodeGenerateSubDivFlag     = "subdiv"
 	locodeGenerateAirportsFlag   = "airports"
 	locodeGenerateCountriesFlag  = "countries"
@@ -25,6 +26,7 @@ const (
 
 var (
 	locodeGenerateInPaths        []string
+	locodeGenerateOverridesPath  string
 	locodeGenerateSubDivPath     string
 	locodeGenerateAirportsPath   string
 	locodeGenerateCountriesPath  string
@@ -35,11 +37,11 @@ var (
 		Use:   "generate",
 		Short: "Generate UN/LOCODE database for FrostFS",
 		Run: func(cmd *cobra.Command, _ []string) {
-
 			locodeDB := csvlocode.New(
 				csvlocode.Prm{
-					Path:       locodeGenerateInPaths[0],
-					SubDivPath: locodeGenerateSubDivPath,
+					Path:          locodeGenerateInPaths[0],
+					SubDivPath:    locodeGenerateSubDivPath,
+					OverridesPath: locodeGenerateOverridesPath,
 				},
 				csvlocode.WithExtraPaths(locodeGenerateInPaths[1:]...),
 			)
@@ -79,6 +81,8 @@ func initUtilLocodeGenerateCmd() {
 	flags.StringSliceVar(&locodeGenerateInPaths, locodeGenerateInputFlag, nil, "List of paths to UN/LOCODE tables (csv)")
 	_ = locodeGenerateCmd.MarkFlagRequired(locodeGenerateInputFlag)
 
+	flags.StringVar(&locodeGenerateOverridesPath, locodeGenerateOverridesFlag, "", "Path to UN/LOCODE override tables (csv)")
+
 	flags.StringVar(&locodeGenerateSubDivPath, locodeGenerateSubDivFlag, "", "Path to UN/LOCODE subdivision database (csv)")
 	_ = locodeGenerateCmd.MarkFlagRequired(locodeGenerateSubDivFlag)
 
diff --git a/pkg/locode/table/csv/calls.go b/pkg/locode/table/csv/calls.go
index 5ce0e7b..7d4e816 100644
--- a/pkg/locode/table/csv/calls.go
+++ b/pkg/locode/table/csv/calls.go
@@ -3,6 +3,7 @@ package csvlocode
 import (
 	"encoding/csv"
 	"errors"
+	"fmt"
 	"io"
 	"os"
 	"strings"
@@ -40,6 +41,10 @@ func (t *Table) IterateAll(f func(locode.Record) error) error {
 			Remarks:          words[11],
 		}
 
+		if err := t.Override(&record); err != nil {
+			return fmt.Errorf("override: %w", err)
+		}
+
 		return f(record)
 	})
 }
@@ -84,6 +89,14 @@ func (t *Table) SubDivName(countryCode *locodedb.CountryCode, code string) (stri
 	return rec.name, nil
 }
 
+func (t *Table) Override(record *locode.Record) error {
+	override, found := t.overrides[record.LOCODE]
+	if found {
+		*record = override
+	}
+	return nil
+}
+
 func (t *Table) initSubDiv() (err error) {
 	t.subDivOnce.Do(func() {
 		t.mSubDiv = make(map[subDivKey]subDivRecord)
@@ -103,6 +116,40 @@ func (t *Table) initSubDiv() (err error) {
 	return
 }
 
+func (t *Table) initOverrides(overridesPath string) error {
+	const wordsPerRecord = 12
+
+	t.overrides = make(map[locode.LOCODE]locode.Record)
+
+	if overridesPath == "" {
+		return nil
+	}
+
+	return t.scanWords([]string{overridesPath}, wordsPerRecord, func(words []string) error {
+		lc, err := locode.FromString(strings.Join(words[1:3], " "))
+		if err != nil {
+			return err
+		}
+
+		record := locode.Record{
+			Ch:               words[0],
+			LOCODE:           *lc,
+			Name:             words[3],
+			NameWoDiacritics: words[4],
+			SubDiv:           words[5],
+			Function:         words[6],
+			Status:           words[7],
+			Date:             words[8],
+			IATA:             words[9],
+			Coordinates:      words[10],
+			Remarks:          words[11],
+		}
+
+		t.overrides[record.LOCODE] = record
+		return nil
+	})
+}
+
 var errScanInt = errors.New("interrupt scan")
 
 func (t *Table) scanWords(paths []string, fpr int, wordsHandler func([]string) error) error {
diff --git a/pkg/locode/table/csv/table.go b/pkg/locode/table/csv/table.go
index b84c2b7..53299df 100644
--- a/pkg/locode/table/csv/table.go
+++ b/pkg/locode/table/csv/table.go
@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"io/fs"
 	"sync"
+
+	"git.frostfs.info/TrueCloudLab/frostfs-locode-db/pkg/locode"
 )
 
 // Prm groups the required parameters of the Table's constructor.
@@ -21,6 +23,11 @@ type Prm struct {
 	//
 	// Must not be empty.
 	SubDivPath string
+
+	// Path to a csv table with UN/LOCODE overrides.
+	//
+	// Optional.
+	OverridesPath string
 }
 
 // Table is a descriptor of the UN/LOCODE table in csv format.
@@ -39,6 +46,8 @@ type Table struct {
 	subDivOnce sync.Once
 
 	mSubDiv map[subDivKey]subDivRecord
+
+	overrides map[locode.LOCODE]locode.Record
 }
 
 const invalidPrmValFmt = "invalid parameter %s (%T):%v"
@@ -67,9 +76,15 @@ func New(prm Prm, opts ...Option) *Table {
 		opts[i](o)
 	}
 
-	return &Table{
+	t := &Table{
 		paths:      append(o.extraPaths, prm.Path),
 		mode:       o.mode,
 		subDivPath: prm.SubDivPath,
 	}
+
+	if err := t.initOverrides(prm.OverridesPath); err != nil {
+		panic(fmt.Errorf("init overrides: %w", err))
+	}
+
+	return t
 }