story/internal/idempotency/executor.go

package idempotency

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

var (
	ErrKeyConflict = errors.New("idempotency key was reused with different request content")
	ErrInProgress  = errors.New("idempotent operation is still in progress")
)

type Request struct {
	TenantID     *string
	ActorKey     string
	Scope        string
	Key          string
	RequestHash  string
	ExpiresAfter time.Duration
}

type Result struct {
	StatusCode int
	Body       json.RawMessage
	ResourceID *string
	Replayed   bool
}

type Executor struct {
	pool *pgxpool.Pool
}

func NewExecutor(pool *pgxpool.Pool) *Executor { return &Executor{pool: pool} }

func Hash(body []byte) string {
	digest := sha256.Sum256(body)
	return hex.EncodeToString(digest[:])
}

// Lookup returns a completed result without opening the operation transaction.
// Callers with an expensive, side-effect-free precomputation can use it to
// avoid repeating that work on ordinary sequential idempotency replays. Execute
// remains the authority for races and must still be called when found is false.
func (e *Executor) Lookup(ctx context.Context, request Request) (result Result, found bool, err error) {
	if err = validateRequest(request); err != nil {
		return Result{}, false, err
	}
	var storedHash, state string
	err = e.pool.QueryRow(ctx, `
		SELECT request_hash,state,response_status,response_body,resource_id::text
		FROM idempotency_record
		WHERE actor_key=$1 AND scope=$2 AND idempotency_key=$3`,
		request.ActorKey, request.Scope, request.Key).
		Scan(&storedHash, &state, &result.StatusCode, &result.Body, &result.ResourceID)
	if errors.Is(err, pgx.ErrNoRows) {
		return Result{}, false, nil
	}
	if err != nil {
		return Result{}, false, fmt.Errorf("lookup idempotency record: %w", err)
	}
	if storedHash != request.RequestHash {
		return Result{}, true, ErrKeyConflict
	}
	if state != "completed" {
		return Result{}, true, ErrInProgress
	}
	result.Replayed = true
	return result, true, nil
}

func (e *Executor) Execute(
	ctx context.Context,
	request Request,
	operation func(context.Context, pgx.Tx) (Result, error),
) (Result, error) {
	return e.ExecuteGuarded(ctx, request, nil, operation)
}

// Guard executes in the operation transaction before mutation or replay. On
// replay it also receives the stored resource for current authorization checks.
type Guard func(context.Context, pgx.Tx, *Result) error

func (e *Executor) ExecuteGuarded(ctx context.Context, request Request, guard Guard, operation func(context.Context, pgx.Tx) (Result, error)) (Result, error) {
	if err := validateRequest(request); err != nil {
		return Result{}, err
	}
	tx, err := e.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return Result{}, fmt.Errorf("begin idempotency transaction: %w", err)
	}
	defer func() { _ = tx.Rollback(context.Background()) }()
	if guard != nil {
		if err := guard(ctx, tx, nil); err != nil {
			return Result{}, err
		}
	}

	tag, err := tx.Exec(ctx, `
		INSERT INTO idempotency_record (
			tenant_id, actor_key, scope, idempotency_key, request_hash,
			state, expires_at
		) VALUES ($1::uuid, $2, $3, $4, $5, 'processing', now() + $6::interval)
		ON CONFLICT (actor_key, scope, idempotency_key) DO NOTHING`,
		request.TenantID, request.ActorKey, request.Scope, request.Key,
		request.RequestHash, postgresInterval(request.ExpiresAfter),
	)
	if err != nil {
		return Result{}, fmt.Errorf("claim idempotency key: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return e.replay(ctx, tx, request, guard)
	}

	result, err := operation(ctx, tx)
	if err != nil {
		return Result{}, err
	}
	if !json.Valid(result.Body) {
		return Result{}, errors.New("idempotent response body must be valid JSON")
	}
	_, err = tx.Exec(ctx, `
		UPDATE idempotency_record
		SET state = 'completed', response_status = $4, response_body = $5::jsonb,
			resource_id = $6::uuid, updated_at = now()
		WHERE actor_key = $1 AND scope = $2 AND idempotency_key = $3`,
		request.ActorKey, request.Scope, request.Key, result.StatusCode, result.Body, result.ResourceID,
	)
	if err != nil {
		return Result{}, fmt.Errorf("complete idempotency record: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return Result{}, fmt.Errorf("commit idempotent operation: %w", err)
	}
	return result, nil
}

func (e *Executor) replay(ctx context.Context, tx pgx.Tx, request Request, guard Guard) (Result, error) {
	var storedHash, state string
	var result Result
	err := tx.QueryRow(ctx, `
		SELECT request_hash, state, response_status, response_body, resource_id::text
		FROM idempotency_record
		WHERE actor_key = $1 AND scope = $2 AND idempotency_key = $3
		FOR UPDATE`, request.ActorKey, request.Scope, request.Key,
	).Scan(&storedHash, &state, &result.StatusCode, &result.Body, &result.ResourceID)
	if err != nil {
		return Result{}, fmt.Errorf("read idempotency record: %w", err)
	}
	if storedHash != request.RequestHash {
		return Result{}, ErrKeyConflict
	}
	if state != "completed" {
		return Result{}, ErrInProgress
	}
	result.Replayed = true
	if guard != nil {
		if err := guard(ctx, tx, &result); err != nil {
			return Result{}, err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return Result{}, fmt.Errorf("commit idempotency replay: %w", err)
	}
	return result, nil
}

func validateRequest(request Request) error {
	if request.ActorKey == "" || request.Scope == "" || request.Key == "" || request.RequestHash == "" {
		return errors.New("actor key, scope, idempotency key and request hash are required")
	}
	if request.ExpiresAfter <= 0 {
		return errors.New("idempotency expiry must be positive")
	}
	return nil
}

func postgresInterval(duration time.Duration) string {
	return fmt.Sprintf("%f seconds", duration.Seconds())
}