EvoBGP/internal/repository/postgres_prefix_snapshot.go

package repository

import (
	"context"
	"errors"
	"strings"

	"evobgp/internal/store"

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

func prefixSnapshotTableExists(ctx context.Context, q queryRower) bool {
	var n int
	err := q.QueryRow(ctx, `
		SELECT 1 FROM information_schema.tables
		WHERE table_schema = 'public' AND table_name = 'prefix_snapshot'
		LIMIT 1`).Scan(&n)
	return err == nil
}

func normalizeSnapshotHash(contentHash string) string {
	h := strings.TrimSpace(contentHash)
	if strings.HasPrefix(h, "sha256:") {
		h = strings.TrimPrefix(h, "sha256:")
	}
	if len(h) > 64 {
		h = h[:64]
	}
	if len(h) < 64 {
		h = h + strings.Repeat("0", 64-len(h))
	}
	return h
}

func (p *Postgres) revisionPrefixSnapshotID(ctx context.Context, revisionID string) (string, bool) {
	if !prefixSnapshotTableExists(ctx, p.pool) {
		return "", false
	}
	var snap *string
	err := p.pool.QueryRow(ctx, `
		SELECT prefix_snapshot_id::text FROM config_revision
		WHERE id = $1::uuid AND prefix_snapshot_id IS NOT NULL`, revisionID).Scan(&snap)
	if err != nil || snap == nil || strings.TrimSpace(*snap) == "" {
		return "", false
	}
	return *snap, true
}

func (p *Postgres) ensurePrefixSnapshot(ctx context.Context, db execQuerier, contentHash string, prefixes []store.PrefixRow) (string, error) {
	if !prefixSnapshotTableExists(ctx, db) {
		return "", nil
	}
	hash := normalizeSnapshotHash(contentHash)
	snapID, err := p.resolvePrefixSnapshotID(ctx, db, hash)
	if err != nil {
		return "", err
	}
	var rowCount int
	_ = db.QueryRow(ctx, `SELECT COUNT(*)::int FROM prefix_snapshot_row WHERE snapshot_id = $1::uuid`, snapID).Scan(&rowCount)
	if rowCount > 0 {
		return snapID, nil
	}
	for i, pr := range prefixes {
		var comm any
		if pr.CommunityID != nil && strings.TrimSpace(*pr.CommunityID) != "" {
			comm = strings.TrimSpace(*pr.CommunityID)
		}
		src := pr.Source
		if strings.TrimSpace(src) == "" {
			src = "render"
		}
		if _, err := db.Exec(ctx, `
			INSERT INTO prefix_snapshot_row (snapshot_id, ord, prefix, community_id, source)
			VALUES ($1::uuid, $2, $3, $4::uuid, $5)`,
			snapID, i, strings.TrimSpace(pr.Prefix), comm, src); err != nil {
			return "", err
		}
	}
	return snapID, nil
}

func (p *Postgres) resolvePrefixSnapshotID(ctx context.Context, db execQuerier, hash string) (string, error) {
	var existing string
	err := db.QueryRow(ctx, `SELECT id::text FROM prefix_snapshot WHERE content_hash = $1`, hash).Scan(&existing)
	if err == nil && existing != "" {
		return existing, nil
	}
	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return "", err
	}
	snapID := uuid.NewString()
	if _, err := db.Exec(ctx, `
		INSERT INTO prefix_snapshot (id, content_hash) VALUES ($1::uuid, $2)
		ON CONFLICT (content_hash) DO NOTHING`, snapID, hash); err != nil {
		return "", err
	}
	if err := db.QueryRow(ctx, `SELECT id::text FROM prefix_snapshot WHERE content_hash = $1`, hash).Scan(&snapID); err != nil {
		return "", err
	}
	return snapID, nil
}

func (p *Postgres) listSnapshotPrefixes(ctx context.Context, snapshotID, cursor string, limit int) ([]store.PrefixRow, string, bool) {
	afterOrd, off, useOffset := store.ParsePrefixPageCursor(cursor)
	var rows pgx.Rows
	var err error
	if useOffset {
		rows, err = p.pool.Query(ctx, `
			SELECT ord, prefix::text, community_id::text, source
			FROM prefix_snapshot_row
			WHERE snapshot_id = $1::uuid
			ORDER BY ord
			LIMIT $2 OFFSET $3`, snapshotID, limit+1, off)
	} else {
		var afterArg any
		if afterOrd != nil {
			afterArg = int(*afterOrd)
		}
		rows, err = p.pool.Query(ctx, `
			SELECT ord, prefix::text, community_id::text, source
			FROM prefix_snapshot_row
			WHERE snapshot_id = $1::uuid AND ($2::int IS NULL OR ord > $2::int)
			ORDER BY ord
			LIMIT $3`, snapshotID, afterArg, limit+1)
	}
	if err != nil {
		return nil, "", false
	}
	defer rows.Close()
	var all []store.PrefixRow
	var ords []int64
	for rows.Next() {
		var ord int
		var pr store.PrefixRow
		var comm *string
		if err := rows.Scan(&ord, &pr.Prefix, &comm, &pr.Source); err != nil {
			continue
		}
		pr.CommunityID = comm
		ords = append(ords, int64(ord))
		all = append(all, pr)
	}
	more := len(all) > limit
	if more {
		all = all[:limit]
		ords = ords[:limit]
	}
	next := ""
	if more && len(ords) > 0 {
		next = store.FormatPrefixPageCursor(ords[len(ords)-1])
	}
	if len(all) == 0 {
		return nil, "", false
	}
	return all, next, more
}

func (p *Postgres) linkRevisionPrefixSnapshot(ctx context.Context, db execQuerier, revisionID, snapshotID string) error {
	if snapshotID == "" || !prefixSnapshotTableExists(ctx, db) {
		return nil
	}
	_, err := db.Exec(ctx, `
		UPDATE config_revision SET prefix_snapshot_id = $2::uuid WHERE id = $1::uuid`,
		revisionID, snapshotID)
	return err
}

func (p *Postgres) copyRevisionPrefixSnapshotRef(ctx context.Context, db execQuerier, dstRevisionID, srcRevisionID string) error {
	if !prefixSnapshotTableExists(ctx, db) {
		return nil
	}
	_, err := db.Exec(ctx, `
		UPDATE config_revision dst
		SET prefix_snapshot_id = src.prefix_snapshot_id
		FROM config_revision src
		WHERE dst.id = $1::uuid AND src.id = $2::uuid AND src.prefix_snapshot_id IS NOT NULL`,
		dstRevisionID, srcRevisionID)
	return err
}

func sqlRevisionPrefixes(revParam string) string {
	return `SELECT psr.prefix FROM config_revision cr
JOIN prefix_snapshot_row psr ON psr.snapshot_id = cr.prefix_snapshot_id
WHERE cr.id = ` + revParam + `::uuid AND cr.prefix_snapshot_id IS NOT NULL
UNION ALL
SELECT rmp.prefix FROM revision_materialized_prefix rmp
WHERE rmp.revision_id = ` + revParam + `::uuid
  AND NOT EXISTS (
    SELECT 1 FROM config_revision cr2
    WHERE cr2.id = ` + revParam + `::uuid AND cr2.prefix_snapshot_id IS NOT NULL
  )`
}