Agent-System/internal/traceimpact/pg_store.go

package traceimpact

import (
	"context"
	"fmt"

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

type pgStore struct {
	pool *pgxpool.Pool
}

func newPGStore(pool *pgxpool.Pool) *pgStore {
	return &pgStore{pool: pool}
}

func (p *pgStore) UpsertItem(ctx context.Context, item Item) (Item, error) {
	if item.Lifecycle == "" {
		item.Lifecycle = "active"
	}
	row := p.pool.QueryRow(ctx, `
		INSERT INTO project_trace_items
			(project_id, business_key, kind, lifecycle_state, review_state, tombstone, statement, source_doc_id)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
		ON CONFLICT (project_id, kind, business_key) DO UPDATE SET
			review_state = EXCLUDED.review_state,
			tombstone = EXCLUDED.tombstone,
			statement = CASE WHEN EXCLUDED.statement <> '' THEN EXCLUDED.statement ELSE project_trace_items.statement END,
			source_doc_id = CASE WHEN EXCLUDED.source_doc_id <> '' THEN EXCLUDED.source_doc_id ELSE project_trace_items.source_doc_id END
		RETURNING id, project_id, business_key, kind, lifecycle_state, review_state, tombstone, statement, source_doc_id
	`, item.ProjectID, item.BusinessKey, item.Kind, item.Lifecycle, item.ReviewState, item.Tombstone, item.Statement, item.SourceDocID)
	return scanItem(row)
}

func scanItem(row pgx.Row) (Item, error) {
	var item Item
	err := row.Scan(&item.ID, &item.ProjectID, &item.BusinessKey, &item.Kind, &item.Lifecycle, &item.ReviewState, &item.Tombstone, &item.Statement, &item.SourceDocID)
	return item, err
}

func (p *pgStore) ListItems(ctx context.Context, projectID int64) ([]Item, error) {
	rows, err := p.pool.Query(ctx, `
		SELECT id, project_id, business_key, kind, lifecycle_state, review_state, tombstone, statement, source_doc_id
		FROM project_trace_items WHERE project_id = $1 ORDER BY id
	`, projectID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Item
	for rows.Next() {
		var item Item
		if err := rows.Scan(&item.ID, &item.ProjectID, &item.BusinessKey, &item.Kind, &item.Lifecycle, &item.ReviewState, &item.Tombstone, &item.Statement, &item.SourceDocID); err != nil {
			return nil, err
		}
		out = append(out, item)
	}
	return out, rows.Err()
}

func (p *pgStore) UpsertEdge(ctx context.Context, edge Edge) (Edge, error) {
	row := p.pool.QueryRow(ctx, `
		INSERT INTO project_trace_edges
			(project_id, from_item_id, to_item_id, relation_type, provenance, review_state, evidence_doc_id)
		VALUES ($1,$2,$3,$4,$5,$6,$7)
		ON CONFLICT (project_id, from_item_id, to_item_id, relation_type, provenance) DO UPDATE SET
			review_state = CASE
				WHEN project_trace_edges.review_state = 'rejected' AND EXCLUDED.review_state = 'candidate'
					THEN project_trace_edges.review_state
				WHEN project_trace_edges.review_state = 'confirmed'
					THEN project_trace_edges.review_state
				ELSE EXCLUDED.review_state
			END
		RETURNING id, project_id, from_item_id, to_item_id, relation_type, provenance, review_state, evidence_doc_id
	`, edge.ProjectID, edge.FromItemID, edge.ToItemID, edge.RelationType, edge.Provenance, edge.ReviewState, edge.EvidenceDoc)
	var out Edge
	err := row.Scan(&out.ID, &out.ProjectID, &out.FromItemID, &out.ToItemID, &out.RelationType, &out.Provenance, &out.ReviewState, &out.EvidenceDoc)
	return out, err
}

func (p *pgStore) ListEdges(ctx context.Context, projectID int64) ([]Edge, error) {
	rows, err := p.pool.Query(ctx, `
		SELECT id, project_id, from_item_id, to_item_id, relation_type, provenance, review_state, evidence_doc_id
		FROM project_trace_edges WHERE project_id = $1 ORDER BY id
	`, projectID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Edge
	for rows.Next() {
		var e Edge
		if err := rows.Scan(&e.ID, &e.ProjectID, &e.FromItemID, &e.ToItemID, &e.RelationType, &e.Provenance, &e.ReviewState, &e.EvidenceDoc); err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

func (p *pgStore) PutIndex(ctx context.Context, st IndexState) error {
	if st.Extractor == "" {
		st.Extractor = ExtractorVersion
	}
	_, err := p.pool.Exec(ctx, `
		INSERT INTO project_trace_index_state
			(project_id, doc_id, source_snapshot_id, status, extractor_version, item_count, error)
		VALUES ($1,$2,$3,$4,$5,$6,$7)
		ON CONFLICT (project_id, doc_id, extractor_version) DO UPDATE SET
			source_snapshot_id = EXCLUDED.source_snapshot_id,
			status = EXCLUDED.status,
			item_count = EXCLUDED.item_count,
			error = EXCLUDED.error,
			indexed_at = now()
	`, st.ProjectID, st.DocID, st.SnapshotID, st.Status, st.Extractor, st.ItemCount, st.Error)
	return err
}

func (p *pgStore) IndexStatus(ctx context.Context, projectID int64, docID string) (IndexState, bool, error) {
	row := p.pool.QueryRow(ctx, `
		SELECT project_id, doc_id, source_snapshot_id, status, extractor_version, item_count, error
		FROM project_trace_index_state
		WHERE project_id = $1 AND doc_id = $2
		ORDER BY indexed_at DESC LIMIT 1
	`, projectID, docID)
	var st IndexState
	err := row.Scan(&st.ProjectID, &st.DocID, &st.SnapshotID, &st.Status, &st.Extractor, &st.ItemCount, &st.Error)
	if err == pgx.ErrNoRows {
		return IndexState{}, false, nil
	}
	if err != nil {
		return IndexState{}, false, err
	}
	return st, true, nil
}

func (p *pgStore) ListIndex(ctx context.Context, projectID int64) ([]IndexState, error) {
	rows, err := p.pool.Query(ctx, `
		SELECT project_id, doc_id, source_snapshot_id, status, extractor_version, item_count, error
		FROM project_trace_index_state WHERE project_id = $1 ORDER BY doc_id
	`, projectID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []IndexState
	for rows.Next() {
		var st IndexState
		if err := rows.Scan(&st.ProjectID, &st.DocID, &st.SnapshotID, &st.Status, &st.Extractor, &st.ItemCount, &st.Error); err != nil {
			return nil, err
		}
		out = append(out, st)
	}
	return out, rows.Err()
}

func (p *pgStore) SetEdgeReview(ctx context.Context, edgeID int64, review string) error {
	tag, err := p.pool.Exec(ctx, `UPDATE project_trace_edges SET review_state = $2 WHERE id = $1`, edgeID, review)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("traceimpact: edge not found")
	}
	return nil
}

func (p *pgStore) AddDecision(ctx context.Context, d Decision) (Decision, error) {
	if d.ProjectID == 0 {
		_ = p.pool.QueryRow(ctx, `SELECT project_id FROM project_trace_edges WHERE id = $1`, d.EdgeID).Scan(&d.ProjectID)
	}
	err := p.pool.QueryRow(ctx, `
		INSERT INTO project_trace_edge_decisions (project_id, edge_id, decision, reason, decided_by)
		VALUES ($1,$2,$3,$4,$5) RETURNING id
	`, d.ProjectID, d.EdgeID, d.Decision, d.Reason, d.DecidedBy).Scan(&d.ID)
	return d, err
}

func (p *pgStore) ListDecisions(ctx context.Context, projectID int64) ([]Decision, error) {
	rows, err := p.pool.Query(ctx, `
		SELECT id, project_id, edge_id, decision, reason, decided_by
		FROM project_trace_edge_decisions WHERE project_id = $1 ORDER BY id
	`, projectID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Decision
	for rows.Next() {
		var d Decision
		if err := rows.Scan(&d.ID, &d.ProjectID, &d.EdgeID, &d.Decision, &d.Reason, &d.DecidedBy); err != nil {
			return nil, err
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

func (p *pgStore) PutReopen(ctx context.Context, req ReopenRequest) (ReopenRequest, error) {
	if req.ID == 0 {
		err := p.pool.QueryRow(ctx, `
			INSERT INTO stage_reopen_requests (project_id, instance_id, stage_key, doc_task_keys, reason, requested_by, status)
			VALUES ($1,$2,$3,$4,$5,$6,$7) RETURNING id
		`, req.ProjectID, req.InstanceID, req.StageKey, req.DocTasks, req.Reason, req.RequestedBy, req.Status).Scan(&req.ID)
		return req, err
	}
	_, err := p.pool.Exec(ctx, `UPDATE stage_reopen_requests SET status = $2 WHERE id = $1`, req.ID, req.Status)
	return req, err
}

func (p *pgStore) GetReopen(ctx context.Context, id int64) (ReopenRequest, error) {
	var req ReopenRequest
	err := p.pool.QueryRow(ctx, `
		SELECT id, project_id, instance_id, stage_key, doc_task_keys, reason, requested_by, status
		FROM stage_reopen_requests WHERE id = $1
	`, id).Scan(&req.ID, &req.ProjectID, &req.InstanceID, &req.StageKey, &req.DocTasks, &req.Reason, &req.RequestedBy, &req.Status)
	if err == pgx.ErrNoRows {
		return ReopenRequest{}, fmt.Errorf("traceimpact: reopen request not found")
	}
	return req, err
}

func (p *pgStore) PutMark(ctx context.Context, mark RecheckMark) (RecheckMark, error) {
	if mark.Status == "" {
		mark.Status = MarkOpen
	}
	err := p.pool.QueryRow(ctx, `
		INSERT INTO stage_recheck_marks (project_id, stage_key, source_doc_id, status, reason)
		VALUES ($1,$2,$3,$4,$5)
		ON CONFLICT (project_id, stage_key, source_doc_id) WHERE status = 'open'
		DO UPDATE SET reason = stage_recheck_marks.reason
		RETURNING id, project_id, stage_key, source_doc_id, status, reason, COALESCE(acked_by, 0)
	`, mark.ProjectID, mark.StageKey, mark.SourceDoc, mark.Status, mark.Reason).Scan(
		&mark.ID, &mark.ProjectID, &mark.StageKey, &mark.SourceDoc, &mark.Status, &mark.Reason, &mark.AckedBy,
	)
	return mark, err
}

func (p *pgStore) ListOpenMarks(ctx context.Context, projectID int64, stageKey string) ([]RecheckMark, error) {
	rows, err := p.pool.Query(ctx, `
		SELECT id, project_id, stage_key, source_doc_id, status, reason, COALESCE(acked_by, 0)
		FROM stage_recheck_marks
		WHERE project_id = $1 AND status = 'open' AND ($2 = '' OR stage_key = $2)
		ORDER BY id
	`, projectID, stageKey)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []RecheckMark
	for rows.Next() {
		var m RecheckMark
		if err := rows.Scan(&m.ID, &m.ProjectID, &m.StageKey, &m.SourceDoc, &m.Status, &m.Reason, &m.AckedBy); err != nil {
			return nil, err
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

func (p *pgStore) AckMark(ctx context.Context, id, userID int64) error {
	tag, err := p.pool.Exec(ctx, `UPDATE stage_recheck_marks SET acked_by = $2 WHERE id = $1 AND status = 'open'`, id, userID)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("traceimpact: mark not found")
	}
	return nil
}

func (p *pgStore) SetMarkStatus(ctx context.Context, id int64, status, reason string) error {
	tag, err := p.pool.Exec(ctx, `UPDATE stage_recheck_marks SET status = $2, reason = CASE WHEN $3 <> '' THEN $3 ELSE reason END WHERE id = $1`, id, status, reason)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("traceimpact: mark not found")
	}
	return nil
}

func (p *pgStore) GetEdge(ctx context.Context, edgeID int64) (Edge, error) {
	var e Edge
	err := p.pool.QueryRow(ctx, `
		SELECT id, project_id, from_item_id, to_item_id, relation_type, provenance, review_state, evidence_doc_id
		FROM project_trace_edges WHERE id = $1
	`, edgeID).Scan(&e.ID, &e.ProjectID, &e.FromItemID, &e.ToItemID, &e.RelationType, &e.Provenance, &e.ReviewState, &e.EvidenceDoc)
	if err == pgx.ErrNoRows {
		return Edge{}, fmt.Errorf("traceimpact: edge not found")
	}
	return e, err
}

func (p *pgStore) AppendRevision(ctx context.Context, rev ItemRevision) error {
	if rev.Extractor == "" {
		rev.Extractor = ExtractorVersion
	}
	_, err := p.pool.Exec(ctx, `
		INSERT INTO project_trace_item_revisions
			(item_id, revision_no, source_doc_id, source_snapshot_id, content_hash, statement, extractor_version, tombstone)
		SELECT $1,
			COALESCE((SELECT MAX(revision_no) + 1 FROM project_trace_item_revisions WHERE item_id = $1), 1),
			$2, $3, $4, $5, $6, $7
		WHERE NOT EXISTS (
			SELECT 1 FROM project_trace_item_revisions
			WHERE item_id = $1 AND source_snapshot_id = $3 AND content_hash = $4 AND extractor_version = $6
		)
	`, rev.ItemID, rev.SourceDocID, rev.SnapshotID, rev.ContentHash, rev.Statement, rev.Extractor, rev.Tombstone)
	return err
}

func (p *pgStore) RevisionCount(ctx context.Context, itemID int64) (int, error) {
	var n int
	err := p.pool.QueryRow(ctx, `SELECT COUNT(*) FROM project_trace_item_revisions WHERE item_id = $1`, itemID).Scan(&n)
	return n, err
}

func (p *pgStore) AddProposal(ctx context.Context, projectID, userID int64, text string) (int64, error) {
	var id int64
	err := p.pool.QueryRow(ctx, `
		INSERT INTO trace_revision_proposals (project_id, author_id, body) VALUES ($1,$2,$3) RETURNING id
	`, projectID, userID, text).Scan(&id)
	return id, err
}