Use shadow-table swap for IP country blocks import

Replace TRUNCATE+COPY inside a single transaction with a staging table
approach so that SELECT queries from the cookie banner handler are never
blocked during import. The ACCESS EXCLUSIVE lock is now only held for
the sub-millisecond DROP+RENAME swap.

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

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 {
+const ipCountryBlocksStagingTable = "common_ip_country_blocks_staging"
-	_, err := conn.Exec(ctx, "TRUNCATE common_ip_country_blocks")
+
-	if err != nil {
+func CreateIPCountryBlocksStaging(ctx context.Context, conn pg.Querier) error {
-		return fmt.Errorf("cannot truncate common_ip_country_blocks: %w", err)
+	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 {
+func CopyIPCountryBlocksStaging(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"},
+		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

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(
+	if err := s.pgClient.WithConn(ctx, func(ctx context.Context, conn pg.Querier) error {
-		ctx,
+		if err := coredata.CreateIPCountryBlocksStaging(ctx, conn); err != nil {
-		func(ctx context.Context, tx pg.Tx) error {
+			return fmt.Errorf("cannot create staging table: %w", err)
-			if err := coredata.TruncateIPCountryBlocks(ctx, tx); err != nil {
-				return err
 		}
 
-			if err := coredata.CopyIPCountryBlocks(ctx, tx, blocks); err != nil {
+		if err := coredata.CopyIPCountryBlocksStaging(ctx, conn, blocks); err != nil {
-				return err
+			return fmt.Errorf("cannot copy IP country blocks to staging: %w", err)
+		}
+
+		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)
+			if cc, err = coredata.LookupCountryByIP(ctx, conn, ip); err != nil {
-			return err
+				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)
+			if populated, err = coredata.IsIPCountryBlocksPopulated(ctx, conn); err != nil {
-			return err
+				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)