1 files changed, 81 insertions, 0 deletions
diff --git a/internal/controlplane/limits.go b/internal/controlplane/limits.go
new file mode 100644
index 0000000..5b5825e
--- /dev/null
+++ b/internal/controlplane/limits.go
@@ -0,0 +1,81 @@
+package controlplane
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/jackc/pgx/v5"
+)
+
+func (s *Store) ListProjectLimits(ctx context.Context, tenantID string) ([]ProjectLimit, error) {
+	query := `SELECT l.tenant_id::text, l.project_id::text, p.name, t.name,
+		l.requests_per_minute, l.tokens_per_minute, l.concurrent_requests, l.monthly_spend_micros, l.updated_at
+		FROM project_limits l JOIN projects p ON p.id=l.project_id JOIN tenants t ON t.id=l.tenant_id`
+	args := []any{}
+	if strings.TrimSpace(tenantID) != "" {
+		query += ` WHERE l.tenant_id=$1`
+		args = append(args, tenantID)
+	}
+	query += ` ORDER BY t.name, p.name`
+	rows, err := s.db.Query(ctx, query, args...)
+	if err != nil {
+		return nil, fmt.Errorf("query project limits: %w", err)
+	}
+	defer rows.Close()
+	result := make([]ProjectLimit, 0)
+	for rows.Next() {
+		var item ProjectLimit
+		if err := rows.Scan(&item.TenantID, &item.ProjectID, &item.ProjectName, &item.TenantName,
+			&item.RequestsPerMinute, &item.TokensPerMinute, &item.Concurrent, &item.MonthlySpendMicros, &item.UpdatedAt); err != nil {
+			return nil, fmt.Errorf("scan project limit: %w", err)
+		}
+		result = append(result, item)
+	}
+	return result, rows.Err()
+}
+
+func (s *Store) SetProjectLimit(ctx context.Context, projectID string, input SetProjectLimitInput) (ProjectLimit, int64, error) {
+	if strings.TrimSpace(projectID) == "" || input.RequestsPerMinute < 0 || input.TokensPerMinute < 0 || input.ConcurrentRequests < 0 || input.MonthlySpendMicros < 0 {
+		return ProjectLimit{}, 0, errors.New("limit values cannot be negative and project_id is required")
+	}
+	tx, err := s.db.Begin(ctx)
+	if err != nil {
+		return ProjectLimit{}, 0, err
+	}
+	defer tx.Rollback(ctx)
+	var result ProjectLimit
+	err = tx.QueryRow(ctx, `
+		WITH project AS (
+			SELECT p.id, p.tenant_id, p.name AS project_name, t.name AS tenant_name
+			FROM projects p JOIN tenants t ON t.id=p.tenant_id WHERE p.id=$1
+		), updated AS (
+			INSERT INTO project_limits (project_id, tenant_id, requests_per_minute, tokens_per_minute, concurrent_requests, monthly_spend_micros)
+			SELECT p.id, p.tenant_id, $2, $3, $4, $5 FROM project p
+			ON CONFLICT (project_id) DO UPDATE SET requests_per_minute=EXCLUDED.requests_per_minute,
+				tokens_per_minute=EXCLUDED.tokens_per_minute, concurrent_requests=EXCLUDED.concurrent_requests,
+				monthly_spend_micros=EXCLUDED.monthly_spend_micros, updated_at=now()
+			RETURNING tenant_id, project_id, requests_per_minute, tokens_per_minute, concurrent_requests, monthly_spend_micros, updated_at
+		)
+		SELECT u.tenant_id::text, u.project_id::text, p.project_name, p.tenant_name,
+			u.requests_per_minute, u.tokens_per_minute, u.concurrent_requests, u.monthly_spend_micros, u.updated_at
+		FROM updated u JOIN project p ON p.id=u.project_id`,
+		projectID, input.RequestsPerMinute, input.TokensPerMinute, input.ConcurrentRequests, input.MonthlySpendMicros,
+	).Scan(&result.TenantID, &result.ProjectID, &result.ProjectName, &result.TenantName,
+		&result.RequestsPerMinute, &result.TokensPerMinute, &result.Concurrent, &result.MonthlySpendMicros, &result.UpdatedAt)
+	if errors.Is(err, pgx.ErrNoRows) {
+		return ProjectLimit{}, 0, ErrNotFound
+	}
+	if err != nil {
+		return ProjectLimit{}, 0, fmt.Errorf("set project limit: %w", err)
+	}
+	generation, err := bumpGeneration(ctx, tx)
+	if err != nil {
+		return ProjectLimit{}, 0, err
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return ProjectLimit{}, 0, err
+	}
+	return result, generation, nil
+}