Agent-System/internal/dataset/store.go

package dataset

import (
	"context"

	"chat-agent-system/internal/artifactvault"

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

type PG struct {
	Pool *pgxpool.Pool
	Env  string
}

func (p *PG) UpsertDataset(ctx context.Context, d Dataset) error {
	if err := ValidateClass(d.Class); err != nil {
		return err
	}
	_, err := p.Pool.Exec(ctx, `
		INSERT INTO dataset_registry (dataset_id, class, display_name, storage_identity, source_schema, source_commit, status, project_count)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
		ON CONFLICT (dataset_id) DO UPDATE SET
			class=EXCLUDED.class, display_name=EXCLUDED.display_name, storage_identity=EXCLUDED.storage_identity,
			source_schema=EXCLUDED.source_schema, source_commit=EXCLUDED.source_commit, status=EXCLUDED.status,
			project_count=EXCLUDED.project_count
	`, d.ID, d.Class, d.DisplayName, d.StorageIdentity, d.SourceSchema, d.SourceCommit, nz(d.Status, "approved"), d.ProjectCount)
	return err
}

func (p *PG) EnsureJob(ctx context.Context, job *Job) error {
	_, err := p.Pool.Exec(ctx, `
		INSERT INTO dataset_backfill_jobs (job_id, status, source_schema, source_commit, storage_identity)
		VALUES ($1,$2,$3,$4,$5)
		ON CONFLICT (job_id) DO NOTHING
	`, job.ID, nz(job.Status, JobRunning), job.SourceSchema, job.SourceCommit, job.StorageIdentity)
	return err
}

func (p *PG) SetJobStatus(ctx context.Context, jobID, status string) error {
	_, err := p.Pool.Exec(ctx, `
		UPDATE dataset_backfill_jobs
		SET status=$2,
		    paused_at=CASE WHEN $2='paused' THEN now() ELSE paused_at END,
		    updated_at=now()
		WHERE job_id=$1`, jobID, status)
	return err
}

func (p *PG) EnqueueKind(ctx context.Context, job *Job, kind string, shard int) (int, error) {
	total := 0
	last := int64(0)
	for {
		items, err := artifactvault.Enumerate(ctx, p.Pool, kind, last, artifactvault.BackfillBatchSize)
		if err != nil {
			return total, err
		}
		if len(items) == 0 {
			return total, nil
		}
		for _, c := range items {
			_, err := p.Pool.Exec(ctx, `
				INSERT INTO dataset_backfill_queue (job_id, source_kind, source_pk, shard, status, logical_id)
				VALUES ($1,$2,$3,$4,'pending',$5)
				ON CONFLICT (job_id, source_kind, source_pk) DO NOTHING
			`, job.ID, kind, c.SourcePK, shard, c.LogicalID)
			if err != nil {
				return total, err
			}
			last = c.SourcePK
			total++
		}
	}
}

const claimSQL = `
WITH cte AS (
	SELECT job_id, source_kind, source_pk
	FROM dataset_backfill_queue
	WHERE job_id=$1 AND status='pending'
	ORDER BY source_pk
	FOR UPDATE SKIP LOCKED
	LIMIT $2
)
UPDATE dataset_backfill_queue q
SET status='claimed', claimed_at=now()
FROM cte
WHERE q.job_id=cte.job_id AND q.source_kind=cte.source_kind AND q.source_pk=cte.source_pk
RETURNING q.source_kind, q.source_pk, q.logical_id, q.shard
`

func (p *PG) ClaimAndProcess(ctx context.Context, job *Job, limit int) (JobStats, error) {
	if p.Env == "" {
		p.Env = "dev"
	}
	tx, err := p.Pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return JobStats{}, err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	rows, err := tx.Query(ctx, claimSQL, job.ID, limit)
	if err != nil {
		return JobStats{}, err
	}
	type claimed struct {
		kind, logical string
		pk            int64
		shard         int
	}
	var batch []claimed
	for rows.Next() {
		var it claimed
		if err := rows.Scan(&it.kind, &it.pk, &it.logical, &it.shard); err != nil {
			rows.Close()
			return JobStats{}, err
		}
		batch = append(batch, it)
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return JobStats{}, err
	}

	stats := JobStats{}
	for _, it := range batch {
		c, err := p.loadCandidate(ctx, tx, it.kind, it.pk)
		if err != nil {
			return stats, err
		}
		c.LogicalID = nz(c.LogicalID, it.logical)
		c.SourceSchema = job.SourceSchema
		res, err := artifactvault.ProcessCandidate(ctx, tx, p.Env, c)
		if err != nil {
			return stats, err
		}
		stats.Scanned++
		status := "done"
		succ, quar := 1, 0
		if res.Status == "quarantine" {
			status = "quarantine"
			succ, quar = 0, 1
			stats.Quarantined++
		} else {
			stats.Succeeded++
		}
		if _, err := tx.Exec(ctx, `
			UPDATE dataset_backfill_queue SET status=$4 WHERE job_id=$1 AND source_kind=$2 AND source_pk=$3
		`, job.ID, it.kind, it.pk, status); err != nil {
			return stats, err
		}
		if _, err := tx.Exec(ctx, `
			INSERT INTO dataset_backfill_checkpoints (job_id, shard, source_kind, last_pk, scanned, succeeded, quarantined)
			VALUES ($1,$2,$3,$4,1,$5,$6)
			ON CONFLICT (job_id, shard, source_kind) DO UPDATE SET
				last_pk=GREATEST(dataset_backfill_checkpoints.last_pk, EXCLUDED.last_pk),
				scanned=dataset_backfill_checkpoints.scanned+1,
				succeeded=dataset_backfill_checkpoints.succeeded+EXCLUDED.succeeded,
				quarantined=dataset_backfill_checkpoints.quarantined+EXCLUDED.quarantined
		`, job.ID, it.shard, it.kind, it.pk, succ, quar); err != nil {
			return stats, err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return stats, err
	}
	return stats, nil
}

func (p *PG) loadCandidate(ctx context.Context, tx pgx.Tx, kind string, pk int64) (artifactvault.Candidate, error) {
	items, err := artifactvault.Enumerate(ctx, tx, kind, pk-1, 1)
	if err != nil {
		return artifactvault.Candidate{}, err
	}
	if len(items) == 0 || items[0].SourcePK != pk {
		return artifactvault.Candidate{Kind: kind, SourcePK: pk}, nil
	}
	return items[0], nil
}

func (p *PG) RecalcJob(ctx context.Context, jobID string) (JobStats, error) {
	var s JobStats
	err := p.Pool.QueryRow(ctx, `
		SELECT COALESCE(SUM(scanned),0), COALESCE(SUM(succeeded),0), COALESCE(SUM(quarantined),0)
		FROM dataset_backfill_checkpoints WHERE job_id=$1
	`, jobID).Scan(&s.Scanned, &s.Succeeded, &s.Quarantined)
	return s, err
}

func (p *PG) CountQueue(ctx context.Context, jobID, status string) (int, error) {
	var n int
	err := p.Pool.QueryRow(ctx, `SELECT COUNT(*) FROM dataset_backfill_queue WHERE job_id=$1 AND status=$2`, jobID, status).Scan(&n)
	return n, err
}

func (p *PG) CountArtifacts(ctx context.Context) (int, error) {
	var n int
	err := p.Pool.QueryRow(ctx, `SELECT COUNT(*) FROM final_artifacts`).Scan(&n)
	return n, err
}