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, `"`, `""`) + `"` }