gito/services/core/internal/storage/postgres.go

package storage

import (
	"context"
	"fmt"
	"strings"
	"time"

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

	"git.toki-labs.com/toki/gito/services/core/internal/core"
)

// PgStore is a pgxpool-backed Store implementation.
type PgStore struct {
	pool       *pgxpool.Pool
	watches    *pgBranchWatchStore
	cursors    *pgRevisionCursorStore
	deliveries *pgProviderDeliveryStore
}

// NewPgStore opens a pgxpool connection and optionally applies a migration.
// migrationSQL is the raw SQL content (e.g. from //go:embed); pass an empty
// string to skip migration.
func NewPgStore(ctx context.Context, dsn string, migrationSQL string) (*PgStore, error) {
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, fmt.Errorf("pgstore: open pool: %w", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("pgstore: ping: %w", err)
	}
	if migrationSQL != "" {
		if err := applyMigrationSQL(ctx, pool, migrationSQL); err != nil {
			pool.Close()
			return nil, fmt.Errorf("pgstore: migration: %w", err)
		}
	}
	s := &PgStore{pool: pool}
	s.watches = &pgBranchWatchStore{pool: pool}
	s.cursors = &pgRevisionCursorStore{pool: pool}
	s.deliveries = &pgProviderDeliveryStore{pool: pool}
	return s, nil
}

func (s *PgStore) Close() { s.pool.Close() }

func (s *PgStore) Ping(ctx context.Context) error { return s.pool.Ping(ctx) }

func (s *PgStore) Repos() RepoStore                       { return nil }
func (s *PgStore) WorkspaceLeases() WorkspaceLeaseStore   { return nil }
func (s *PgStore) Operations() OperationStore             { return nil }
func (s *PgStore) OperationEvents() OperationEventStore   { return nil }
func (s *PgStore) BranchWatches() BranchWatchStore        { return s.watches }
func (s *PgStore) RevisionCursors() RevisionCursorStore   { return s.cursors }
func (s *PgStore) ProviderDeliveries() ProviderDeliveryStore { return s.deliveries }

// applyMigrationSQL executes the Up block extracted from goose-style SQL content.
func applyMigrationSQL(ctx context.Context, pool *pgxpool.Pool, content string) error {
	sql := extractUpBlock(content)
	if sql == "" {
		return nil
	}
	_, err := pool.Exec(ctx, sql)
	return err
}

func extractUpBlock(content string) string {
	const beginMarker = "-- +goose StatementBegin"
	const endMarker = "-- +goose StatementEnd"
	const upMarker = "-- +goose Up"

	upIdx := strings.Index(content, upMarker)
	if upIdx < 0 {
		return ""
	}
	after := content[upIdx:]
	beginIdx := strings.Index(after, beginMarker)
	endIdx := strings.Index(after, endMarker)
	if beginIdx < 0 || endIdx < 0 || endIdx <= beginIdx {
		return ""
	}
	return strings.TrimSpace(after[beginIdx+len(beginMarker) : endIdx])
}

// pgBranchWatchStore

type pgBranchWatchStore struct{ pool *pgxpool.Pool }

func (s *pgBranchWatchStore) UpsertBranchWatch(ctx context.Context, watch core.BranchWatch) (core.BranchWatch, error) {
	const q = `
		INSERT INTO branch_watches (id, provider, repo_id, branch, created_at)
		VALUES ($1, $2, $3, $4, $5)
		ON CONFLICT (provider, repo_id, branch) DO UPDATE
			SET id = branch_watches.id, created_at = branch_watches.created_at
		RETURNING id, provider, repo_id, branch, created_at
	`
	var out core.BranchWatch
	row := s.pool.QueryRow(ctx, q, watch.ID, watch.Provider, watch.RepoID, watch.Branch, watch.CreatedAt)
	if err := row.Scan(&out.ID, &out.Provider, &out.RepoID, &out.Branch, &out.CreatedAt); err != nil {
		return core.BranchWatch{}, fmt.Errorf("upsert branch watch: %w", err)
	}
	return out, nil
}

func (s *pgBranchWatchStore) FindBranchWatch(ctx context.Context, provider, repoID, branch string) (core.BranchWatch, bool, error) {
	const q = `SELECT id, provider, repo_id, branch, created_at FROM branch_watches WHERE provider=$1 AND repo_id=$2 AND branch=$3`
	var out core.BranchWatch
	row := s.pool.QueryRow(ctx, q, provider, repoID, branch)
	if err := row.Scan(&out.ID, &out.Provider, &out.RepoID, &out.Branch, &out.CreatedAt); err != nil {
		if isNoRows(err) {
			return core.BranchWatch{}, false, nil
		}
		return core.BranchWatch{}, false, fmt.Errorf("find branch watch: %w", err)
	}
	return out, true, nil
}

func (s *pgBranchWatchStore) ListBranchWatches(ctx context.Context) ([]core.BranchWatch, error) {
	const q = `SELECT id, provider, repo_id, branch, created_at FROM branch_watches ORDER BY id`
	rows, err := s.pool.Query(ctx, q)
	if err != nil {
		return nil, fmt.Errorf("list branch watches: %w", err)
	}
	defer rows.Close()
	var out []core.BranchWatch
	for rows.Next() {
		var w core.BranchWatch
		if err := rows.Scan(&w.ID, &w.Provider, &w.RepoID, &w.Branch, &w.CreatedAt); err != nil {
			return nil, fmt.Errorf("scan branch watch: %w", err)
		}
		out = append(out, w)
	}
	return out, rows.Err()
}

// pgRevisionCursorStore

type pgRevisionCursorStore struct{ pool *pgxpool.Pool }

func (s *pgRevisionCursorStore) UpsertRevisionCursor(ctx context.Context, cursor core.RevisionCursor) error {
	const q = `
		INSERT INTO revision_cursors (repo_id, branch, revision, observed_at)
		VALUES ($1, $2, $3, $4)
		ON CONFLICT (repo_id, branch) DO UPDATE
			SET revision = EXCLUDED.revision, observed_at = EXCLUDED.observed_at
	`
	_, err := s.pool.Exec(ctx, q, cursor.RepoID, cursor.Branch, cursor.Revision, cursor.ObservedAt)
	if err != nil {
		return fmt.Errorf("upsert revision cursor: %w", err)
	}
	return nil
}

func (s *pgRevisionCursorStore) GetRevisionCursor(ctx context.Context, repoID, branch string) (core.RevisionCursor, bool, error) {
	const q = `SELECT repo_id, branch, revision, observed_at FROM revision_cursors WHERE repo_id=$1 AND branch=$2`
	var out core.RevisionCursor
	row := s.pool.QueryRow(ctx, q, repoID, branch)
	if err := row.Scan(&out.RepoID, &out.Branch, &out.Revision, &out.ObservedAt); err != nil {
		if isNoRows(err) {
			return core.RevisionCursor{}, false, nil
		}
		return core.RevisionCursor{}, false, fmt.Errorf("get revision cursor: %w", err)
	}
	return out, true, nil
}

// pgProviderDeliveryStore

type pgProviderDeliveryStore struct{ pool *pgxpool.Pool }

func (s *pgProviderDeliveryStore) RecordOnce(ctx context.Context, delivery core.ProviderDelivery) (DeliveryResult, error) {
	const insert = `
		INSERT INTO provider_deliveries
			(id, provider, delivery_id, dedupe_key, event_id, repo_id, branch, revision, created_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
		ON CONFLICT (provider, dedupe_key) DO NOTHING
	`
	tag, err := s.pool.Exec(ctx, insert,
		delivery.ID, delivery.Provider, delivery.DeliveryID, delivery.DedupeKey,
		delivery.EventID, delivery.RepoID, delivery.Branch, delivery.Revision,
		delivery.CreatedAt,
	)
	if err != nil {
		return DeliveryResult{}, fmt.Errorf("record delivery: %w", err)
	}
	if tag.RowsAffected() == 1 {
		return DeliveryResult{First: true}, nil
	}
	// Duplicate: fetch existing event_id and created_at for idempotent response.
	const sel = `SELECT event_id, created_at FROM provider_deliveries WHERE provider=$1 AND dedupe_key=$2`
	var existingEventID string
	var existingCreatedAt time.Time
	row := s.pool.QueryRow(ctx, sel, delivery.Provider, delivery.DedupeKey)
	if err := row.Scan(&existingEventID, &existingCreatedAt); err != nil {
		return DeliveryResult{}, fmt.Errorf("fetch existing delivery: %w", err)
	}
	return DeliveryResult{
		First:             false,
		ExistingEventID:   existingEventID,
		ExistingCreatedAt: existingCreatedAt,
	}, nil
}

func isNoRows(err error) bool {
	return err != nil && strings.Contains(err.Error(), "no rows")
}