Agent-System/internal/supplementimpact/pg_store.go

package supplementimpact

import (
	"context"
	"fmt"

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

// PGStore persists supplement-impact batches in PostgreSQL (migration 00070).
type PGStore struct {
	pool *pgxpool.Pool
}

// NewPGStore builds a Store backed by pool.
func NewPGStore(pool *pgxpool.Pool) *PGStore {
	return &PGStore{pool: pool}
}

func (p *PGStore) InsertBatch(ctx context.Context, b *Batch) error {
	if p == nil || p.pool == nil {
		return fmt.Errorf("supplementimpact: pg store not available")
	}
	if b == nil {
		return fmt.Errorf("supplementimpact: nil batch")
	}
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)

	_, err = tx.Exec(ctx, `
		INSERT INTO supplement_impact_batches (
			batch_id, project_id, session_id, initiator_user_id, source_kind, source_summary,
			coverage_status, status, decision_nonce, decision_digest, revision, expires_at
		) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12)
	`, b.BatchID, b.ProjectID, b.SessionID, b.InitiatorUserID, b.SourceKind, b.SourceSummary,
		b.CoverageStatus, b.Status, b.DecisionNonce, b.DecisionDigest, b.Revision, b.ExpiresAt)
	if err != nil {
		return err
	}
	for i := range b.Items {
		it := &b.Items[i]
		err = tx.QueryRow(ctx, `
			INSERT INTO supplement_impact_items (
				batch_id, doc_id, title, item_key, impact_type, depth, reason,
				author_claimed_snapshot, author_id_snapshot, busy_snapshot, locked_snapshot, change_summary,
				selected, outcome, proposal_id, run_id
			) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)
			RETURNING id
		`, it.BatchID, it.DocID, it.Title, it.ItemKey, it.ImpactType, it.Depth, it.Reason,
			it.AuthorClaimedSnapshot, it.AuthorIDSnapshot, it.BusySnapshot, it.LockedSnapshot, it.ChangeSummary,
			it.Selected, it.Outcome, it.ProposalID, it.RunID).Scan(&it.ID)
		if err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}

func (p *PGStore) GetBatch(ctx context.Context, batchID string) (*Batch, error) {
	if p == nil || p.pool == nil {
		return nil, fmt.Errorf("supplementimpact: pg store not available")
	}
	var b Batch
	err := p.pool.QueryRow(ctx, `
		SELECT batch_id, project_id, session_id, initiator_user_id, source_kind, source_summary,
		       coverage_status, status, decision_nonce, decision_digest, revision, expires_at
		FROM supplement_impact_batches WHERE batch_id = $1
	`, batchID).Scan(
		&b.BatchID, &b.ProjectID, &b.SessionID, &b.InitiatorUserID, &b.SourceKind, &b.SourceSummary,
		&b.CoverageStatus, &b.Status, &b.DecisionNonce, &b.DecisionDigest, &b.Revision, &b.ExpiresAt,
	)
	if err != nil {
		return nil, fmt.Errorf("supplementimpact: batch not found: %w", err)
	}
	rows, err := p.pool.Query(ctx, `
		SELECT id, batch_id, doc_id, title, item_key, impact_type, depth, reason,
		       author_claimed_snapshot, author_id_snapshot, busy_snapshot, locked_snapshot, change_summary,
		       selected, outcome, proposal_id, run_id
		FROM supplement_impact_items WHERE batch_id = $1 ORDER BY id
	`, batchID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var it Item
		if err := rows.Scan(
			&it.ID, &it.BatchID, &it.DocID, &it.Title, &it.ItemKey, &it.ImpactType, &it.Depth, &it.Reason,
			&it.AuthorClaimedSnapshot, &it.AuthorIDSnapshot, &it.BusySnapshot, &it.LockedSnapshot, &it.ChangeSummary,
			&it.Selected, &it.Outcome, &it.ProposalID, &it.RunID,
		); err != nil {
			return nil, err
		}
		b.Items = append(b.Items, it)
	}
	return &b, rows.Err()
}

func (p *PGStore) SaveDecision(ctx context.Context, batchID string, revision int64, nonce, digest, status string, selected map[string]bool) error {
	if p == nil || p.pool == nil {
		return fmt.Errorf("supplementimpact: pg store not available")
	}
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)

	tag, err := tx.Exec(ctx, `
		UPDATE supplement_impact_batches
		SET status = $1, decision_nonce = $2, decision_digest = $3, revision = $4, updated_at = now()
		WHERE batch_id = $5
	`, status, nonce, digest, revision, batchID)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("supplementimpact: batch not found")
	}
	if selected != nil {
		rows, qerr := tx.Query(ctx, `SELECT id, doc_id FROM supplement_impact_items WHERE batch_id = $1`, batchID)
		if qerr != nil {
			return qerr
		}
		type pair struct {
			id    int64
			docID string
		}
		var pairs []pair
		for rows.Next() {
			var pr pair
			if err := rows.Scan(&pr.id, &pr.docID); err != nil {
				rows.Close()
				return err
			}
			pairs = append(pairs, pr)
		}
		rows.Close()
		if err := rows.Err(); err != nil {
			return err
		}
		for _, pr := range pairs {
			_, err = tx.Exec(ctx, `
				UPDATE supplement_impact_items SET selected = $1, updated_at = now() WHERE id = $2
			`, selected[pr.docID], pr.id)
			if err != nil {
				return err
			}
		}
	}
	return tx.Commit(ctx)
}

func (p *PGStore) UpdateItemOutcome(ctx context.Context, batchID, docID, outcome, proposalID, runID string) error {
	if p == nil || p.pool == nil {
		return fmt.Errorf("supplementimpact: pg store not available")
	}
	tag, err := p.pool.Exec(ctx, `
		UPDATE supplement_impact_items
		SET outcome = $1, proposal_id = $2, run_id = $3, updated_at = now()
		WHERE batch_id = $4 AND doc_id = $5
	`, outcome, proposalID, runID, batchID, docID)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("supplementimpact: item not found")
	}
	return nil
}