Prism/internal/store/country_repo.go

package store

import (
	"context"

	"github.com/prism/proxy/internal/country"
)

type CountryStat struct {
	Code  string `json:"code"`
	Name  string `json:"name"`
	Count int    `json:"count"`
}

func (s *Store) ListCountryStats(ctx context.Context) ([]CountryStat, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT country, COUNT(*) FROM proxy_nodes
		WHERE enabled=1 AND country != ''
		GROUP BY country ORDER BY COUNT(*) DESC, country`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []CountryStat
	for rows.Next() {
		var st CountryStat
		if err := rows.Scan(&st.Code, &st.Count); err != nil {
			return nil, err
		}
		st.Name = country.Name(st.Code)
		out = append(out, st)
	}
	return out, rows.Err()
}

func (s *Store) BackfillNodeCountries(ctx context.Context) (int, error) {
	nodes, err := s.ListProxyNodes(ctx, "", nil)
	if err != nil {
		return 0, err
	}
	updated := 0
	for _, n := range nodes {
		code := n.Country
		if code == "" {
			code = country.Detect(n.Name)
		}
		if code == "" || code == n.Country {
			continue
		}
		if err := s.UpdateProxyNodeCountry(ctx, n.ID, code); err != nil {
			return updated, err
		}
		updated++
	}
	return updated, nil
}

func (s *Store) UpdateProxyNodeCountry(ctx context.Context, id int64, code string) error {
	_, err := s.db.ExecContext(ctx, `UPDATE proxy_nodes SET country=?, updated_at=datetime('now') WHERE id=?`, code, id)
	return err
}

func (s *Store) RebuildAllNodeCountries(ctx context.Context) (int, error) {
	nodes, err := s.ListProxyNodes(ctx, "", nil)
	if err != nil {
		return 0, err
	}
	updated := 0
	for _, n := range nodes {
		code := country.Detect(n.Name)
		if code == n.Country {
			continue
		}
		if err := s.UpdateProxyNodeCountry(ctx, n.ID, code); err != nil {
			return updated, err
		}
		updated++
	}
	return updated, nil
}