EvoBGP/internal/repository/asn_cache.go

package repository

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/netip"
	"sort"
	"strings"
	"sync"
	"time"

	"evobgp/internal/store"

	"github.com/jackc/pgx/v5"
)

// asnCacheRowExistsOnce caches the schema check for the process lifetime (see moduleSnapshotRowTableExists).
var (
	asnCacheRowExistsOnce   sync.Once
	asnCacheRowExistsCached bool
)

func asnCacheRowTableExists(ctx context.Context, q queryRower) bool {
	asnCacheRowExistsOnce.Do(func() {
		var n int
		err := q.QueryRow(ctx, `
			SELECT 1 FROM information_schema.tables
			WHERE table_schema = 'public' AND table_name = 'asn_prefix_cache_row'
			LIMIT 1`).Scan(&n)
		if err != nil {
			asnCacheRowExistsOnce = sync.Once{}
			return
		}
		asnCacheRowExistsCached = true
	})
	return asnCacheRowExistsCached
}

func (p *Postgres) GetASNPrefixCache(asn int64) (*store.ASNPrefixCacheEntry, bool, error) {
	ctx := context.Background()
	var holder string
	var fetchedAt time.Time
	err := p.pool.QueryRow(ctx, `
		SELECT holder, fetched_at FROM asn_prefix_cache WHERE asn = $1`, asn).
		Scan(&holder, &fetchedAt)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, false, nil
		}
		return nil, false, err
	}
	var prefixes []string
	if asnCacheRowTableExists(ctx, p.pool) {
		rows, qerr := p.pool.Query(ctx, `
			SELECT prefix::text FROM asn_prefix_cache_row WHERE asn = $1 ORDER BY prefix`, asn)
		if qerr != nil {
			return nil, false, qerr
		}
		defer rows.Close()
		for rows.Next() {
			var s string
			if err := rows.Scan(&s); err != nil {
				continue
			}
			prefixes = append(prefixes, s)
		}
	} else {
		var raw []byte
		if err := p.pool.QueryRow(ctx, `SELECT prefixes_json FROM asn_prefix_cache WHERE asn = $1`, asn).Scan(&raw); err == nil && len(raw) > 0 {
			_ = json.Unmarshal(raw, &prefixes)
		}
	}
	return &store.ASNPrefixCacheEntry{
		ASN:       asn,
		Holder:    holder,
		Prefixes:  prefixes,
		FetchedAt: fetchedAt.UTC(),
	}, true, nil
}

func (p *Postgres) SetASNPrefixCache(asn int64, holder string, prefixes []string) error {
	ctx := context.Background()
	normalized, err := canonicalASNPrefixes(prefixes)
	if err != nil {
		return err
	}
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	_, err = tx.Exec(ctx, `
			INSERT INTO asn_prefix_cache (asn, holder, fetched_at)
			VALUES ($1, $2, now())
			ON CONFLICT (asn) DO UPDATE SET
				holder = EXCLUDED.holder,
				fetched_at = EXCLUDED.fetched_at`, asn, holder)
	if err != nil {
		return err
	}
	if asnCacheRowTableExists(ctx, tx) {
		existing, err := listASNCachePrefixes(ctx, tx, asn)
		if err != nil {
			return err
		}
		if prefixSetEqual(existing, normalized) {
			return tx.Commit(ctx)
		}
		if _, err := tx.Exec(ctx, `DELETE FROM asn_prefix_cache_row WHERE asn = $1`, asn); err != nil {
			return err
		}
		if len(normalized) > 0 {
			if _, err := tx.CopyFrom(ctx, pgx.Identifier{"asn_prefix_cache_row"}, []string{"asn", "prefix"}, pgx.CopyFromSlice(len(normalized), func(i int) ([]any, error) {
				pfx, err := netip.ParsePrefix(normalized[i])
				if err != nil {
					return nil, fmt.Errorf("repository: asn cache prefix %q: %w", normalized[i], err)
				}
				return []any{asn, pfx.Masked()}, nil
			})); err != nil {
				return fmt.Errorf("repository: copy asn cache rows: %w", err)
			}
		}
	} else {
		raw, err := json.Marshal(normalized)
		if err != nil {
			return err
		}
		if _, err := tx.Exec(ctx, `
				UPDATE asn_prefix_cache SET prefixes_json = $2::jsonb WHERE asn = $1`, asn, string(raw)); err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}

func listASNCachePrefixes(ctx context.Context, tx pgx.Tx, asn int64) ([]string, error) {
	rows, err := tx.Query(ctx, `SELECT prefix::text FROM asn_prefix_cache_row WHERE asn = $1`, asn)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []string
	for rows.Next() {
		var s string
		if err := rows.Scan(&s); err != nil {
			return nil, err
		}
		out = append(out, s)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	return canonicalASNPrefixes(out)
}

func canonicalASNPrefixes(in []string) ([]string, error) {
	seen := make(map[string]struct{}, len(in))
	out := make([]string, 0, len(in))
	for _, raw := range in {
		raw = strings.TrimSpace(raw)
		if raw == "" {
			continue
		}
		pfx, err := netip.ParsePrefix(raw)
		if err != nil {
			return nil, fmt.Errorf("repository: asn cache prefix %q: %w", raw, err)
		}
		s := pfx.Masked().String()
		if _, ok := seen[s]; ok {
			continue
		}
		seen[s] = struct{}{}
		out = append(out, s)
	}
	sort.Strings(out)
	return out, nil
}

func prefixSetEqual(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}