geoip-rest/internal/userprogram/ip_geoinfo.go

package userprogram

import (
	"context"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"

	"geoip-rest/internal/geo"
)

func EnsureIPGeoInfoTable(ctx context.Context, conn *pgx.Conn, schema string) error {
	ddl := fmt.Sprintf(`
CREATE TABLE IF NOT EXISTS %s.ip_geoinfo (
  id bigserial PRIMARY KEY,
  ip inet UNIQUE NOT NULL,
  country text,
  region text,
  city text,
  address text,
  latitude double precision,
  longitude double precision
);`, pgx.Identifier{schema}.Sanitize())
	_, err := conn.Exec(ctx, ddl)
	return err
}

const defaultSeedPath = "/initial_data/ip_geoinfo_seed_20251208.sql"

// SeedIPGeoInfoIfMissing applies the seed SQL when ip_geoinfo is absent.
func SeedIPGeoInfoIfMissing(ctx context.Context, conn *pgx.Conn, schema string) error {
	exists, err := ipGeoInfoExists(ctx, conn, schema)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}
	if _, err := os.Stat(defaultSeedPath); err == nil {
		if err := ExecuteSQLFile(ctx, conn, defaultSeedPath); err != nil {
			return fmt.Errorf("execute seed sql: %w", err)
		}
	}
	return EnsureIPGeoInfoTable(ctx, conn, schema)
}

func ipGeoInfoExists(ctx context.Context, conn *pgx.Conn, schema string) (bool, error) {
	var exists bool
	err := conn.QueryRow(ctx, `
SELECT EXISTS (
  SELECT 1 FROM information_schema.tables
  WHERE table_schema = $1 AND table_name = 'ip_geoinfo'
);`, schema).Scan(&exists)
	return exists, err
}

// ExportPublicIPs writes distinct login_public_ip values to a CSV file with header.
func ExportPublicIPs(ctx context.Context, conn *pgx.Conn, schema, path string) error {
	rows, err := conn.Query(ctx, fmt.Sprintf(`
SELECT DISTINCT login_public_ip
FROM %s.user_program_info_replica
WHERE login_public_ip IS NOT NULL AND login_public_ip <> ''
ORDER BY login_public_ip;`, pgx.Identifier{schema}.Sanitize()))
	if err != nil {
		return err
	}
	defer rows.Close()

	var ips []string
	for rows.Next() {
		var ip string
		if err := rows.Scan(&ip); err != nil {
			return err
		}
		ips = append(ips, ip)
	}
	if rows.Err() != nil {
		return rows.Err()
	}

	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return err
	}
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	if _, err := f.WriteString(`"login_public_ip"` + "\n"); err != nil {
		return err
	}
	for _, ip := range ips {
		if _, err := f.WriteString(fmt.Sprintf(`"%s"`+"\n", ip)); err != nil {
			return err
		}
	}
	return nil
}

// GenerateIPGeoInfoSQL builds an upsert SQL file for IPs. If onlyNew is true, it skips
// IPs already present in ip_geoinfo.
func GenerateIPGeoInfoSQL(ctx context.Context, conn *pgx.Conn, schema string, resolver geo.Resolver, output string, onlyNew bool) (int, error) {
	query := fmt.Sprintf(`
SELECT DISTINCT login_public_ip
FROM %s.user_program_info_replica r
WHERE login_public_ip IS NOT NULL AND login_public_ip <> ''`, pgx.Identifier{schema}.Sanitize())
	if onlyNew {
		query += fmt.Sprintf(`
  AND NOT EXISTS (
    SELECT 1 FROM %s.ip_geoinfo g WHERE g.ip = r.login_public_ip::inet
  )`, pgx.Identifier{schema}.Sanitize())
	}
	query += ";"

	rows, err := conn.Query(ctx, query)
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	var ips []string
	for rows.Next() {
		var ip string
		if err := rows.Scan(&ip); err != nil {
			return 0, err
		}
		ips = append(ips, ip)
	}
	if rows.Err() != nil {
		return 0, rows.Err()
	}
	if len(ips) == 0 {
		return 0, nil
	}

	sort.Strings(ips)

	if err := os.MkdirAll(filepath.Dir(output), 0o755); err != nil {
		return 0, err
	}
	f, err := os.Create(output)
	if err != nil {
		return 0, err
	}
	defer f.Close()

	header := fmt.Sprintf("-- generated at %s KST\n", time.Now().In(kst()).Format(time.RFC3339))
	header += fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s;\n", schemaIdent(schema))
	header += fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.ip_geoinfo (
  id bigserial PRIMARY KEY,
  ip inet UNIQUE NOT NULL,
  country text,
  region text,
  city text,
  address text,
  latitude double precision,
  longitude double precision
);`+"\n", schemaIdent(schema))
	if _, err := f.WriteString(header); err != nil {
		return 0, err
	}

	count := 0
	for _, ip := range ips {
		loc, err := resolver.Lookup(ip)
		if err != nil {
			continue
		}
		stmt := fmt.Sprintf(`INSERT INTO %s.ip_geoinfo (ip, country, region, city, address, latitude, longitude)
VALUES ('%s', %s, %s, %s, %s, %f, %f)
ON CONFLICT (ip) DO UPDATE SET
  country = EXCLUDED.country,
  region = EXCLUDED.region,
  city = EXCLUDED.city,
  address = EXCLUDED.address,
  latitude = EXCLUDED.latitude,
  longitude = EXCLUDED.longitude;
`, schemaIdent(schema), toHostCIDR(ip), quote(loc.Country), quote(loc.Region), quote(loc.City), quote(loc.Address), loc.Latitude, loc.Longitude)
		if _, err := f.WriteString(stmt); err != nil {
			return count, err
		}
		count++
	}

	return count, nil
}

func ExecuteSQLFile(ctx context.Context, conn *pgx.Conn, path string) error {
	content, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	_, err = conn.Exec(ctx, string(content))
	return err
}

func toHostCIDR(ipStr string) string {
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return ""
	}
	if ip.To4() != nil {
		return ip.String() + "/32"
	}
	return ip.String() + "/128"
}

func quote(s string) string {
	return fmt.Sprintf("'%s'", strings.ReplaceAll(s, "'", "''"))
}

func schemaIdent(s string) string {
	return `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
}