EvoBGP/internal/repository/module_snapshot.go

package repository

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"strconv"
	"strings"
	"time"

	"evobgp/internal/store"

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

// ListModuleSnapshotSourceRows returns persisted rows for one source key.
func (p *Postgres) ListModuleSnapshotSourceRows(ctx context.Context, tenantID, moduleID, source string) ([]store.PrefixRow, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	rows, err := p.pool.Query(ctx, `
		SELECT prefix, community_id::text, source
		FROM module_prefix_snapshot_row
		WHERE tenant_id = $1::uuid AND module_id = $2::uuid AND source = $3
		ORDER BY ord`, tenantID, moduleID, source)
	if err != nil {
		return nil, fmt.Errorf("repository: read snapshot source rows: %w", err)
	}
	defer rows.Close()
	out, err := readModuleSnapshotRows(rows)
	if err != nil {
		return nil, err
	}
	return out, nil
}

func (p *Postgres) GetModulePrefixSnapshot(ctx context.Context, tenantID, moduleID string) (*store.ModulePrefixSnapshot, bool, error) {
	tx, err := p.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.RepeatableRead, AccessMode: pgx.ReadOnly})
	if err != nil {
		return nil, false, fmt.Errorf("repository: begin snapshot read: %w", err)
	}
	defer rollbackSnapshotTx(tx)
	snap, ok, err := readModuleSnapshotTx(ctx, tx, tenantID, moduleID)
	if err != nil {
		return nil, false, err
	}
	if err := tx.Commit(ctx); err != nil {
		return nil, false, fmt.Errorf("repository: commit snapshot read: %w", err)
	}
	return snap, ok, nil
}

func readModuleSnapshotTx(ctx context.Context, tx pgx.Tx, tenantID, moduleID string) (*store.ModulePrefixSnapshot, bool, error) {
	var inputHash string
	var collectedAt time.Time
	var version int64
	err := tx.QueryRow(ctx, `
		SELECT input_hash, collected_at, snapshot_version
		FROM module_prefix_snapshot
		WHERE tenant_id = $1 AND module_id = $2`,
		tenantID, moduleID).Scan(&inputHash, &collectedAt, &version)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, false, nil
		}
		return nil, false, err
	}
	rows, qerr := tx.Query(ctx, `
			SELECT prefix, community_id::text, source
			FROM module_prefix_snapshot_row
			WHERE tenant_id = $1::uuid AND module_id = $2::uuid
			ORDER BY ord`, tenantID, moduleID)
	if qerr != nil {
		return nil, false, fmt.Errorf("repository: read snapshot rows: %w", qerr)
	}
	defer rows.Close()
	prefixes, err := readModuleSnapshotRows(rows)
	if err != nil {
		return nil, false, err
	}
	return &store.ModulePrefixSnapshot{
		InputHash:   inputHash,
		CollectedAt: collectedAt.UTC(),
		Version:     version,
		Prefixes:    prefixes,
	}, true, nil
}

type moduleSnapshotRows interface {
	Next() bool
	Scan(...any) error
	Err() error
}

func readModuleSnapshotRows(rows moduleSnapshotRows) ([]store.PrefixRow, error) {
	var prefixes []store.PrefixRow
	for rows.Next() {
		var pr store.PrefixRow
		var comm *string
		if err := rows.Scan(&pr.Prefix, &comm, &pr.Source); err != nil {
			return nil, fmt.Errorf("repository: scan snapshot row: %w", err)
		}
		pr.CommunityID = comm
		prefixes = append(prefixes, pr)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("repository: iterate snapshot rows: %w", err)
	}
	return prefixes, nil
}

func (p *Postgres) SetModulePrefixSnapshot(ctx context.Context, tenantID, moduleID, inputHash string, prefixes []store.PrefixRow) error {
	if strings.TrimSpace(tenantID) == "" || strings.TrimSpace(moduleID) == "" || strings.TrimSpace(inputHash) == "" {
		return store.ErrInvalidInput
	}
	return p.UpdateModulePrefixSnapshot(ctx, tenantID, moduleID, func(*store.ModulePrefixSnapshot) (*store.ModulePrefixSnapshot, error) {
		return &store.ModulePrefixSnapshot{InputHash: inputHash, Prefixes: prefixes}, nil
	})
}

func rollbackSnapshotTx(tx pgx.Tx) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_ = tx.Rollback(ctx)
}

// UpdateModulePrefixSnapshot serializes read, mutation and write on the same transaction connection.
func (p *Postgres) UpdateModulePrefixSnapshot(ctx context.Context, tenantID, moduleID string, update store.ModuleSnapshotUpdate) error {
	if update == nil || strings.TrimSpace(tenantID) == "" || strings.TrimSpace(moduleID) == "" {
		return store.ErrInvalidInput
	}
	ctx, cancel, timeoutInstalled := store.SnapshotWriteBudget(ctx)
	defer cancel()
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("repository: begin snapshot update: %w", err)
	}
	defer rollbackSnapshotTx(tx)
	if err := setSnapshotStatementTimeout(ctx, tx, timeoutInstalled); err != nil {
		return err
	}
	if _, err = tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1::bigint)`, moduleSnapshotAdvisoryKey(tenantID, moduleID)); err != nil {
		return fmt.Errorf("repository: lock snapshot transaction: %w", err)
	}
	current, _, err := readModuleSnapshotTx(ctx, tx, tenantID, moduleID)
	if err != nil {
		return err
	}
	next, err := update(current)
	if err != nil {
		return err
	}
	if next != nil && strings.TrimSpace(next.InputHash) == "" {
		return store.ErrInvalidInput
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if next == nil {
		if _, err := tx.Exec(ctx, `DELETE FROM module_prefix_snapshot WHERE tenant_id=$1::uuid AND module_id=$2::uuid`, tenantID, moduleID); err != nil {
			return fmt.Errorf("repository: delete module snapshot: %w", err)
		}
	} else if !store.SameSnapshotContent(current, next) {
		if err := writeModuleSnapshotTx(ctx, tx, tenantID, moduleID, next); err != nil {
			return err
		}
	}
	hash := ""
	if next != nil {
		hash = next.InputHash
	}
	tag, err := tx.Exec(ctx, `UPDATE module SET input_hash=$3 WHERE tenant_id=$1::uuid AND id=$2::uuid AND deleted_at IS NULL`, tenantID, moduleID, hash)
	if err != nil {
		return fmt.Errorf("repository: update module input hash: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return store.ErrNotFound
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("repository: commit module snapshot: %w", err)
	}
	return nil
}

func writeModuleSnapshotTx(ctx context.Context, tx pgx.Tx, tenantID, moduleID string, next *store.ModulePrefixSnapshot) error {
	_, err := tx.Exec(ctx, `
		INSERT INTO module_prefix_snapshot (tenant_id, module_id, input_hash, collected_at)
		VALUES ($1::uuid, $2::uuid, $3, now())
		ON CONFLICT (tenant_id, module_id) DO UPDATE SET
			input_hash = EXCLUDED.input_hash,
			collected_at = EXCLUDED.collected_at`,
		tenantID, moduleID, next.InputHash)
	if err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `
			DELETE FROM module_prefix_snapshot_row
			WHERE tenant_id = $1::uuid AND module_id = $2::uuid`, tenantID, moduleID); err != nil {
		return err
	}
	if len(next.Prefixes) > 0 {
		_, err = tx.CopyFrom(
			ctx,
			pgx.Identifier{"module_prefix_snapshot_row"},
			[]string{"tenant_id", "module_id", "ord", "prefix", "community_id", "source"},
			pgx.CopyFromSlice(len(next.Prefixes), func(i int) ([]any, error) {
				pr := next.Prefixes[i]
				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"
				}
				return []any{tenantID, moduleID, i, strings.TrimSpace(pr.Prefix), comm, src}, nil
			}),
		)
		if err != nil {
			return err
		}
	}
	return nil
}

func (p *Postgres) DeleteModulePrefixSnapshot(ctx context.Context, tenantID, moduleID string) error {
	return p.UpdateModulePrefixSnapshot(ctx, tenantID, moduleID, func(*store.ModulePrefixSnapshot) (*store.ModulePrefixSnapshot, error) {
		return nil, nil
	})
}

func (p *Postgres) SetModuleInputHash(ctx context.Context, tenantID, moduleID, hash string) error {
	ctx, cancel, timeoutInstalled := store.SnapshotWriteBudget(ctx)
	defer cancel()
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("repository: begin module hash update: %w", err)
	}
	defer rollbackSnapshotTx(tx)
	if err := setSnapshotStatementTimeout(ctx, tx, timeoutInstalled); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1::bigint)`, moduleSnapshotAdvisoryKey(tenantID, moduleID)); err != nil {
		return fmt.Errorf("repository: lock module hash update: %w", err)
	}
	tag, err := tx.Exec(ctx, `
		UPDATE module SET input_hash = $3
		WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL`,
		moduleID, tenantID, hash)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return store.ErrNotFound
	}
	return tx.Commit(ctx)
}

// setSnapshotStatementTimeout bounds a lock when the caller did not set a deadline.
func setSnapshotStatementTimeout(ctx context.Context, tx pgx.Tx, installed bool) error {
	if !installed {
		return nil
	}
	ms := store.SnapshotStatementTimeout(ctx).Milliseconds()
	if ms < 1 {
		ms = 1
	}
	if _, err := tx.Exec(ctx, `SELECT set_config('statement_timeout', $1, true)`, strconv.FormatInt(ms, 10)); err != nil {
		return fmt.Errorf("repository: set snapshot statement timeout: %w", err)
	}
	return nil
}

func moduleSnapshotAdvisoryKey(tenantID, moduleID string) int64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte("evobgp-module-snapshot-v1\x00"))
	_, _ = h.Write([]byte(strings.TrimSpace(tenantID)))
	_, _ = h.Write([]byte{0})
	_, _ = h.Write([]byte(strings.TrimSpace(moduleID)))
	return int64(h.Sum64())
}