ITWorx-Pulse-Public/internal/alert/repository.go

package alert

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/itworx/pulse/internal/metriccatalog"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
	"reflect"
)

type Store interface {
	Create(context.Context, string, Document, string) (Rule, Version, error)
	Get(context.Context, string) (Rule, error)
	List(context.Context, int) ([]Rule, error)
	Update(context.Context, string, string, int64, Document, string) (Rule, error)
	Versions(context.Context, string, int) ([]Version, error)
	SetEnabled(context.Context, string, int64, bool) (Rule, error)
}

type Repository struct {
	Pool     *pgxpool.Pool
	Registry metriccatalog.Registry
}

func (r Repository) Create(ctx context.Context, actor string, document Document, changeSummary string) (Rule, Version, error) {
	if r.Pool == nil {
		return Rule{}, Version{}, ErrUnavailable
	}
	if err := document.Validate(r.Registry); err != nil {
		return Rule{}, Version{}, err
	}
	if changeSummary == "" {
		changeSummary = "initial version"
	}
	tx, err := r.Pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return Rule{}, Version{}, fmt.Errorf("begin alert rule create: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	docJSON, err := document.MarshalCanonical()
	if err != nil {
		return Rule{}, Version{}, fmt.Errorf("marshal alert rule: %w", err)
	}
	conditionJSON, _ := json.Marshal(document.Condition)
	scopeJSON, _ := json.Marshal(nonNilMap(document.Scope))
	groupJSON, _ := json.Marshal(nonNilStrings(document.GroupBy))
	suppressJSON, _ := json.Marshal(nonNilStrings(document.SuppressWhen))
	messageJSON, _ := json.Marshal(document.Message)
	if _, err = tx.Exec(ctx, `INSERT INTO alert_rules (id,schema_version,name,enabled,severity,scope,condition,evaluation_interval_seconds,pending_seconds,resolve_seconds,cooldown_seconds,unknown_behavior,group_by,suppress_when,message,revision,created_by) VALUES ($1,$2,$3,$4,$5,$6::jsonb,$7::jsonb,$8,$9,$10,$11,$12,$13::jsonb,$14::jsonb,$15::jsonb,1,(SELECT id FROM users WHERE external_subject=$16))`, document.ID, document.SchemaVersion, document.Name, document.Enabled, document.Severity, scopeJSON, conditionJSON, document.EvaluationIntervalSeconds, document.PendingSeconds, document.ResolveSeconds, document.CooldownSeconds, document.UnknownBehavior, groupJSON, suppressJSON, messageJSON, actor); err != nil {
		return Rule{}, Version{}, mapError(fmt.Errorf("create alert rule: %w", err))
	}
	versionID := NewID()
	if _, err = tx.Exec(ctx, `INSERT INTO alert_rule_versions (id,rule_id,version_number,document,change_summary,created_by) VALUES ($1,$2,1,$3::jsonb,$4,(SELECT id FROM users WHERE external_subject=$5))`, versionID, document.ID, docJSON, changeSummary, actor); err != nil {
		return Rule{}, Version{}, mapError(fmt.Errorf("create alert rule version: %w", err))
	}
	if _, err = tx.Exec(ctx, `UPDATE alert_rules SET current_version_id=$1 WHERE id=$2`, versionID, document.ID); err != nil {
		return Rule{}, Version{}, fmt.Errorf("set current alert rule version: %w", err)
	}
	if err = tx.Commit(ctx); err != nil {
		return Rule{}, Version{}, fmt.Errorf("commit alert rule create: %w", err)
	}
	rule, err := r.Get(ctx, document.ID)
	if err != nil {
		return Rule{}, Version{}, err
	}
	versions, err := r.Versions(ctx, document.ID, 1)
	if err != nil || len(versions) == 0 {
		return Rule{}, Version{}, err
	}
	return rule, versions[0], nil
}

func (r Repository) Get(ctx context.Context, id string) (Rule, error) {
	if r.Pool == nil {
		return Rule{}, ErrUnavailable
	}
	var rule Rule
	var scopeJSON, conditionJSON, groupJSON, suppressJSON, messageJSON []byte
	var createdBy *string
	err := r.Pool.QueryRow(ctx, `SELECT r.id,r.schema_version,r.name,r.enabled,r.severity,r.scope,r.condition,r.evaluation_interval_seconds,r.pending_seconds,r.resolve_seconds,r.cooldown_seconds,r.unknown_behavior,r.group_by,r.suppress_when,r.message,r.revision,v.version_number,COALESCE(u.external_subject,''),r.created_at,r.updated_at FROM alert_rules r JOIN alert_rule_versions v ON v.id=r.current_version_id LEFT JOIN users u ON u.id=r.created_by WHERE r.id=$1`, id).Scan(&rule.ID, &rule.SchemaVersion, &rule.Name, &rule.Enabled, &rule.Severity, &scopeJSON, &conditionJSON, &rule.EvaluationIntervalSeconds, &rule.PendingSeconds, &rule.ResolveSeconds, &rule.CooldownSeconds, &rule.UnknownBehavior, &groupJSON, &suppressJSON, &messageJSON, &rule.Revision, &rule.CurrentVersion, &createdBy, &rule.CreatedAt, &rule.UpdatedAt)
	if errors.Is(err, pgx.ErrNoRows) {
		return Rule{}, ErrNotFound
	}
	if err != nil {
		return Rule{}, fmt.Errorf("get alert rule: %w", err)
	}
	rule.CreatedBy = valueOrEmpty(createdBy)
	if err := decodeStored(&rule.Document, scopeJSON, conditionJSON, groupJSON, suppressJSON, messageJSON, rule.CooldownSeconds); err != nil {
		return Rule{}, err
	}
	return rule, nil
}

func (r Repository) List(ctx context.Context, limit int) ([]Rule, error) {
	if r.Pool == nil {
		return nil, ErrUnavailable
	}
	if limit < 1 || limit > 100 {
		return nil, errors.New("alert rule limit is invalid")
	}
	rows, err := r.Pool.Query(ctx, `SELECT r.id,r.schema_version,r.name,r.enabled,r.severity,r.scope,r.condition,r.evaluation_interval_seconds,r.pending_seconds,r.resolve_seconds,r.cooldown_seconds,r.unknown_behavior,r.group_by,r.suppress_when,r.message,r.revision,v.version_number,COALESCE(u.external_subject,''),r.created_at,r.updated_at FROM alert_rules r JOIN alert_rule_versions v ON v.id=r.current_version_id LEFT JOIN users u ON u.id=r.created_by ORDER BY r.name ASC,r.id ASC LIMIT $1`, limit)
	if err != nil {
		return nil, fmt.Errorf("list alert rules: %w", err)
	}
	defer rows.Close()
	result := make([]Rule, 0, limit)
	for rows.Next() {
		var rule Rule
		var scopeJSON, conditionJSON, groupJSON, suppressJSON, messageJSON []byte
		if err := rows.Scan(&rule.ID, &rule.SchemaVersion, &rule.Name, &rule.Enabled, &rule.Severity, &scopeJSON, &conditionJSON, &rule.EvaluationIntervalSeconds, &rule.PendingSeconds, &rule.ResolveSeconds, &rule.CooldownSeconds, &rule.UnknownBehavior, &groupJSON, &suppressJSON, &messageJSON, &rule.Revision, &rule.CurrentVersion, &rule.CreatedBy, &rule.CreatedAt, &rule.UpdatedAt); err != nil {
			return nil, fmt.Errorf("scan alert rule: %w", err)
		}
		if err := decodeStored(&rule.Document, scopeJSON, conditionJSON, groupJSON, suppressJSON, messageJSON, rule.CooldownSeconds); err != nil {
			return nil, err
		}
		result = append(result, rule)
	}
	return result, rows.Err()
}

func (r Repository) Update(ctx context.Context, id, actor string, expected int64, document Document, changeSummary string) (Rule, error) {
	if r.Pool == nil {
		return Rule{}, ErrUnavailable
	}
	document.ID = id
	if err := document.Validate(r.Registry); err != nil {
		return Rule{}, err
	}
	tx, err := r.Pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return Rule{}, fmt.Errorf("begin alert rule update: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	var currentRevision int64
	var currentVersionID string
	var currentJSON []byte
	err = tx.QueryRow(ctx, `SELECT revision,current_version_id FROM alert_rules WHERE id=$1 FOR UPDATE`, id).Scan(&currentRevision, &currentVersionID)
	if errors.Is(err, pgx.ErrNoRows) {
		return Rule{}, ErrNotFound
	}
	if err != nil {
		return Rule{}, fmt.Errorf("lock alert rule: %w", err)
	}
	if currentRevision != expected {
		return Rule{}, ErrConflict
	}
	if err = tx.QueryRow(ctx, `SELECT document FROM alert_rule_versions WHERE id=$1`, currentVersionID).Scan(&currentJSON); err != nil {
		return Rule{}, fmt.Errorf("read current alert rule version: %w", err)
	}
	nextJSON, err := document.MarshalCanonical()
	if err != nil {
		return Rule{}, err
	}
	if sameJSON(currentJSON, nextJSON) {
		if err := tx.Commit(ctx); err != nil {
			return Rule{}, err
		}
		return r.Get(ctx, id)
	}
	var currentVersion int
	if err := tx.QueryRow(ctx, `SELECT version_number FROM alert_rule_versions WHERE id=$1`, currentVersionID).Scan(&currentVersion); err != nil {
		return Rule{}, err
	}
	if changeSummary == "" {
		changeSummary = "rule update"
	}
	versionID := NewID()
	if _, err = tx.Exec(ctx, `INSERT INTO alert_rule_versions (id,rule_id,version_number,document,change_summary,created_by) VALUES ($1,$2,$3,$4::jsonb,$5,(SELECT id FROM users WHERE external_subject=$6))`, versionID, id, currentVersion+1, nextJSON, changeSummary, actor); err != nil {
		return Rule{}, mapError(err)
	}
	conditionJSON, _ := json.Marshal(document.Condition)
	scopeJSON, _ := json.Marshal(nonNilMap(document.Scope))
	groupJSON, _ := json.Marshal(nonNilStrings(document.GroupBy))
	suppressJSON, _ := json.Marshal(nonNilStrings(document.SuppressWhen))
	messageJSON, _ := json.Marshal(document.Message)
	tag, err := tx.Exec(ctx, `UPDATE alert_rules SET schema_version=$1,name=$2,enabled=$3,severity=$4,scope=$5::jsonb,condition=$6::jsonb,evaluation_interval_seconds=$7,pending_seconds=$8,resolve_seconds=$9,cooldown_seconds=$10,unknown_behavior=$11,group_by=$12::jsonb,suppress_when=$13::jsonb,message=$14::jsonb,current_version_id=$15,revision=revision+1,updated_at=now() WHERE id=$16 AND revision=$17`, document.SchemaVersion, document.Name, document.Enabled, document.Severity, scopeJSON, conditionJSON, document.EvaluationIntervalSeconds, document.PendingSeconds, document.ResolveSeconds, document.CooldownSeconds, document.UnknownBehavior, groupJSON, suppressJSON, messageJSON, versionID, id, expected)
	if err != nil {
		return Rule{}, mapError(err)
	}
	if tag.RowsAffected() != 1 {
		return Rule{}, ErrConflict
	}
	if err = tx.Commit(ctx); err != nil {
		return Rule{}, fmt.Errorf("commit alert rule update: %w", err)
	}
	return r.Get(ctx, id)
}

func (r Repository) Versions(ctx context.Context, id string, limit int) ([]Version, error) {
	if r.Pool == nil {
		return nil, ErrUnavailable
	}
	if limit < 1 || limit > 100 {
		return nil, errors.New("version limit is invalid")
	}
	rows, err := r.Pool.Query(ctx, `SELECT v.id,v.rule_id,v.version_number,v.document,v.change_summary,COALESCE(u.external_subject,''),v.created_at FROM alert_rule_versions v LEFT JOIN users u ON u.id=v.created_by WHERE v.rule_id=$1 ORDER BY v.version_number DESC,v.id ASC LIMIT $2`, id, limit)
	if err != nil {
		return nil, fmt.Errorf("list alert rule versions: %w", err)
	}
	defer rows.Close()
	result := make([]Version, 0, limit)
	for rows.Next() {
		var version Version
		var raw []byte
		if err := rows.Scan(&version.ID, &version.RuleID, &version.VersionNumber, &raw, &version.ChangeSummary, &version.CreatedBy, &version.CreatedAt); err != nil {
			return nil, fmt.Errorf("scan alert rule version: %w", err)
		}
		var document Document
		if _, err := DecodeDocument(raw, r.Registry); err != nil {
			return nil, fmt.Errorf("decode stored alert rule version: %w", err)
		} else {
			document = mustDecode(raw)
		}
		version.Document = document
		result = append(result, version)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if len(result) == 0 {
		var exists bool
		if err := r.Pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM alert_rules WHERE id=$1)`, id).Scan(&exists); err != nil {
			return nil, err
		}
		if !exists {
			return nil, ErrNotFound
		}
	}
	return result, nil
}

func (r Repository) SetEnabled(ctx context.Context, id string, expected int64, enabled bool) (Rule, error) {
	if r.Pool == nil {
		return Rule{}, ErrUnavailable
	}
	tx, err := r.Pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return Rule{}, err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	var revision int64
	var current bool
	if err := tx.QueryRow(ctx, `SELECT revision,enabled FROM alert_rules WHERE id=$1 FOR UPDATE`, id).Scan(&revision, &current); errors.Is(err, pgx.ErrNoRows) {
		return Rule{}, ErrNotFound
	} else if err != nil {
		return Rule{}, err
	}
	if revision != expected {
		return Rule{}, ErrConflict
	}
	if current == enabled {
		if err := tx.Commit(ctx); err != nil {
			return Rule{}, err
		}
		return r.Get(ctx, id)
	}
	if _, err := tx.Exec(ctx, `UPDATE alert_rules SET enabled=$1,revision=revision+1,updated_at=now() WHERE id=$2 AND revision=$3`, enabled, id, expected); err != nil {
		return Rule{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return Rule{}, err
	}
	return r.Get(ctx, id)
}

func (r Repository) ListEnabled(ctx context.Context, limit int) ([]Rule, error) {
	if r.Pool == nil {
		return nil, ErrUnavailable
	}
	if limit < 1 || limit > 100 {
		return nil, errors.New("enabled alert rule limit is invalid")
	}
	rows, err := r.Pool.Query(ctx, `SELECT r.id,r.schema_version,r.name,r.enabled,r.severity,r.scope,r.condition,r.evaluation_interval_seconds,r.pending_seconds,r.resolve_seconds,r.cooldown_seconds,r.unknown_behavior,r.group_by,r.suppress_when,r.message,r.revision,v.version_number,COALESCE(u.external_subject,''),r.created_at,r.updated_at FROM alert_rules r JOIN alert_rule_versions v ON v.id=r.current_version_id LEFT JOIN users u ON u.id=r.created_by WHERE r.enabled=true ORDER BY r.evaluation_interval_seconds ASC,r.name ASC,r.id ASC LIMIT $1`, limit)
	if err != nil {
		return nil, fmt.Errorf("list enabled alert rules: %w", err)
	}
	defer rows.Close()
	result := make([]Rule, 0, limit)
	for rows.Next() {
		var rule Rule
		var scopeJSON, conditionJSON, groupJSON, suppressJSON, messageJSON []byte
		if err := rows.Scan(&rule.ID, &rule.SchemaVersion, &rule.Name, &rule.Enabled, &rule.Severity, &scopeJSON, &conditionJSON, &rule.EvaluationIntervalSeconds, &rule.PendingSeconds, &rule.ResolveSeconds, &rule.CooldownSeconds, &rule.UnknownBehavior, &groupJSON, &suppressJSON, &messageJSON, &rule.Revision, &rule.CurrentVersion, &rule.CreatedBy, &rule.CreatedAt, &rule.UpdatedAt); err != nil {
			return nil, fmt.Errorf("scan enabled alert rule: %w", err)
		}
		if err := decodeStored(&rule.Document, scopeJSON, conditionJSON, groupJSON, suppressJSON, messageJSON, rule.CooldownSeconds); err != nil {
			return nil, err
		}
		result = append(result, rule)
	}
	return result, rows.Err()
}

func decodeStored(document *Document, scopeJSON, conditionJSON, groupJSON, suppressJSON, messageJSON []byte, cooldownSeconds int) error {
	if err := json.Unmarshal(scopeJSON, &document.Scope); err != nil {
		return errors.New("invalid stored alert rule scope")
	}
	if err := json.Unmarshal(conditionJSON, &document.Condition); err != nil {
		return errors.New("invalid stored alert rule condition")
	}
	if err := json.Unmarshal(groupJSON, &document.GroupBy); err != nil {
		return errors.New("invalid stored alert rule groups")
	}
	if err := json.Unmarshal(suppressJSON, &document.SuppressWhen); err != nil {
		return errors.New("invalid stored alert rule suppression")
	}
	document.CooldownSeconds = cooldownSeconds
	if err := json.Unmarshal(messageJSON, &document.Message); err != nil {
		return errors.New("invalid stored alert rule message")
	}
	return nil
}

func mustDecode(raw []byte) Document {
	var document Document
	_ = json.Unmarshal(raw, &document)
	return document
}

func sameJSON(left, right []byte) bool {
	var a, b any
	if json.Unmarshal(left, &a) != nil || json.Unmarshal(right, &b) != nil {
		return bytes.Equal(bytes.TrimSpace(left), bytes.TrimSpace(right))
	}
	return reflect.DeepEqual(a, b)
}

func nonNilMap(value map[string]any) map[string]any {
	if value == nil {
		return map[string]any{}
	}
	return value
}
func nonNilStrings(value []string) []string {
	if value == nil {
		return []string{}
	}
	return value
}
func valueOrEmpty(value *string) string {
	if value == nil {
		return ""
	}
	return *value
}

func mapError(err error) error {
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "23505" {
		return ErrConflict
	}
	return err
}