diff --git a/pkg/coredata/ip_country_block.go b/pkg/coredata/ip_country_block.go index 5512cda98..0c2d93c47 100644 --- a/pkg/coredata/ip_country_block.go +++ b/pkg/coredata/ip_country_block.go @@ -70,28 +70,67 @@ func IsIPCountryBlocksPopulated(ctx context.Context, conn pg.Querier) (bool, err return populated, nil } -func TruncateIPCountryBlocks(ctx context.Context, conn pg.Querier) error { - _, err := conn.Exec(ctx, "TRUNCATE common_ip_country_blocks") - if err != nil { - return fmt.Errorf("cannot truncate common_ip_country_blocks: %w", err) +const ipCountryBlocksStagingTable = "common_ip_country_blocks_staging" + +func CreateIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error { + q := ` +DROP TABLE IF EXISTS common_ip_country_blocks_staging; +CREATE TABLE common_ip_country_blocks_staging (LIKE common_ip_country_blocks INCLUDING DEFAULTS); +` + + if _, err := conn.Exec(ctx, q); err != nil { + return fmt.Errorf("cannot create staging table: %w", err) } + return nil } -func CopyIPCountryBlocks(ctx context.Context, conn pg.Querier, blocks []IPCountryBlock) error { - rows := make([][]any, len(blocks)) - for i, b := range blocks { - rows[i] = []any{b.CIDR, b.CountryCode.String()} - } - +func CopyIPCountryBlocksStaging(ctx context.Context, conn pg.Querier, blocks []IPCountryBlock) error { _, err := conn.CopyFrom( ctx, - pgx.Identifier{"common_ip_country_blocks"}, + pgx.Identifier{ipCountryBlocksStagingTable}, []string{"cidr", "country_code"}, - pgx.CopyFromRows(rows), + pgx.CopyFromSlice(len(blocks), func(i int) ([]any, error) { + return []any{blocks[i].CIDR, blocks[i].CountryCode.String()}, nil + }), ) if err != nil { - return fmt.Errorf("cannot copy rows into common_ip_country_blocks: %w", err) + return fmt.Errorf("cannot copy rows into staging table: %w", err) + } + + return nil +} + +func FinalizeIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error { + q := ` +CREATE INDEX idx_common_ip_country_blocks_staging_cidr + ON common_ip_country_blocks_staging USING gist (cidr inet_ops); +ANALYZE common_ip_country_blocks_staging; +` + + if _, err := conn.Exec(ctx, q); err != nil { + return fmt.Errorf("cannot finalize staging table: %w", err) + } + + return nil +} + +func SwapIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error { + q := ` +DROP TABLE common_ip_country_blocks; +ALTER TABLE common_ip_country_blocks_staging RENAME TO common_ip_country_blocks; +` + + if _, err := conn.Exec(ctx, q); err != nil { + return fmt.Errorf("cannot swap staging table: %w", err) + } + + return nil +} + +func DropIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error { + if _, err := conn.Exec(ctx, "DROP TABLE IF EXISTS common_ip_country_blocks_staging"); err != nil { + return fmt.Errorf("cannot drop staging table: %w", err) } return nil diff --git a/pkg/geoloc/service.go b/pkg/geoloc/service.go index 80458001a..cd4e4bbb7 100644 --- a/pkg/geoloc/service.go +++ b/pkg/geoloc/service.go @@ -61,7 +61,7 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error { cidrs, err := parseCIDRFile(path) if err != nil { - continue + return fmt.Errorf("cannot parse CIDR file %s: %w", path, err) } for _, cidr := range cidrs { @@ -73,20 +73,37 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error { } } - return s.pgClient.WithTx( - ctx, - func(ctx context.Context, tx pg.Tx) error { - if err := coredata.TruncateIPCountryBlocks(ctx, tx); err != nil { - return err - } + if err := s.pgClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + if err := coredata.CreateIPCountryBlocksStaging(ctx, conn); err != nil { + return fmt.Errorf("cannot create staging table: %w", err) + } - if err := coredata.CopyIPCountryBlocks(ctx, tx, blocks); err != nil { - return err - } + if err := coredata.CopyIPCountryBlocksStaging(ctx, conn, blocks); err != nil { + return fmt.Errorf("cannot copy IP country blocks to staging: %w", err) + } - return nil - }, - ) + if err := coredata.FinalizeIPCountryBlocksStaging(ctx, conn); err != nil { + return fmt.Errorf("cannot finalize staging table: %w", err) + } + + return nil + }); err != nil { + _ = s.pgClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + return coredata.DropIPCountryBlocksStaging(ctx, conn) + }) + return err + } + + if err := s.pgClient.WithTx(ctx, func(ctx context.Context, tx pg.Tx) error { + return coredata.SwapIPCountryBlocksStaging(ctx, tx) + }); err != nil { + _ = s.pgClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error { + return coredata.DropIPCountryBlocksStaging(ctx, conn) + }) + return fmt.Errorf("cannot swap staging table: %w", err) + } + + return nil } func (s *Service) LookupCountry(ctx context.Context, ip string) (coredata.CountryCode, error) { @@ -101,8 +118,11 @@ func (s *Service) LookupCountry(ctx context.Context, ip string) (coredata.Countr ctx, func(ctx context.Context, conn pg.Querier) error { var err error - cc, err = coredata.LookupCountryByIP(ctx, conn, ip) - return err + if cc, err = coredata.LookupCountryByIP(ctx, conn, ip); err != nil { + return fmt.Errorf("cannot lookup country by IP: %w", err) + } + + return nil }, ) if err != nil { @@ -119,8 +139,11 @@ func (s *Service) IsPopulated(ctx context.Context) (bool, error) { ctx, func(ctx context.Context, conn pg.Querier) error { var err error - populated, err = coredata.IsIPCountryBlocksPopulated(ctx, conn) - return err + if populated, err = coredata.IsIPCountryBlocksPopulated(ctx, conn); err != nil { + return fmt.Errorf("cannot check if IP country blocks are populated: %w", err) + } + + return nil }, ) if err != nil { @@ -148,7 +171,7 @@ func parseCIDRFile(path string) ([]string, error) { _, _, err := net.ParseCIDR(line) if err != nil { - continue + return nil, fmt.Errorf("cannot parse file: %w", err) } cidrs = append(cidrs, line)