// 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< 500 { f.Limit = 100 } rows, err := s.pool.Query(ctx, ` SELECT id, provider, service, hostname, provider_hostname, resource_identifier, region, public_semantic, exposure_state, exposure_detail, first_seen, last_seen, probed_at FROM cloud_assets WHERE ($1 = '' OR provider = $1) AND ($2 = '' OR service = $2) AND ($3 = '' OR exposure_state = $3) AND ($4 = '' OR provider_hostname ILIKE '%' || $4 || '%' OR hostname ILIKE '%' || $4 || '%' OR resource_identifier ILIKE '%' || $4 || '%') ORDER BY last_seen DESC LIMIT $5 OFFSET $6`, f.Provider, f.Service, f.Exposure, f.Query, f.Limit, f.Offset) if err != nil { return nil, err } defer rows.Close() out := []Asset{} for rows.Next() { var a Asset if err := rows.Scan(&a.ID, &a.Provider, &a.Service, &a.Hostname, &a.ProviderHost, &a.Resource, &a.Region, &a.PublicSemantic, &a.Exposure, &a.ExposureDetail, &a.FirstSeen, &a.LastSeen, &a.ProbedAt); err != nil { return nil, err } out = append(out, a) } return out, rows.Err() } // Count is a (label, count) pair for the UI summary. type Count struct { Label string `json:"label"` Count int64 `json:"count"` } // Stats summarizes the assets table. type Stats struct { Total int64 `json:"total"` Readable int64 `json:"readable"` ByProvider []Count `json:"by_provider"` ByService []Count `json:"by_service"` ByExposure []Count `json:"by_exposure"` } func (s *Store) Stats(ctx context.Context) (Stats, error) { var st Stats if err := s.pool.QueryRow(ctx, `SELECT count(*), count(*) FILTER (WHERE exposure_state = 'public_read') FROM cloud_assets`).Scan(&st.Total, &st.Readable); err != nil { return st, err } group := func(column string) ([]Count, error) { rows, err := s.pool.Query(ctx, `SELECT `+column+`, count(*) FROM cloud_assets GROUP BY 1 ORDER BY 2 DESC`) if err != nil { return nil, err } defer rows.Close() var counts []Count for rows.Next() { var c Count if err := rows.Scan(&c.Label, &c.Count); err != nil { return nil, err } counts = append(counts, c) } return counts, rows.Err() } var err error if st.ByProvider, err = group("provider"); err != nil { return st, err } if st.ByService, err = group("service"); err != nil { return st, err } if st.ByExposure, err = group("exposure_state"); err != nil { return st, err } return st, nil }