Added automatic pruning for configured Postgres sinks

sinks/postgres.go

@@ -22,6 +22,10 @@ type postgresTx interface {
 	Rollback() error
 }
 
+type postgresExecer interface {
+	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
+}
+
 type postgresDB interface {
 	PingContext(ctx context.Context) error
 	BeginTx(ctx context.Context, opts *sql.TxOptions) (postgresTx, error)
@@ -78,9 +82,10 @@ var openPostgresDB = func(dsn string) (postgresDB, error) {
 }
 
 type PostgresSink struct {
-	name   string
-	db     postgresDB
-	schema postgresSchemaCompiled
+	name        string
+	db          postgresDB
+	schema      postgresSchemaCompiled
+	pruneWindow time.Duration
 }
 
 func NewPostgresSinkFromConfig(cfg config.SinkConfig) (Sink, error) {
@@ -96,6 +101,10 @@ func NewPostgresSinkFromConfig(cfg config.SinkConfig) (Sink, error) {
 	if err != nil {
 		return nil, err
 	}
+	pruneWindow, err := parsePostgresPruneWindow(cfg)
+	if err != nil {
+		return nil, err
+	}
 
 	schema, ok := lookupPostgresSchema(cfg.Name)
 	if !ok {
@@ -112,7 +121,7 @@ func NewPostgresSinkFromConfig(cfg config.SinkConfig) (Sink, error) {
 		return nil, fmt.Errorf("postgres sink %q: open db: %w", cfg.Name, err)
 	}
 
-	s := &PostgresSink{name: cfg.Name, db: db, schema: schema}
+	s := &PostgresSink{name: cfg.Name, db: db, schema: schema, pruneWindow: pruneWindow}
 	if err := s.initialize(); err != nil {
 		_ = db.Close()
 		return nil, err
@@ -168,6 +177,15 @@ func (p *PostgresSink) Consume(ctx context.Context, e event.Event) error {
 			return fmt.Errorf("postgres sink: insert into %q: %w", tbl.name, err)
 		}
 	}
+	if p.pruneWindow > 0 {
+		cutoff := time.Now().UTC().Add(-p.pruneWindow)
+		for _, tableName := range p.schema.tableOrder {
+			tbl := p.schema.tables[tableName]
+			if _, err := execPruneOlderThan(ctx, tx, tbl, cutoff); err != nil {
+				return fmt.Errorf("postgres sink: %w", err)
+			}
+		}
+	}
 
 	if err := tx.Commit(); err != nil {
 		return fmt.Errorf("postgres sink: commit tx: %w", err)
@@ -214,19 +232,9 @@ func (p *PostgresSink) PruneOlderThan(ctx context.Context, table string, cutoff
 		return 0, err
 	}
 
-	query := fmt.Sprintf(
-		`DELETE FROM %s WHERE %s < $1`,
-		quotePostgresIdent(tbl.name),
-		quotePostgresIdent(tbl.pruneColumn),
-	)
-
-	res, err := p.db.ExecContext(ctx, query, cutoff)
+	rows, err := execPruneOlderThan(ctx, p.db, tbl, cutoff)
 	if err != nil {
-		return 0, fmt.Errorf("postgres sink: prune older than table %q: %w", tbl.name, err)
-	}
-	rows, err := res.RowsAffected()
-	if err != nil {
-		return 0, fmt.Errorf("postgres sink: prune older than table %q rows affected: %w", tbl.name, err)
+		return 0, fmt.Errorf("postgres sink: %w", err)
 	}
 	return rows, nil
 }
@@ -309,6 +317,77 @@ func buildPostgresDSN(uri, username, password string) (string, error) {
 	return u.String(), nil
 }
 
+func parsePostgresPruneWindow(cfg config.SinkConfig) (time.Duration, error) {
+	raw, ok := cfg.Params["prune"]
+	if !ok || raw == nil {
+		return 0, nil
+	}
+
+	s, ok := raw.(string)
+	if !ok {
+		return 0, fmt.Errorf("sink %q: params.prune must be a string duration (e.g. 72h, 3d, 2w)", cfg.Name)
+	}
+
+	d, err := parsePostgresPruneDuration(s)
+	if err != nil {
+		return 0, fmt.Errorf("sink %q: params.prune %q is invalid: %w", cfg.Name, s, err)
+	}
+	return d, nil
+}
+
+func parsePostgresPruneDuration(raw string) (time.Duration, error) {
+	s := strings.TrimSpace(raw)
+	if s == "" {
+		return 0, fmt.Errorf("must not be empty")
+	}
+
+	lower := strings.ToLower(s)
+	if strings.HasSuffix(lower, "d") || strings.HasSuffix(lower, "w") {
+		unit := lower[len(lower)-1]
+		n, err := strconv.Atoi(strings.TrimSpace(lower[:len(lower)-1]))
+		if err != nil {
+			return 0, fmt.Errorf("must use a positive integer before %q", string(unit))
+		}
+		if n <= 0 {
+			return 0, fmt.Errorf("must be > 0")
+		}
+		if unit == 'd' {
+			return time.Duration(n) * 24 * time.Hour, nil
+		}
+		return time.Duration(n) * 7 * 24 * time.Hour, nil
+	}
+
+	d, err := time.ParseDuration(s)
+	if err != nil {
+		return 0, fmt.Errorf("must be a Go duration or use d/w suffixes")
+	}
+	if d <= 0 {
+		return 0, fmt.Errorf("must be > 0")
+	}
+	return d, nil
+}
+
+func buildPruneOlderThanSQL(tbl postgresTableCompiled) string {
+	return fmt.Sprintf(
+		`DELETE FROM %s WHERE %s < $1`,
+		quotePostgresIdent(tbl.name),
+		quotePostgresIdent(tbl.pruneColumn),
+	)
+}
+
+func execPruneOlderThan(ctx context.Context, execer postgresExecer, tbl postgresTableCompiled, cutoff time.Time) (int64, error) {
+	query := buildPruneOlderThanSQL(tbl)
+	res, err := execer.ExecContext(ctx, query, cutoff)
+	if err != nil {
+		return 0, fmt.Errorf("prune older than table %q: %w", tbl.name, err)
+	}
+	rows, err := res.RowsAffected()
+	if err != nil {
+		return 0, fmt.Errorf("prune older than table %q rows affected: %w", tbl.name, err)
+	}
+	return rows, nil
+}
+
 func buildCreateTableSQL(tbl postgresTableCompiled) string {
 	defs := make([]string, 0, len(tbl.columnOrder)+1)
 	for _, colName := range tbl.columnOrder {