proxy-pool/internal/adapters/postgresadmin/mutate.go

package postgresadmin

import (
	"context"
	"encoding/json"
	"errors"
	"math"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"

	"proxy-pool/internal/domain/adminstate"
)

const (
	lockRevisionSQL    = `LOCK TABLE control_revisions IN SHARE ROW EXCLUSIVE MODE`
	currentRevisionSQL = `SELECT COALESCE(MAX(revision), 0)::bigint FROM control_revisions`
	insertRevisionSQL  = `
		INSERT INTO control_revisions (revision, kind, created_at)
		VALUES ($1, $2, $3)`
	insertAuditSQL = `
		INSERT INTO admin_audit_log (
			request_id, actor_id, source_ip, action, resource_type, resource_name,
			changed, revision, reason, occurred_at
		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`
	insertOutboxSQL = `
		INSERT INTO admin_outbox (
			revision, event_type, aggregate_type, aggregate_id, payload, occurred_at
		) VALUES ($1, $2, $3, $4, $5, $6)`
)

func (adapter *adapter) CommitConfig(
	ctx context.Context,
	command adminstate.CommitConfigCommand,
) (adminstate.MutationResult, error) {
	result := adminstate.MutationResult{RequestID: command.RequestID}
	if err := contextError(ctx); err != nil {
		return result, err
	}
	if command.Validate() != nil || !adapter.valid() {
		return result, adminstate.ErrInvalidCommand
	}
	command = cloneConfigCommand(command)

	tx, currentRevision, err := adapter.beginMutation(ctx, "commit config")
	if err != nil {
		return result, err
	}
	defer rollback(tx)

	var existingChecksum string
	var isCurrent bool
	err = tx.QueryRow(ctx, `
		SELECT checksum::text,
			revision = (SELECT MAX(revision) FROM config_revisions)
		FROM config_revisions
		WHERE config_version = $1`, command.ConfigVersion).Scan(&existingChecksum, &isCurrent)
	switch {
	case err == nil:
		if !isCurrent || !strings.EqualFold(existingChecksum, command.Checksum) {
			return result, adminstate.ErrConflict
		}
		if err := insertAudit(ctx, tx, command.RequestID, command.Actor,
			adminstate.ActionCommitConfig, "config", command.ConfigVersion, false,
			currentRevision, "", command.OccurredAt, "audit config no-op"); err != nil {
			return result, err
		}
		if err := commit(ctx, tx, "commit config no-op"); err != nil {
			return result, err
		}
		result.Revision = uint64(currentRevision)
		return result, nil
	case !errors.Is(err, pgx.ErrNoRows):
		return result, databaseError(ctx, "read config revision", err)
	}

	nextRevision, err := allocateRevision(ctx, tx, currentRevision, "config", command.OccurredAt)
	if err != nil {
		return result, err
	}
	if _, err := tx.Exec(ctx, `
		INSERT INTO config_revisions (
			revision, config_version, checksum, source, created_at
		) VALUES ($1, $2, $3, $4, $5)`,
		nextRevision, command.ConfigVersion, command.Checksum, command.Source, utc(command.OccurredAt)); err != nil {
		return result, databaseError(ctx, "insert config revision", err)
	}
	if _, err := tx.Exec(ctx, `DELETE FROM routing_admin_state`); err != nil {
		return result, databaseError(ctx, "replace routing state", err)
	}
	if _, err := tx.Exec(ctx, `DELETE FROM upstream_admin_state`); err != nil {
		return result, databaseError(ctx, "replace upstream state", err)
	}
	if err := copyUpstreams(ctx, tx, command.Upstreams, nextRevision, command.OccurredAt); err != nil {
		return result, err
	}
	if err := copyRoutings(ctx, tx, command.Routings, nextRevision, command.OccurredAt); err != nil {
		return result, err
	}
	if err := insertAudit(ctx, tx, command.RequestID, command.Actor,
		adminstate.ActionCommitConfig, "config", command.ConfigVersion, true,
		nextRevision, "", command.OccurredAt, "audit config change"); err != nil {
		return result, err
	}
	payload, err := encodePayload(map[string]any{
		"configVersion": command.ConfigVersion,
		"checksum":      command.Checksum,
		"revision":      nextRevision,
	})
	if err != nil {
		return result, unavailable("encode config event")
	}
	if err := insertOutbox(ctx, tx, nextRevision, "config.committed", "config",
		command.ConfigVersion, payload, command.OccurredAt, "insert config event"); err != nil {
		return result, err
	}
	if err := commit(ctx, tx, "commit config change"); err != nil {
		return result, err
	}
	return adminstate.MutationResult{
		RequestID: command.RequestID,
		Changed:   true,
		Revision:  uint64(nextRevision),
	}, nil
}

func (adapter *adapter) SetUpstreamEnabled(
	ctx context.Context,
	command adminstate.SetUpstreamCommand,
) (adminstate.MutationResult, error) {
	result := adminstate.MutationResult{RequestID: command.RequestID}
	if err := contextError(ctx); err != nil {
		return result, err
	}
	if command.Validate() != nil || !adapter.valid() {
		return result, adminstate.ErrInvalidCommand
	}

	tx, currentRevision, err := adapter.beginMutation(ctx, "set upstream")
	if err != nil {
		return result, err
	}
	defer rollback(tx)

	var enabled bool
	if err := tx.QueryRow(ctx, `
		SELECT enabled
		FROM upstream_admin_state
		WHERE name = $1
		FOR UPDATE`, command.Name).Scan(&enabled); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return result, adminstate.ErrNotFound
		}
		return result, databaseError(ctx, "read upstream state", err)
	}
	if enabled == command.Enabled {
		if err := insertAudit(ctx, tx, command.RequestID, command.Actor,
			adminstate.ActionSetUpstream, "upstream", command.Name, false,
			currentRevision, "", command.OccurredAt, "audit upstream no-op"); err != nil {
			return result, err
		}
		if err := commit(ctx, tx, "commit upstream no-op"); err != nil {
			return result, err
		}
		result.Revision = uint64(currentRevision)
		return result, nil
	}

	nextRevision, err := allocateRevision(ctx, tx, currentRevision, "upstream", command.OccurredAt)
	if err != nil {
		return result, err
	}
	commandTag, err := tx.Exec(ctx, `
		UPDATE upstream_admin_state
		SET enabled = $1, revision = $2, updated_at = $3
		WHERE name = $4`, command.Enabled, nextRevision, utc(command.OccurredAt), command.Name)
	if err != nil {
		return result, databaseError(ctx, "update upstream state", err)
	}
	if commandTag.RowsAffected() != 1 {
		return result, unavailable("update upstream state")
	}
	if err := insertAudit(ctx, tx, command.RequestID, command.Actor,
		adminstate.ActionSetUpstream, "upstream", command.Name, true,
		nextRevision, "", command.OccurredAt, "audit upstream change"); err != nil {
		return result, err
	}
	payload, err := encodePayload(map[string]any{
		"enabled":  command.Enabled,
		"name":     command.Name,
		"revision": nextRevision,
	})
	if err != nil {
		return result, unavailable("encode upstream event")
	}
	if err := insertOutbox(ctx, tx, nextRevision, "upstream.enabled_changed", "upstream",
		command.Name, payload, command.OccurredAt, "insert upstream event"); err != nil {
		return result, err
	}
	if err := commit(ctx, tx, "commit upstream change"); err != nil {
		return result, err
	}
	return adminstate.MutationResult{
		RequestID: command.RequestID,
		Changed:   true,
		Revision:  uint64(nextRevision),
	}, nil
}

func (adapter *adapter) SwitchRouting(
	ctx context.Context,
	command adminstate.SwitchRoutingCommand,
) (adminstate.MutationResult, error) {
	result := adminstate.MutationResult{RequestID: command.RequestID}
	if err := contextError(ctx); err != nil {
		return result, err
	}
	if command.Validate() != nil || !adapter.valid() {
		return result, adminstate.ErrInvalidCommand
	}

	tx, currentRevision, err := adapter.beginMutation(ctx, "switch routing")
	if err != nil {
		return result, err
	}
	defer rollback(tx)

	var enabled bool
	var candidates []string
	var current string
	if err := tx.QueryRow(ctx, `
		SELECT enabled, upstreams, current_upstream
		FROM routing_admin_state
		WHERE name = $1
		FOR UPDATE`, command.Name).Scan(&enabled, &candidates, &current); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return result, adminstate.ErrNotFound
		}
		return result, databaseError(ctx, "read routing state", err)
	}
	if !enabled || current != command.ExpectedCurrent {
		return result, adminstate.ErrConflict
	}
	if !contains(candidates, command.Target) {
		return result, adminstate.ErrInvalidCommand
	}
	if current == command.Target {
		if err := insertAudit(ctx, tx, command.RequestID, command.Actor,
			adminstate.ActionSwitchRoute, "routing", command.Name, false,
			currentRevision, command.Reason, command.OccurredAt, "audit routing no-op"); err != nil {
			return result, err
		}
		if err := commit(ctx, tx, "commit routing no-op"); err != nil {
			return result, err
		}
		result.Revision = uint64(currentRevision)
		return result, nil
	}

	nextRevision, err := allocateRevision(ctx, tx, currentRevision, "routing", command.OccurredAt)
	if err != nil {
		return result, err
	}
	commandTag, err := tx.Exec(ctx, `
		UPDATE routing_admin_state
		SET current_upstream = $1, revision = $2, updated_at = $3
		WHERE name = $4`, command.Target, nextRevision, utc(command.OccurredAt), command.Name)
	if err != nil {
		return result, databaseError(ctx, "update routing state", err)
	}
	if commandTag.RowsAffected() != 1 {
		return result, unavailable("update routing state")
	}
	if err := insertAudit(ctx, tx, command.RequestID, command.Actor,
		adminstate.ActionSwitchRoute, "routing", command.Name, true,
		nextRevision, command.Reason, command.OccurredAt, "audit routing change"); err != nil {
		return result, err
	}
	payload, err := encodePayload(map[string]any{
		"current":  command.Target,
		"name":     command.Name,
		"previous": command.ExpectedCurrent,
		"reason":   command.Reason,
		"revision": nextRevision,
	})
	if err != nil {
		return result, unavailable("encode routing event")
	}
	if err := insertOutbox(ctx, tx, nextRevision, "routing.switched", "routing",
		command.Name, payload, command.OccurredAt, "insert routing event"); err != nil {
		return result, err
	}
	if err := commit(ctx, tx, "commit routing change"); err != nil {
		return result, err
	}
	return adminstate.MutationResult{
		RequestID: command.RequestID,
		Changed:   true,
		Revision:  uint64(nextRevision),
	}, nil
}

func (adapter *adapter) beginMutation(ctx context.Context, operation string) (pgx.Tx, int64, error) {
	tx, err := adapter.begin(ctx, pgx.TxOptions{AccessMode: pgx.ReadWrite}, operation)
	if err != nil {
		return nil, 0, err
	}
	if _, err := tx.Exec(ctx, lockRevisionSQL); err != nil {
		rollback(tx)
		return nil, 0, databaseError(ctx, "lock revision state", err)
	}
	var currentRevision int64
	if err := tx.QueryRow(ctx, currentRevisionSQL).Scan(&currentRevision); err != nil {
		rollback(tx)
		return nil, 0, databaseError(ctx, "read current revision", err)
	}
	if currentRevision < 0 {
		rollback(tx)
		return nil, 0, unavailable("read current revision")
	}
	return tx, currentRevision, nil
}

func allocateRevision(
	ctx context.Context,
	tx pgx.Tx,
	currentRevision int64,
	kind string,
	occurredAt time.Time,
) (int64, error) {
	if currentRevision == math.MaxInt64 {
		return 0, adminstate.ErrUnavailable
	}
	nextRevision := currentRevision + 1
	commandTag, err := tx.Exec(ctx, insertRevisionSQL, nextRevision, kind, utc(occurredAt))
	if err != nil {
		return 0, databaseError(ctx, "allocate revision", err)
	}
	if commandTag.RowsAffected() != 1 {
		return 0, unavailable("allocate revision")
	}
	return nextRevision, nil
}

func copyUpstreams(
	ctx context.Context,
	tx pgx.Tx,
	definitions []adminstate.UpstreamDefinition,
	revision int64,
	occurredAt time.Time,
) error {
	if len(definitions) == 0 {
		return nil
	}
	updatedAt := utc(occurredAt)
	count, err := tx.CopyFrom(ctx, pgx.Identifier{"upstream_admin_state"},
		[]string{"name", "enabled", "revision", "updated_at"},
		pgx.CopyFromSlice(len(definitions), func(index int) ([]any, error) {
			definition := definitions[index]
			return []any{definition.Name, definition.Enabled, revision, updatedAt}, nil
		}))
	if err != nil {
		return databaseError(ctx, "replace upstream state", err)
	}
	if count != int64(len(definitions)) {
		return unavailable("replace upstream state")
	}
	return nil
}

func copyRoutings(
	ctx context.Context,
	tx pgx.Tx,
	definitions []adminstate.RoutingDefinition,
	revision int64,
	occurredAt time.Time,
) error {
	if len(definitions) == 0 {
		return nil
	}
	updatedAt := utc(occurredAt)
	count, err := tx.CopyFrom(ctx, pgx.Identifier{"routing_admin_state"},
		[]string{"name", "enabled", "upstreams", "current_upstream", "revision", "updated_at"},
		pgx.CopyFromSlice(len(definitions), func(index int) ([]any, error) {
			definition := definitions[index]
			return []any{
				definition.Name,
				definition.Enabled,
				definition.Upstreams,
				definition.CurrentUpstream,
				revision,
				updatedAt,
			}, nil
		}))
	if err != nil {
		return databaseError(ctx, "replace routing state", err)
	}
	if count != int64(len(definitions)) {
		return unavailable("replace routing state")
	}
	return nil
}

func insertAudit(
	ctx context.Context,
	tx pgx.Tx,
	requestID string,
	actor adminstate.Actor,
	action adminstate.Action,
	resourceType string,
	resourceName string,
	changed bool,
	revision int64,
	reason string,
	occurredAt time.Time,
	operation string,
) error {
	commandTag, err := tx.Exec(ctx, insertAuditSQL,
		requestID,
		actor.ID,
		sourceIPValue(actor.SourceIP),
		string(action),
		resourceType,
		resourceName,
		changed,
		revision,
		reason,
		utc(occurredAt),
	)
	if err != nil {
		return databaseError(ctx, operation, err)
	}
	if commandTag.RowsAffected() != 1 {
		return unavailable(operation)
	}
	return nil
}

func insertOutbox(
	ctx context.Context,
	tx pgx.Tx,
	revision int64,
	eventType string,
	aggregateType string,
	aggregateID string,
	payload json.RawMessage,
	occurredAt time.Time,
	operation string,
) error {
	commandTag, err := tx.Exec(ctx, insertOutboxSQL,
		revision,
		eventType,
		aggregateType,
		aggregateID,
		payload,
		utc(occurredAt),
	)
	if err != nil {
		return databaseError(ctx, operation, err)
	}
	if commandTag.RowsAffected() != 1 {
		return unavailable(operation)
	}
	return nil
}

func contains(values []string, target string) bool {
	for _, value := range values {
		if value == target {
			return true
		}
	}
	return false
}