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

package store

import (
	"context"
	"strings"

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

const maxIdentityBatchSize = 500

type IdentityBatchInput struct {
	IDs    []string `json:"ids"`
	Action string   `json:"action" enums:"enable,disable,delete"`
}

type IdentityBatchResult struct {
	Action         string   `json:"action"`
	RequestedCount int      `json:"requestedCount"`
	AffectedCount  int      `json:"affectedCount"`
	IDs            []string `json:"ids"`
}

func (s *Store) BatchGatewayUsers(ctx context.Context, input IdentityBatchInput) (IdentityBatchResult, error) {
	input, err := normalizeIdentityBatchInput(input)
	if err != nil {
		return IdentityBatchResult{}, err
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return IdentityBatchResult{}, err
	}
	defer rollbackTransaction(tx)

	locked, err := lockIdentityBatchTargets(ctx, tx, `
SELECT id::text
FROM gateway_users
WHERE id = ANY($1::uuid[])
  AND deleted_at IS NULL
FOR UPDATE`, input.IDs)
	if err != nil {
		return IdentityBatchResult{}, err
	}
	if len(locked) != len(input.IDs) {
		return IdentityBatchResult{}, ErrIdentityBatchTargetNotFound
	}

	switch input.Action {
	case "enable":
		_, err = tx.Exec(ctx, `
UPDATE gateway_users
SET status = 'active', updated_at = now()
WHERE id = ANY($1::uuid[]) AND deleted_at IS NULL`, input.IDs)
	case "disable":
		_, err = tx.Exec(ctx, `
UPDATE gateway_users
SET status = 'disabled', updated_at = now()
WHERE id = ANY($1::uuid[]) AND deleted_at IS NULL`, input.IDs)
	case "delete":
		_, err = tx.Exec(ctx, `
UPDATE gateway_users
SET deleted_at = now(),
    status = 'deleted',
    user_key = user_key || ':deleted:' || left(id::text, 8),
    external_user_id = CASE WHEN source = 'oidc' THEN external_user_id ELSE NULL END,
    email = NULL,
    updated_at = now()
WHERE id = ANY($1::uuid[]) AND deleted_at IS NULL`, input.IDs)
	}
	if err != nil {
		return IdentityBatchResult{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return IdentityBatchResult{}, err
	}
	return identityBatchResult(input), nil
}

func (s *Store) BatchUserGroups(ctx context.Context, input IdentityBatchInput) (IdentityBatchResult, error) {
	input, err := normalizeIdentityBatchInput(input)
	if err != nil {
		return IdentityBatchResult{}, err
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return IdentityBatchResult{}, err
	}
	defer rollbackTransaction(tx)

	locked, err := lockIdentityBatchTargets(ctx, tx, `
SELECT id::text
FROM gateway_user_groups
WHERE id = ANY($1::uuid[])
FOR UPDATE`, input.IDs)
	if err != nil {
		return IdentityBatchResult{}, err
	}
	if len(locked) != len(input.IDs) {
		return IdentityBatchResult{}, ErrIdentityBatchTargetNotFound
	}

	switch input.Action {
	case "enable":
		_, err = tx.Exec(ctx, `
UPDATE gateway_user_groups
SET status = 'active', updated_at = now()
WHERE id = ANY($1::uuid[])`, input.IDs)
	case "disable":
		_, err = tx.Exec(ctx, `
UPDATE gateway_user_groups
SET status = 'disabled', updated_at = now()
WHERE id = ANY($1::uuid[])`, input.IDs)
	case "delete":
		if _, err = tx.Exec(ctx, `
DELETE FROM gateway_access_rules
WHERE subject_type = 'user_group'
  AND subject_id = ANY($1::uuid[])`, input.IDs); err == nil {
			_, err = tx.Exec(ctx, `DELETE FROM gateway_user_groups WHERE id = ANY($1::uuid[])`, input.IDs)
		}
	}
	if err != nil {
		return IdentityBatchResult{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return IdentityBatchResult{}, err
	}
	return identityBatchResult(input), nil
}

func normalizeIdentityBatchInput(input IdentityBatchInput) (IdentityBatchInput, error) {
	input.Action = strings.ToLower(strings.TrimSpace(input.Action))
	if input.Action != "enable" && input.Action != "disable" && input.Action != "delete" {
		return IdentityBatchInput{}, ErrInvalidIdentityBatch
	}
	seen := make(map[string]bool, len(input.IDs))
	ids := make([]string, 0, len(input.IDs))
	for _, value := range input.IDs {
		id := strings.TrimSpace(value)
		parsed, err := uuid.Parse(id)
		if err != nil {
			return IdentityBatchInput{}, ErrInvalidIdentityBatch
		}
		id = parsed.String()
		if seen[id] {
			continue
		}
		seen[id] = true
		ids = append(ids, id)
	}
	if len(ids) == 0 || len(ids) > maxIdentityBatchSize {
		return IdentityBatchInput{}, ErrInvalidIdentityBatch
	}
	input.IDs = ids
	return input, nil
}

func lockIdentityBatchTargets(ctx context.Context, tx pgx.Tx, query string, ids []string) ([]string, error) {
	rows, err := tx.Query(ctx, query, ids)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	locked := make([]string, 0, len(ids))
	for rows.Next() {
		var id string
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		locked = append(locked, id)
	}
	return locked, rows.Err()
}

func identityBatchResult(input IdentityBatchInput) IdentityBatchResult {
	return IdentityBatchResult{
		Action:         input.Action,
		RequestedCount: len(input.IDs),
		AffectedCount:  len(input.IDs),
		IDs:            append([]string(nil), input.IDs...),
	}
}