broad-announce/internal/archiver/archiver.go

// Package archiver is the M7 periodic data-tier job. It
// moves old rows from TimescaleDB `deliveries` to ClickHouse
// `deliveries_archive` in batches. The hot-path services
// (ingestd, routerd, deliverd-*) are unaffected; they
// INSERT into `deliveries` just as before.
//
// The archiver is fail-stop: a partial batch is rolled back
// (no DELETE happens if the ClickHouse INSERT fails).
// Container restart retries the next hour.
//
// Concurrency: only one archiverd should run at a time per
// environment. We use a Postgres advisory lock so two
// replicas can't double-archive. `SKIP LOCKED` in the
// SELECT makes the lock acquisition non-blocking.
package archiver

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"time"

	_ "github.com/jackc/pgx/v5/stdlib"
)

// RunOnce executes one full pass: for each table in scope,
// drain rows older than `olderThan` in batches of
// `batchSize`. Returns the total rows archived, total
// batches, and any error encountered (caller decides
// whether to retry the next interval).
//
// The function is safe to call repeatedly; each call
// acquires the advisory lock, processes whatever's old,
// releases the lock, and returns.
func RunOnce(ctx context.Context, opts RunOptions) (Report, error) {
	rep := Report{StartedAt: time.Now().UTC()}

	conn, err := sql.Open("pgx", opts.PostgresDSN)
	if err != nil {
		return rep, fmt.Errorf("postgres open: %w", err)
	}
	defer conn.Close()
	if err := conn.PingContext(ctx); err != nil {
		return rep, fmt.Errorf("postgres ping: %w", err)
	}

	if err := ensureCHSchema(ctx, opts.ClickHouseURL, opts.Logger); err != nil {
		return rep, fmt.Errorf("clickhouse ensure-schema: %w", err)
	}

	// Per-table drain loop. We process the same table
	// repeatedly until the SELECT returns fewer than
	// `batchSize` rows.
	for _, t := range []string{"deliveries"} {
		n, batches, err := drainTable(ctx, conn, opts, t)
		if err != nil {
			return rep, fmt.Errorf("%s: %w", t, err)
		}
		rep.Tables = append(rep.Tables, TableReport{Name: t, Rows: n, Batches: batches})
	}

	rep.FinishedAt = time.Now().UTC()
	rep.Duration = rep.FinishedAt.Sub(rep.StartedAt)
	return rep, nil
}

// RunOptions is the per-call configuration.
type RunOptions struct {
	// PostgresDSN is the libpq-style DSN (e.g. from
	// config.Common.PostgresDSN).
	PostgresDSN string
	// ClickHouseURL is the HTTP base URL for the
	// ClickHouse server (e.g. http://clickhouse:8123).
	// Note: NO trailing slash.
	ClickHouseURL string
	// OlderThan is the cutoff; rows whose created_at
	// is strictly less than now() - OlderThan are
	// eligible for archive. Default: 7 days.
	OlderThan time.Duration
	// BatchSize is the cap per SELECT/INSERT. The
	// per-run cap is BatchSize * ~10 cycles (the
	// drain loop runs until a SELECT returns <BatchSize
	// rows; the safety cap is 100 cycles).
	BatchSize int
	// Logger is the slog handle.
	Logger *slog.Logger
}

// Report is the result of one RunOnce call.
type Report struct {
	StartedAt  time.Time
	FinishedAt time.Time
	Duration   time.Duration
	Tables     []TableReport
}

// TableReport is the per-table result.
type TableReport struct {
	Name    string
	Rows    int
	Batches int
}

// drainTable moves all rows of `table` older than the
// cutoff in batches of `BatchSize`. Returns the total
// row count and the number of batches.
//
// The drain uses `FOR UPDATE SKIP LOCKED` to be safe
// against concurrent archiverd instances (the advisory
// lock is the primary guard; SKIP LOCKED is a
// belt-and-suspenders).
func drainTable(ctx context.Context, conn *sql.DB, opts RunOptions, table string) (int, int, error) {
	olderThan := opts.OlderThan
	if olderThan == 0 {
		olderThan = 7 * 24 * time.Hour
	}
	batchSize := opts.BatchSize
	if batchSize == 0 {
		batchSize = 10000
	}

	// Acquire the per-table advisory lock. If another
	// archiverd holds it, we return early (no error —
	// this is normal in a multi-replica deploy).
	lockKey := int64(0xBA21B0DA) // arbitrary stable key
	conn2, err := conn.Conn(ctx)
	if err != nil {
		return 0, 0, err
	}
	defer conn2.Close()
	var gotLock bool
	if err := conn2.QueryRowContext(ctx,
		`SELECT pg_try_advisory_lock($1)`, lockKey,
	).Scan(&gotLock); err != nil {
		return 0, 0, fmt.Errorf("advisory lock: %w", err)
	}
	if !gotLock {
		opts.Logger.Info("another archiverd holds the lock; skipping this run")
		return 0, 0, nil
	}
	defer func() {
		_, _ = conn2.ExecContext(ctx, `SELECT pg_advisory_unlock($1)`, lockKey)
	}()

	totalRows := 0
	batches := 0
	safetyCycles := 100
	for cycle := 0; cycle < safetyCycles; cycle++ {
		rows, err := selectBatch(ctx, conn2, table, olderThan, batchSize)
		if err != nil {
			return totalRows, batches, err
		}
		if len(rows) == 0 {
			break
		}
		if err := insertCH(ctx, opts.ClickHouseURL, table, rows); err != nil {
			return totalRows, batches, fmt.Errorf("clickhouse insert: %w", err)
		}
		if err := deleteBatch(ctx, conn2, table, rows); err != nil {
			return totalRows, batches, fmt.Errorf("postgres delete: %w", err)
		}
		totalRows += len(rows)
		batches++
		opts.Logger.Info("archived batch",
			"table", table,
			"batch", batches,
			"rows", len(rows),
			"total_rows", totalRows,
		)
		if len(rows) < batchSize {
			break // drained
		}
	}
	return totalRows, batches, nil
}

// deliveryRow is the wire shape we read from Postgres
// and INSERT into ClickHouse. JSON tags match the CH
// column names (case-insensitive in CH).
//
// Time fields are stored as strings (CH wire format) —
// not as time.Time. We do the conversion in selectBatch
// after scanning. This sidesteps the default time.Time
// JSON encoding (RFC3339Nano with 'T' + 'Z') that CH
// can't parse for DateTime64 columns.
type deliveryRow struct {
	ID            int64           `json:"id"`
	AlertID       string          `json:"alert_id"`
	CompanyID     string          `json:"company_id"`
	IndividualID  string          `json:"individual_id"`
	Channel       string          `json:"channel"`
	Target        string          `json:"target"`
	Status        string          `json:"status"`
	Attempts      int32           `json:"attempts"`
	LastError     string          `json:"last_error"`
	Payload       json.RawMessage `json:"payload"`
	CreatedAt     string          `json:"created_at"`
	SentAt        string          `json:"sent_at,omitempty"`
	NextAttemptAt string          `json:"next_attempt_at,omitempty"`
}

// chTime formats a time as ClickHouse's preferred
// DateTime64(3, 'UTC') wire format: "2026-06-06 21:11:52.036".
// Postgres returns RFC3339Nano with 'T' separator and a 'Z'
// suffix; CH doesn't parse those directly. We use 3-decimal
// precision to match the column type.
func chTime(t time.Time) string {
	return t.UTC().Format("2006-01-02 15:04:05.000")
}

// chTimeOrEmpty renders a nullable time as "" for CH's
// Nullable(DateTime64) column. CH accepts an empty string
// for nullable datetime columns in JSONEachRow.
func chTimeOrEmpty(t *time.Time) string {
	if t == nil {
		return ""
	}
	return chTime(*t)
}

// selectBatch reads up to `limit` rows from the given
// table where the time column is older than `cutoff`.
// The time column is always `created_at` in M7; if M7.5
// adds another table, we may need to vary the column.
func selectBatch(ctx context.Context, conn *sql.Conn, table string, olderThan time.Duration, limit int) ([]deliveryRow, error) {
	cutoff := time.Now().UTC().Add(-olderThan)
	q := `
		SELECT id, alert_id, company_id, individual_id, channel,
		       target, status, attempts, last_error, payload,
		       created_at, sent_at, next_attempt_at
		FROM deliveries
		WHERE created_at < $1
		ORDER BY created_at
		LIMIT $2
		FOR UPDATE SKIP LOCKED
	`
	_ = table // currently only one table in scope
	rows, err := conn.QueryContext(ctx, q, cutoff, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []deliveryRow
	for rows.Next() {
		var r deliveryRow
		var payload sql.NullString
		var sentAt, nextAttemptAt sql.NullTime
		var createdAt time.Time
		if err := rows.Scan(
			&r.ID, &r.AlertID, &r.CompanyID, &r.IndividualID, &r.Channel,
			&r.Target, &r.Status, &r.Attempts, &r.LastError, &payload,
			&createdAt, &sentAt, &nextAttemptAt,
		); err != nil {
			return nil, err
		}
		if payload.Valid {
			r.Payload = json.RawMessage(payload.String)
		} else {
			r.Payload = json.RawMessage("null")
		}
		// Convert Postgres time.Time to CH wire format.
		// We marshal as a string in the row, not via
		// the default time.Time JSON (which is RFC3339Nano
		// with Z suffix; CH can't parse that).
		r.CreatedAt = chTime(createdAt)
		if sentAt.Valid {
			t := sentAt.Time
			r.SentAt = chTimeOrEmpty(&t)
		}
		if nextAttemptAt.Valid {
			t := nextAttemptAt.Time
			r.NextAttemptAt = chTimeOrEmpty(&t)
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

// insertCH writes the rows to ClickHouse via the HTTP
// interface. Each batch is one INSERT; CH is atomic per
// INSERT so a partial failure rolls back the whole
// batch.
func insertCH(ctx context.Context, chURL, table string, rows []deliveryRow) error {
	// CH table for the deliveries archive.
	chTable := "ba_archive.deliveries_archive"
	_ = table // currently only one in scope

	// Build a JSONEachRow payload. CH accepts one JSON
	// object per line.
	var buf bytes.Buffer
	for _, r := range rows {
		// CH's DateTime64(3, 'UTC') expects RFC3339Nano
		// format. time.Time's default MarshalJSON gives
		// RFC3339Nano.
		j, err := json.Marshal(r)
		if err != nil {
			return fmt.Errorf("marshal row %d: %w", r.ID, err)
		}
		buf.Write(j)
		buf.WriteByte('\n')
	}

	u := chURL + "/?" + url.Values{
		"query": {"INSERT INTO " + chTable + " FORMAT JSONEachRow"},
	}.Encode()
	req, err := http.NewRequestWithContext(ctx, "POST", u, &buf)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-ndjson")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("clickhouse HTTP %d: %s", resp.StatusCode, string(body))
	}
	return nil
}

// deleteBatch removes the just-archived rows from
// Postgres. The `created_at < cutoff` is the same
// predicate we used in selectBatch, but we constrain on
// `id IN (...)` to avoid accidentally re-deleting rows
// that arrived between the SELECT and the DELETE.
func deleteBatch(ctx context.Context, conn *sql.Conn, table string, rows []deliveryRow) error {
	_ = table
	ids := make([]int64, 0, len(rows))
	for _, r := range rows {
		ids = append(ids, r.ID)
	}
	// Cap chunk size at 1000 ids per query to stay
	// within pgx's parameter limit and not blow up
	// statement parsing on large batches.
	const chunk = 1000
	for i := 0; i < len(ids); i += chunk {
		end := i + chunk
		if end > len(ids) {
			end = len(ids)
		}
		q := `DELETE FROM deliveries WHERE id = ANY($1::bigint[])`
		_, err := conn.ExecContext(ctx, q, ids[i:end])
		if err != nil {
			return err
		}
	}
	return nil
}

// ensureCHSchema runs the 007_clickhouse.up.sql DDL on the
// ClickHouse server. Idempotent — every statement is
// CREATE ... IF NOT EXISTS. We do this at the start of
// every RunOnce so a fresh deploy auto-creates the
// schema without operator intervention.
func ensureCHSchema(ctx context.Context, chURL string, logger *slog.Logger) error {
	stmts := []string{
		`CREATE DATABASE IF NOT EXISTS ba_archive`,
		`CREATE TABLE IF NOT EXISTS ba_archive.deliveries_archive (
			id              BIGINT,
			alert_id        String,
			company_id      String,
			individual_id   String,
			channel         LowCardinality(String),
			target          String,
			status          LowCardinality(String),
			attempts        UInt32,
			last_error      String,
			payload         String,
			created_at      DateTime64(3, 'UTC'),
			sent_at         Nullable(DateTime64(3, 'UTC')),
			next_attempt_at Nullable(DateTime64(3, 'UTC')),
			archived_at     DateTime DEFAULT now()
		) ENGINE = MergeTree
		  PARTITION BY toYYYYMM(created_at)
		  ORDER BY (company_id, created_at, id)
		  TTL toDateTime(created_at) + INTERVAL 365 DAY`,
		`CREATE MATERIALIZED VIEW IF NOT EXISTS
			ba_archive.deliveries_per_company_daily_mv
		ENGINE = SummingMergeTree
		  PARTITION BY toYYYYMM(day)
		  ORDER BY (company_id, day, channel, status)
		AS
		SELECT
			company_id,
			toDate(created_at) AS day,
			channel,
			status,
			count()           AS n
		FROM ba_archive.deliveries_archive
		GROUP BY company_id, day, channel, status`,
	}
	for _, s := range stmts {
		req, err := http.NewRequestWithContext(ctx, "POST", chURL+"/", bytes.NewBufferString(s))
		if err != nil {
			return err
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return fmt.Errorf("ch stmt: %w", err)
		}
		if resp.StatusCode != 200 {
			body, _ := io.ReadAll(resp.Body)
			resp.Body.Close()
			return fmt.Errorf("ch stmt HTTP %d: %s", resp.StatusCode, string(body))
		}
		resp.Body.Close()
		logger.Debug("clickhouse schema ok", "stmt_prefix", firstLine(s))
	}
	return nil
}

// firstLine returns up to 80 chars of the first line of
// a SQL statement, for log readability.
func firstLine(s string) string {
	for i, c := range s {
		if c == '\n' {
			if i > 80 {
				return s[:80] + "..."
			}
			return s[:i]
		}
	}
	if len(s) > 80 {
		return s[:80] + "..."
	}
	return s
}