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") }