Files
probo/cmd/geoloc-import/main.go
Émile Ré ad22fec81d Add IP-to-country geolocation service
Introduce a geoloc package that stores CIDR-to-country mappings in
PostgreSQL using the native cidr type with a GiST index for fast
containment lookups. Data comes from the ipverse/country-ip-blocks
dataset added as a git submodule.

A standalone geoloc-import command reads the TXT files from disk
and bulk-loads them via COPY. probod wires the service and logs a
warning when the table is empty.

Signed-off-by: Émile Ré <emile@getprobo.com>
2026-05-07 15:13:41 +04:00

105 lines
2.5 KiB
Go

// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
// Command geoloc-import loads IP-to-country CIDR blocks from the
// ipverse/country-ip-blocks dataset into the common_ip_country_blocks table.
package main
import (
"context"
"flag"
"fmt"
"net/url"
"os"
"go.gearno.de/kit/pg"
"go.probo.inc/probo/pkg/geoloc"
)
func main() {
if err := run(); err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(1)
}
}
func run() error {
var (
pgDSN string
dataDir string
)
flag.StringVar(
&pgDSN,
"pg-dsn",
os.Getenv("DATABASE_URL"),
"PostgreSQL connection URL (default: DATABASE_URL env)",
)
flag.StringVar(
&dataDir,
"data-dir",
"pkg/geoloc/data/country-ip-blocks",
"path to ipverse country-ip-blocks checkout",
)
flag.Parse()
if pgDSN == "" {
return fmt.Errorf("set -pg-dsn or DATABASE_URL")
}
ctx := context.Background()
pgClient, err := newPgClientFromDSN(pgDSN)
if err != nil {
return fmt.Errorf("cannot create pg client: %w", err)
}
svc := geoloc.NewService(pgClient)
fmt.Printf("importing IP country blocks from %s\n", dataDir)
if err := svc.ImportFromDir(ctx, dataDir); err != nil {
return fmt.Errorf("cannot import geoloc data: %w", err)
}
fmt.Println("done")
return nil
}
func newPgClientFromDSN(dsn string) (*pg.Client, error) {
u, err := url.Parse(dsn)
if err != nil {
return nil, fmt.Errorf("cannot parse DSN: %w", err)
}
var opts []pg.Option
if u.Host != "" {
opts = append(opts, pg.WithAddr(u.Host))
}
if u.User != nil {
opts = append(opts, pg.WithUser(u.User.Username()))
if password, ok := u.User.Password(); ok {
opts = append(opts, pg.WithPassword(password))
}
}
if len(u.Path) > 1 {
opts = append(opts, pg.WithDatabase(u.Path[1:]))
}
return pg.NewClient(opts...)
}