Move IP country block SQL to coredata

Extract struct and query functions into
pkg/coredata/ip_country_block.go following the convention that
all raw SQL lives in coredata. The geoloc service now delegates
to these functions.

Signed-off-by: Émile Ré <emile@getprobo.com>

pkg/coredata/ip_country_block.go

@@ -0,0 +1,98 @@
+// 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.
+
+package coredata
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5"
+	"go.gearno.de/kit/pg"
+)
+
+type IPCountryBlock struct {
+	CIDR        string      `db:"cidr"`
+	CountryCode CountryCode `db:"country_code"`
+}
+
+func LookupCountryByIP(ctx context.Context, conn pg.Querier, ip string) (CountryCode, error) {
+	q := `
+SELECT country_code
+FROM common_ip_country_blocks
+WHERE cidr >>= @ip::inet
+ORDER BY masklen(cidr) DESC
+LIMIT 1;
+`
+
+	args := pgx.StrictNamedArgs{"ip": ip}
+
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return "", fmt.Errorf("cannot query ip country blocks: %w", err)
+	}
+
+	cc, err := pgx.CollectExactlyOneRow(rows, pgx.RowTo[CountryCode])
+	if err != nil {
+		if err == pgx.ErrNoRows {
+			return "", nil
+		}
+		return "", fmt.Errorf("cannot collect ip country block row: %w", err)
+	}
+
+	return cc, nil
+}
+
+func IsIPCountryBlocksPopulated(ctx context.Context, conn pg.Querier) (bool, error) {
+	q := `SELECT EXISTS (SELECT 1 FROM common_ip_country_blocks);`
+
+	rows, err := conn.Query(ctx, q)
+	if err != nil {
+		return false, fmt.Errorf("cannot check if ip country blocks is populated: %w", err)
+	}
+
+	populated, err := pgx.CollectExactlyOneRow(rows, pgx.RowTo[bool])
+	if err != nil {
+		return false, fmt.Errorf("cannot collect populated check: %w", 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)
+	}
+	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()}
+	}
+
+	_, err := conn.CopyFrom(
+		ctx,
+		pgx.Identifier{"common_ip_country_blocks"},
+		[]string{"cidr", "country_code"},
+		pgx.CopyFromRows(rows),
+	)
+	if err != nil {
+		return fmt.Errorf("cannot copy rows into common_ip_country_blocks: %w", err)
+	}
+
+	return nil
+}

pkg/geoloc/service.go

@@ -23,7 +23,6 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/jackc/pgx/v5"
 	"go.gearno.de/kit/pg"
 	"go.probo.inc/probo/pkg/coredata"
 )
@@ -44,12 +43,7 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error {
 		return fmt.Errorf("cannot read country directory: %w", err)
 	}
 
-	type row struct {
+	var blocks []coredata.IPCountryBlock
-		cidr        string
-		countryCode coredata.CountryCode
-	}
-
-	var rows []row
 
 	for _, entry := range entries {
 		if !entry.IsDir() {
@@ -71,7 +65,10 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error {
 			}
 
 			for _, cidr := range cidrs {
-				rows = append(rows, row{cidr: cidr, countryCode: cc})
+				blocks = append(blocks, coredata.IPCountryBlock{
+					CIDR:        cidr,
+					CountryCode: cc,
+				})
 			}
 		}
 	}
@@ -79,24 +76,12 @@ func (s *Service) ImportFromDir(ctx context.Context, dataDir string) error {
 	return s.pgClient.WithTx(
 		ctx,
 		func(ctx context.Context, tx pg.Tx) error {
-			_, err := tx.Exec(ctx, "TRUNCATE common_ip_country_blocks")
+			if err := coredata.TruncateIPCountryBlocks(ctx, tx); err != nil {
-			if err != nil {
+				return err
-				return fmt.Errorf("cannot truncate common_ip_country_blocks: %w", err)
 			}
 
-			pgxRows := make([][]any, len(rows))
+			if err := coredata.CopyIPCountryBlocks(ctx, tx, blocks); err != nil {
-			for i, r := range rows {
+				return err
-				pgxRows[i] = []any{r.cidr, r.countryCode.String()}
-			}
-
-			_, err = tx.CopyFrom(
-				ctx,
-				pgx.Identifier{"common_ip_country_blocks"},
-				[]string{"cidr", "country_code"},
-				pgx.CopyFromRows(pgxRows),
-			)
-			if err != nil {
-				return fmt.Errorf("cannot copy rows into common_ip_country_blocks: %w", err)
 			}
 
 			return nil
@@ -110,46 +95,11 @@ func (s *Service) LookupCountry(ctx context.Context, conn pg.Querier, ip string)
 		return "", fmt.Errorf("cannot parse IP address: %q", ip)
 	}
 
-	q := `
+	return coredata.LookupCountryByIP(ctx, conn, ip)
-SELECT country_code
-FROM common_ip_country_blocks
-WHERE cidr >>= @ip::inet
-ORDER BY masklen(cidr) DESC
-LIMIT 1;
-`
-
-	args := pgx.StrictNamedArgs{"ip": ip}
-
-	rows, err := conn.Query(ctx, q, args)
-	if err != nil {
-		return "", fmt.Errorf("cannot query ip country blocks: %w", err)
-	}
-
-	cc, err := pgx.CollectExactlyOneRow(rows, pgx.RowTo[coredata.CountryCode])
-	if err != nil {
-		if err == pgx.ErrNoRows {
-			return "", nil
-		}
-		return "", fmt.Errorf("cannot collect ip country block row: %w", err)
-	}
-
-	return cc, nil
 }
 
 func (s *Service) IsPopulated(ctx context.Context, conn pg.Querier) (bool, error) {
-	q := `SELECT EXISTS (SELECT 1 FROM common_ip_country_blocks);`
+	return coredata.IsIPCountryBlocksPopulated(ctx, conn)
-
-	rows, err := conn.Query(ctx, q)
-	if err != nil {
-		return false, fmt.Errorf("cannot check if ip country blocks is populated: %w", err)
-	}
-
-	populated, err := pgx.CollectExactlyOneRow(rows, pgx.RowTo[bool])
-	if err != nil {
-		return false, fmt.Errorf("cannot collect populated check: %w", err)
-	}
-
-	return populated, nil
 }
 
 func parseCIDRFile(path string) ([]string, error) {