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( "a.ID, "a.WorkspaceID, "a.ResourceType, "a.HardLimit, "a.SoftLimit, "a.Used, "a.CreatedAt, "a.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( "a.ID, "a.WorkspaceID, "a.ResourceType, "a.HardLimit, "a.SoftLimit, "a.Used, "a.CreatedAt, "a.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( "a.ID, "a.WorkspaceID, "a.ResourceType, "a.HardLimit, "a.SoftLimit, "a.Used, "a.CreatedAt, "a.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 }