GitHop/services/control-plane/internal/store/store.go

package store

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"sort"
	"time"

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

	"hopweb/internal/gitea"
	"hopweb/internal/ids"
)

//go:embed migrations/*.sql
var migrations embed.FS

type Repository struct {
	ID            string    `json:"id"`
	GiteaID       int64     `json:"gitea_id"`
	Owner         string    `json:"owner"`
	Name          string    `json:"name"`
	FullName      string    `json:"full_name"`
	CloneURL      string    `json:"clone_url"`
	DefaultBranch string    `json:"default_branch"`
	CreatedAt     time.Time `json:"created_at"`
	UpdatedAt     time.Time `json:"updated_at"`
}

type WebhookDelivery struct {
	DeliveryID    string
	Event         string
	EventType     string
	PayloadSHA256 string
	Repository    *gitea.Repository
}

type Prompt struct {
	ID              string          `json:"id"`
	GiteaUserID     int64           `json:"-"`
	TaskID          string          `json:"task_id,omitempty"`
	AttemptID       string          `json:"attempt_id,omitempty"`
	StateID         string          `json:"state_id,omitempty"`
	Prompt          string          `json:"prompt"`
	AgentName       string          `json:"agent_name,omitempty"`
	AgentModel      string          `json:"agent_model,omitempty"`
	Status          string          `json:"status"`
	ResponseSummary string          `json:"response_summary,omitempty"`
	Metadata        json.RawMessage `json:"metadata"`
	CreatedAt       time.Time       `json:"created_at"`
	CompletedAt     *time.Time      `json:"completed_at,omitempty"`
}

type CreatePromptInput struct {
	Owner           string
	Repository      string
	GiteaUserID     int64
	TaskID          string
	AttemptID       string
	StateID         string
	Prompt          string
	AgentName       string
	AgentModel      string
	Status          string
	ResponseSummary string
	Metadata        json.RawMessage
	CreatedAt       *time.Time
	CompletedAt     *time.Time
}

type Store interface {
	Ping(context.Context) error
	ListRepositories(context.Context) ([]Repository, error)
	UpsertRepository(context.Context, gitea.Repository) (Repository, error)
	ListPrompts(context.Context, string, string, int64, int) ([]Prompt, error)
	CreatePrompt(context.Context, CreatePromptInput) (Prompt, error)
	RecordWebhook(context.Context, WebhookDelivery) (bool, error)
}

type Postgres struct {
	pool *pgxpool.Pool
}

func Open(ctx context.Context, databaseURL string) (*Postgres, error) {
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		return nil, err
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, err
	}
	return &Postgres{pool: pool}, nil
}

func (p *Postgres) Close() { p.pool.Close() }

func (p *Postgres) Ping(ctx context.Context) error { return p.pool.Ping(ctx) }

func (p *Postgres) Migrate(ctx context.Context) error {
	if _, err := p.pool.Exec(ctx, `CREATE TABLE IF NOT EXISTS schema_migrations (name text PRIMARY KEY, applied_at timestamptz NOT NULL DEFAULT now())`); err != nil {
		return err
	}
	entries, err := migrations.ReadDir("migrations")
	if err != nil {
		return err
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].Name() < entries[j].Name() })
	for _, entry := range entries {
		name := entry.Name()
		var exists bool
		if err := p.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM schema_migrations WHERE name=$1)`, name).Scan(&exists); err != nil {
			return err
		}
		if exists {
			continue
		}
		sql, err := migrations.ReadFile("migrations/" + name)
		if err != nil {
			return err
		}
		tx, err := p.pool.Begin(ctx)
		if err != nil {
			return err
		}
		if _, err = tx.Exec(ctx, string(sql)); err == nil {
			_, err = tx.Exec(ctx, `INSERT INTO schema_migrations (name) VALUES ($1)`, name)
		}
		if err != nil {
			_ = tx.Rollback(ctx)
			return fmt.Errorf("apply migration %s: %w", name, err)
		}
		if err := tx.Commit(ctx); err != nil {
			return err
		}
	}
	return nil
}

func (p *Postgres) ListRepositories(ctx context.Context) ([]Repository, error) {
	rows, err := p.pool.Query(ctx, `SELECT id, gitea_id, owner, name, full_name, clone_url, default_branch, created_at, updated_at FROM repositories ORDER BY full_name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	repositories := make([]Repository, 0)
	for rows.Next() {
		var repository Repository
		if err := rows.Scan(&repository.ID, &repository.GiteaID, &repository.Owner, &repository.Name, &repository.FullName, &repository.CloneURL, &repository.DefaultBranch, &repository.CreatedAt, &repository.UpdatedAt); err != nil {
			return nil, err
		}
		repositories = append(repositories, repository)
	}
	return repositories, rows.Err()
}

func (p *Postgres) UpsertRepository(ctx context.Context, repository gitea.Repository) (Repository, error) {
	id, err := ids.Repository()
	if err != nil {
		return Repository{}, err
	}
	return upsertRepository(ctx, p.pool, id, repository)
}

func (p *Postgres) ListPrompts(ctx context.Context, owner, name string, giteaUserID int64, limit int) ([]Prompt, error) {
	rows, err := p.pool.Query(ctx, `
		SELECT p.id, p.gitea_user_id, p.task_id, p.attempt_id, p.state_id, p.prompt, p.agent_name, p.agent_model,
			p.status, p.response_summary, p.metadata, p.created_at, p.completed_at
		FROM prompts p
		JOIN repositories r ON r.id = p.repository_id
		WHERE r.owner=$1 AND r.name=$2 AND p.gitea_user_id=$3
		ORDER BY p.created_at DESC
		LIMIT $4`, owner, name, giteaUserID, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	prompts := make([]Prompt, 0)
	for rows.Next() {
		var prompt Prompt
		if err := rows.Scan(&prompt.ID, &prompt.GiteaUserID, &prompt.TaskID, &prompt.AttemptID, &prompt.StateID, &prompt.Prompt, &prompt.AgentName, &prompt.AgentModel, &prompt.Status, &prompt.ResponseSummary, &prompt.Metadata, &prompt.CreatedAt, &prompt.CompletedAt); err != nil {
			return nil, err
		}
		prompts = append(prompts, prompt)
	}
	return prompts, rows.Err()
}

func (p *Postgres) CreatePrompt(ctx context.Context, input CreatePromptInput) (Prompt, error) {
	id, err := ids.Prompt()
	if err != nil {
		return Prompt{}, err
	}
	metadata := input.Metadata
	if len(metadata) == 0 {
		metadata = json.RawMessage(`{}`)
	}
	var prompt Prompt
	err = p.pool.QueryRow(ctx, `
		INSERT INTO prompts (id, repository_id, gitea_user_id, task_id, attempt_id, state_id, prompt, agent_name, agent_model, status, response_summary, metadata, created_at, completed_at)
		SELECT $1, r.id, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, COALESCE($14, now()), $15
		FROM repositories r WHERE r.owner=$2 AND r.name=$3
		ON CONFLICT (repository_id, gitea_user_id, state_id)
			WHERE gitea_user_id IS NOT NULL AND state_id <> ''
		DO UPDATE SET
			task_id=EXCLUDED.task_id, attempt_id=EXCLUDED.attempt_id,
			prompt=EXCLUDED.prompt, agent_name=EXCLUDED.agent_name, agent_model=EXCLUDED.agent_model,
			status=EXCLUDED.status, response_summary=EXCLUDED.response_summary, metadata=EXCLUDED.metadata,
			created_at=EXCLUDED.created_at, completed_at=EXCLUDED.completed_at
		RETURNING id, gitea_user_id, task_id, attempt_id, state_id, prompt, agent_name, agent_model, status, response_summary, metadata, created_at, completed_at`,
		id, input.Owner, input.Repository, input.GiteaUserID, input.TaskID, input.AttemptID, input.StateID, input.Prompt, input.AgentName, input.AgentModel, input.Status, input.ResponseSummary, metadata, input.CreatedAt, input.CompletedAt,
	).Scan(&prompt.ID, &prompt.GiteaUserID, &prompt.TaskID, &prompt.AttemptID, &prompt.StateID, &prompt.Prompt, &prompt.AgentName, &prompt.AgentModel, &prompt.Status, &prompt.ResponseSummary, &prompt.Metadata, &prompt.CreatedAt, &prompt.CompletedAt)
	if err != nil {
		return Prompt{}, err
	}
	return prompt, nil
}

type querier interface {
	QueryRow(context.Context, string, ...any) pgx.Row
}

func upsertRepository(ctx context.Context, q querier, id string, repository gitea.Repository) (Repository, error) {
	owner := repository.OwnerName()
	fullName := repository.FullName
	if fullName == "" {
		fullName = owner + "/" + repository.Name
	}
	var saved Repository
	err := q.QueryRow(ctx, `
		INSERT INTO repositories (id, gitea_id, owner, name, full_name, clone_url, default_branch)
		VALUES ($1, $2, $3, $4, $5, $6, $7)
		ON CONFLICT (gitea_id) DO UPDATE SET
			owner=EXCLUDED.owner, name=EXCLUDED.name, full_name=EXCLUDED.full_name,
			clone_url=EXCLUDED.clone_url, default_branch=EXCLUDED.default_branch, updated_at=now()
		RETURNING id, gitea_id, owner, name, full_name, clone_url, default_branch, created_at, updated_at`,
		id, repository.ID, owner, repository.Name, fullName, repository.CloneURL, repository.DefaultBranch,
	).Scan(&saved.ID, &saved.GiteaID, &saved.Owner, &saved.Name, &saved.FullName, &saved.CloneURL, &saved.DefaultBranch, &saved.CreatedAt, &saved.UpdatedAt)
	return saved, err
}

func (p *Postgres) RecordWebhook(ctx context.Context, delivery WebhookDelivery) (bool, error) {
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return false, err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	result, err := tx.Exec(ctx, `INSERT INTO webhook_deliveries (delivery_id, event, event_type, payload_sha256) VALUES ($1,$2,$3,$4) ON CONFLICT DO NOTHING`, delivery.DeliveryID, delivery.Event, delivery.EventType, delivery.PayloadSHA256)
	if err != nil {
		return false, err
	}
	if result.RowsAffected() == 0 {
		var existingEvent, existingEventType, existingDigest string
		if err := tx.QueryRow(ctx, `SELECT event, event_type, payload_sha256 FROM webhook_deliveries WHERE delivery_id=$1`, delivery.DeliveryID).Scan(&existingEvent, &existingEventType, &existingDigest); err != nil {
			return false, err
		}
		if existingEvent != delivery.Event || existingEventType != delivery.EventType || existingDigest != delivery.PayloadSHA256 {
			return false, fmt.Errorf("webhook delivery ID %q was reused with different content", delivery.DeliveryID)
		}
		return true, nil
	}
	if delivery.Repository != nil && delivery.Repository.ID > 0 {
		id, err := ids.Repository()
		if err != nil {
			return false, err
		}
		if _, err := upsertRepository(ctx, tx, id, *delivery.Repository); err != nil {
			return false, err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return false, err
	}
	return false, nil
}

var _ Store = (*Postgres)(nil)