easyai-ai-gateway/apps/api/internal/store/runtime_policies.go

package store

import (
	"context"
	"encoding/json"
	"strings"

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

const runtimePolicyColumns = `
id::text, policy_key, name, COALESCE(description, ''), rate_limit_policy, retry_policy,
auto_disable_policy, degrade_policy, metadata, status, created_at, updated_at`

type RuntimePolicySetInput struct {
	PolicyKey         string         `json:"policyKey"`
	Name              string         `json:"name"`
	Description       string         `json:"description"`
	RateLimitPolicy   map[string]any `json:"rateLimitPolicy"`
	RetryPolicy       map[string]any `json:"retryPolicy"`
	AutoDisablePolicy map[string]any `json:"autoDisablePolicy"`
	DegradePolicy     map[string]any `json:"degradePolicy"`
	Metadata          map[string]any `json:"metadata"`
	Status            string         `json:"status"`
}

type runtimePolicyScanner interface {
	Scan(dest ...any) error
}

func (s *Store) ListRuntimePolicySets(ctx context.Context) ([]RuntimePolicySet, error) {
	rows, err := s.pool.Query(ctx, `SELECT `+runtimePolicyColumns+` FROM model_runtime_policy_sets ORDER BY policy_key ASC`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := make([]RuntimePolicySet, 0)
	for rows.Next() {
		item, err := scanRuntimePolicySet(rows)
		if err != nil {
			return nil, err
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) CreateRuntimePolicySet(ctx context.Context, input RuntimePolicySetInput) (RuntimePolicySet, error) {
	input = normalizeRuntimePolicyInput(input)
	rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.RateLimitPolicy))
	retryPolicy, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
	autoDisablePolicy, _ := json.Marshal(emptyObjectIfNil(input.AutoDisablePolicy))
	degradePolicy, _ := json.Marshal(emptyObjectIfNil(input.DegradePolicy))
	metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
	return scanRuntimePolicySet(s.pool.QueryRow(ctx, `
INSERT INTO model_runtime_policy_sets (
  policy_key, name, description, rate_limit_policy, retry_policy, auto_disable_policy, degrade_policy, metadata, status
)
VALUES ($1, $2, NULLIF($3, ''), $4, $5, $6, $7, $8, $9)
RETURNING `+runtimePolicyColumns,
		input.PolicyKey, input.Name, input.Description, rateLimitPolicy, retryPolicy,
		autoDisablePolicy, degradePolicy, metadata, input.Status,
	))
}

func (s *Store) UpdateRuntimePolicySet(ctx context.Context, id string, input RuntimePolicySetInput) (RuntimePolicySet, error) {
	input = normalizeRuntimePolicyInput(input)
	rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.RateLimitPolicy))
	retryPolicy, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
	autoDisablePolicy, _ := json.Marshal(emptyObjectIfNil(input.AutoDisablePolicy))
	degradePolicy, _ := json.Marshal(emptyObjectIfNil(input.DegradePolicy))
	metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
	return scanRuntimePolicySet(s.pool.QueryRow(ctx, `
UPDATE model_runtime_policy_sets
SET policy_key = $2,
    name = $3,
    description = NULLIF($4, ''),
    rate_limit_policy = $5,
    retry_policy = $6,
    auto_disable_policy = $7,
    degrade_policy = $8,
    metadata = $9,
    status = $10,
    updated_at = now()
WHERE id = $1::uuid
RETURNING `+runtimePolicyColumns,
		id, input.PolicyKey, input.Name, input.Description, rateLimitPolicy, retryPolicy,
		autoDisablePolicy, degradePolicy, metadata, input.Status,
	))
}

func (s *Store) DeleteRuntimePolicySet(ctx context.Context, id string) error {
	var policyKey string
	if err := s.pool.QueryRow(ctx, `SELECT policy_key FROM model_runtime_policy_sets WHERE id = $1::uuid`, id).Scan(&policyKey); err != nil {
		return err
	}
	if policyKey == "default-runtime-v1" {
		return ErrProtectedDefault
	}
	result, err := s.pool.Exec(ctx, `DELETE FROM model_runtime_policy_sets WHERE id = $1::uuid`, id)
	if err != nil {
		return err
	}
	if result.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}

func scanRuntimePolicySet(scanner runtimePolicyScanner) (RuntimePolicySet, error) {
	var item RuntimePolicySet
	var rateLimitPolicy []byte
	var retryPolicy []byte
	var autoDisablePolicy []byte
	var degradePolicy []byte
	var metadata []byte
	if err := scanner.Scan(
		&item.ID,
		&item.PolicyKey,
		&item.Name,
		&item.Description,
		&rateLimitPolicy,
		&retryPolicy,
		&autoDisablePolicy,
		&degradePolicy,
		&metadata,
		&item.Status,
		&item.CreatedAt,
		&item.UpdatedAt,
	); err != nil {
		return RuntimePolicySet{}, err
	}
	item.RateLimitPolicy = decodeObject(rateLimitPolicy)
	item.RetryPolicy = decodeObject(retryPolicy)
	item.AutoDisablePolicy = decodeObject(autoDisablePolicy)
	item.DegradePolicy = decodeObject(degradePolicy)
	item.Metadata = decodeObject(metadata)
	return item, nil
}

func normalizeRuntimePolicyInput(input RuntimePolicySetInput) RuntimePolicySetInput {
	input.PolicyKey = strings.TrimSpace(input.PolicyKey)
	input.Name = strings.TrimSpace(input.Name)
	input.Description = strings.TrimSpace(input.Description)
	input.Status = strings.TrimSpace(input.Status)
	if input.Status == "" {
		input.Status = "active"
	}
	return input
}