// 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<