package repository import ( "context" "encoding/json" "errors" "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() 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) { if _, err := tx.Exec(ctx, `DELETE FROM asn_prefix_cache_row WHERE asn = $1`, asn); err != nil { return err } for _, pfx := range prefixes { if _, err := tx.Exec(ctx, ` INSERT INTO asn_prefix_cache_row (asn, prefix) VALUES ($1, $2::cidr)`, asn, pfx); err != nil { return err } } } else { raw, err := json.Marshal(prefixes) 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) }