bucketeer/internal/store/store.go

// Package store owns the PostgreSQL schema, the batched findings writer and the
// restart-safe CT checkpoint state provider.
package store

import (
	"context"
	"fmt"
	"time"

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

	"tea.karlcom.de/karl/bucketeer/internal/metrics"
	"tea.karlcom.de/karl/bucketeer/internal/model"
	"tea.karlcom.de/karl/bucketeer/internal/pipeline"
	"tea.karlcom.de/karl/bucketeer/migrations"
)

// Store wraps a pgx pool with the operations ctcloud needs.
type Store struct {
	pool    *pgxpool.Pool
	metrics *metrics.Metrics
}

// Open connects to PostgreSQL. The pool is bounded; the caller configures the
// maximum number of connections.
func Open(ctx context.Context, databaseURL string, maxConns int) (*Store, error) {
	cfg, err := pgxpool.ParseConfig(databaseURL)
	if err != nil {
		return nil, fmt.Errorf("invalid DATABASE_URL: %w", err)
	}
	if maxConns > 0 {
		cfg.MaxConns = int32(maxConns)
	}
	cfg.MaxConnLifetime = 30 * time.Minute
	cfg.MinConns = 1
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, err
	}
	return &Store{pool: pool}, nil
}

func (s *Store) SetMetrics(m *metrics.Metrics) { s.metrics = m }

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

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

func (s *Store) Pool() *pgxpool.Pool { return s.pool }

// Migrate applies the embedded schema idempotently.
func (s *Store) Migrate(ctx context.Context) error {
	for _, name := range migrations.SQL() {
		sqlBytes, err := migrations.FS.ReadFile(name)
		if err != nil {
			return err
		}
		if _, err := s.pool.Exec(ctx, string(sqlBytes)); err != nil {
			return fmt.Errorf("applying %s: %w", name, err)
		}
	}
	return nil
}

const upsertFindingSQL = `
WITH asset AS (
    INSERT INTO cloud_assets
        (provider, service, hostname, provider_hostname, resource_identifier,
         region, resource_specific, public_semantic, first_seen, last_seen)
    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, now(), now())
    ON CONFLICT (provider, service, provider_hostname)
    DO UPDATE SET
        last_seen           = now(),
        hostname            = EXCLUDED.hostname,
        resource_identifier = CASE WHEN EXCLUDED.resource_identifier = '' THEN cloud_assets.resource_identifier ELSE EXCLUDED.resource_identifier END,
        region              = CASE WHEN EXCLUDED.region = '' THEN cloud_assets.region ELSE EXCLUDED.region END
    RETURNING id
)
INSERT INTO observations
    (asset_id, source_hostname, source_type, cname_chain, ipv4, ipv6)
SELECT asset.id, $9, $10, $11, $12, $13 FROM asset`

// WriteFindings writes one batch of findings in a single round trip. Asset rows
// are upserted and an observation is appended for each finding.
func (s *Store) WriteFindings(ctx context.Context, batch []pipeline.Finding) error {
	if len(batch) == 0 {
		return nil
	}
	var lastErr error
	// ponytail: three quick attempts; the pipeline keeps running and the
	// bounded DB queue sheds load if PostgreSQL stays down.
	for attempt := 0; attempt < 3; attempt++ {
		if err := s.writeBatch(ctx, batch); err != nil {
			lastErr = err
			if ctx.Err() != nil {
				return err
			}
			backoff := time.Duration(1<<attempt) * 100 * time.Millisecond
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(backoff):
			}
			continue
		}
		return nil
	}
	return lastErr
}

func (s *Store) writeBatch(ctx context.Context, batch []pipeline.Finding) error {
	b := &pgx.Batch{}
	for i := range batch {
		f := &batch[i]
		b.Queue(upsertFindingSQL,
			f.Provider.String(), f.Service.String(), f.Original, f.Hostname,
			f.Resource, f.Region, f.ResourceSpecific, f.PublicSemantic.String(),
			f.Original, f.Source.String(),
			nonNil(f.CNAMEChain), nonNil(f.IPv4), nonNil(f.IPv6),
		)
	}
	br := s.pool.SendBatch(ctx, b)
	defer br.Close()
	for range batch {
		if _, err := br.Exec(); err != nil {
			return err
		}
	}
	return nil
}

func nonNil(s []string) []string {
	if s == nil {
		return []string{}
	}
	return s
}

// UpdateExposure records the outcome of an optional reachability probe.
func (s *Store) UpdateExposure(ctx context.Context, provider, service, providerHostname string, state model.ExposureState) error {
	_, err := s.pool.Exec(ctx,
		`UPDATE cloud_assets SET exposure_state = $4 WHERE provider = $1 AND service = $2 AND provider_hostname = $3`,
		provider, service, providerHostname, state.String())
	return err
}