ocdp-go/backend/internal/adapter/output/persistence/postgres/quota_repository.go

package postgres

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"github.com/google/uuid"
	"github.com/ocdp/cluster-service/internal/domain/entity"
	"github.com/ocdp/cluster-service/internal/domain/repository"
)

// QuotaRepository PostgreSQL 配额仓储实现
type QuotaRepository struct {
	db *DB
}

// NewQuotaRepository 创建 PostgreSQL 配额仓储
func NewQuotaRepository(db *DB) repository.QuotaRepository {
	return &QuotaRepository{db: db}
}

// Create 创建配额
func (r *QuotaRepository) Create(ctx context.Context, quota *entity.WorkspaceQuota) error {
	if quota.ID == "" {
		quota.ID = uuid.New().String()
	}

	query := `
		INSERT INTO workspace_quotas (id, workspace_id, resource_type, hard_limit, soft_limit, used, created_at, updated_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
		ON CONFLICT (workspace_id, resource_type) DO UPDATE
		SET hard_limit = $4, soft_limit = $5, updated_at = $8
	`

	_, err := r.db.conn.ExecContext(ctx, query,
		quota.ID,
		quota.WorkspaceID,
		quota.ResourceType,
		quota.HardLimit,
		quota.SoftLimit,
		quota.Used,
		quota.CreatedAt,
		quota.UpdatedAt,
	)

	if err != nil {
		return fmt.Errorf("failed to create quota: %w", err)
	}

	return nil
}

// GetByID 根据 ID 获取配额
func (r *QuotaRepository) GetByID(ctx context.Context, id string) (*entity.WorkspaceQuota, error) {
	query := `
		SELECT id, workspace_id, resource_type, hard_limit, soft_limit, used, created_at, updated_at
		FROM workspace_quotas
		WHERE id = $1
	`

	quota := &entity.WorkspaceQuota{}
	err := r.db.conn.QueryRowContext(ctx, query, id).Scan(
		&quota.ID,
		&quota.WorkspaceID,
		&quota.ResourceType,
		&quota.HardLimit,
		&quota.SoftLimit,
		&quota.Used,
		&quota.CreatedAt,
		&quota.UpdatedAt,
	)

	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get quota: %w", err)
	}

	return quota, nil
}

// GetByWorkspaceAndType 根据 workspace 和资源类型获取配额
func (r *QuotaRepository) GetByWorkspaceAndType(ctx context.Context, workspaceID string, resourceType entity.ResourceType) (*entity.WorkspaceQuota, error) {
	query := `
		SELECT id, workspace_id, resource_type, hard_limit, soft_limit, used, created_at, updated_at
		FROM workspace_quotas
		WHERE workspace_id = $1 AND resource_type = $2
	`

	quota := &entity.WorkspaceQuota{}
	err := r.db.conn.QueryRowContext(ctx, query, workspaceID, resourceType).Scan(
		&quota.ID,
		&quota.WorkspaceID,
		&quota.ResourceType,
		&quota.HardLimit,
		&quota.SoftLimit,
		&quota.Used,
		&quota.CreatedAt,
		&quota.UpdatedAt,
	)

	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to get quota: %w", err)
	}

	return quota, nil
}

// GetByWorkspace 获取 workspace 的所有配额
func (r *QuotaRepository) GetByWorkspace(ctx context.Context, workspaceID string) ([]*entity.WorkspaceQuota, error) {
	query := `
		SELECT id, workspace_id, resource_type, hard_limit, soft_limit, used, created_at, updated_at
		FROM workspace_quotas
		WHERE workspace_id = $1
		ORDER BY resource_type
	`

	rows, err := r.db.conn.QueryContext(ctx, query, workspaceID)
	if err != nil {
		return nil, fmt.Errorf("failed to list quotas: %w", err)
	}
	defer rows.Close()

	quotas := make([]*entity.WorkspaceQuota, 0)
	for rows.Next() {
		quota := &entity.WorkspaceQuota{}
		err := rows.Scan(
			&quota.ID,
			&quota.WorkspaceID,
			&quota.ResourceType,
			&quota.HardLimit,
			&quota.SoftLimit,
			&quota.Used,
			&quota.CreatedAt,
			&quota.UpdatedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan quota: %w", err)
		}
		quotas = append(quotas, quota)
	}

	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("rows iteration error: %w", err)
	}

	return quotas, nil
}

// Update 更新配额
func (r *QuotaRepository) Update(ctx context.Context, quota *entity.WorkspaceQuota) error {
	quota.UpdatedAt = time.Now()

	query := `
		UPDATE workspace_quotas
		SET hard_limit = $1, soft_limit = $2, used = $3, updated_at = $4
		WHERE id = $5
	`

	result, err := r.db.conn.ExecContext(ctx, query,
		quota.HardLimit,
		quota.SoftLimit,
		quota.Used,
		quota.UpdatedAt,
		quota.ID,
	)

	if err != nil {
		return fmt.Errorf("failed to update quota: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get affected rows: %w", err)
	}

	if rows == 0 {
		return fmt.Errorf("quota not found")
	}

	return nil
}

// Delete 删除配额
func (r *QuotaRepository) Delete(ctx context.Context, id string) error {
	query := `DELETE FROM workspace_quotas WHERE id = $1`

	_, err := r.db.conn.ExecContext(ctx, query, id)
	if err != nil {
		return fmt.Errorf("failed to delete quota: %w", err)
	}

	return nil
}

// DeleteByWorkspace 删除 workspace 的所有配额
func (r *QuotaRepository) DeleteByWorkspace(ctx context.Context, workspaceID string) error {
	query := `DELETE FROM workspace_quotas WHERE workspace_id = $1`

	_, err := r.db.conn.ExecContext(ctx, query, workspaceID)
	if err != nil {
		return fmt.Errorf("failed to delete quotas: %w", err)
	}

	return nil
}