Просмотр исходного кода

add usage enforcement to backend, basic plan by default

Alexander Belanger 4 лет назад
Родитель
Сommit
8155d7144a

+ 147 - 0
api/server/handlers/project/get_usage.go

@@ -0,0 +1,147 @@
+package project
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+
+	"github.com/porter-dev/porter/api/server/authz"
+	"github.com/porter-dev/porter/api/server/handlers"
+	"github.com/porter-dev/porter/api/server/shared"
+	"github.com/porter-dev/porter/api/server/shared/apierrors"
+	"github.com/porter-dev/porter/api/server/shared/config"
+	"github.com/porter-dev/porter/api/types"
+	"github.com/porter-dev/porter/internal/kubernetes"
+	"github.com/porter-dev/porter/internal/kubernetes/nodes"
+	"github.com/porter-dev/porter/internal/models"
+	"gorm.io/gorm"
+)
+
+type ProjectGetUsageHandler struct {
+	handlers.PorterHandlerWriter
+}
+
+func NewProjectGetUsageHandler(
+	config *config.Config,
+	writer shared.ResultWriter,
+) *ProjectGetUsageHandler {
+	return &ProjectGetUsageHandler{
+		PorterHandlerWriter: handlers.NewDefaultPorterHandler(config, nil, writer),
+	}
+}
+
+func (p *ProjectGetUsageHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	res := &types.GetProjectUsageResponse{}
+
+	currUsage, limit, err := GetUsage(p.Config(), r)
+
+	if err != nil {
+		p.HandleAPIError(w, r, apierrors.NewErrInternal(err))
+		return
+	}
+
+	res.Current = *currUsage
+	res.Limit = *limit
+
+	p.WriteResult(w, r, res)
+}
+
+// GetUsage gets a project's current usage and usage limit
+func GetUsage(config *config.Config, r *http.Request) (
+	current, limit *types.ProjectUsage,
+	err error,
+) {
+	// read the project from the request
+	proj, _ := r.Context().Value(types.ProjectScope).(*models.Project)
+
+	// query for the project limit; if not found, default to basic
+	limitModel, err := config.Repo.ProjectUsage().ReadProjectUsage(proj.ID)
+
+	if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
+		copyBasic := types.BasicPlan
+		limit = &copyBasic
+	} else if err != nil {
+		return nil, nil, err
+	} else {
+		limit = limitModel.ToProjectUsageType()
+	}
+
+	// query for the linked cluster counts
+	clusters, err := config.Repo.Cluster().ListClustersByProjectID(proj.ID)
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// query for the linked user counts
+	roles, err := config.Repo.Project().ListProjectRoles(proj.ID)
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	usageCache, err := config.Repo.ProjectUsage().ReadProjectUsageCache(proj.ID)
+	isCacheFound := true
+
+	if isCacheFound = !errors.Is(err, gorm.ErrRecordNotFound); err != nil && isCacheFound {
+		return nil, nil, err
+	}
+
+	// if the usage cache is 24 hours old, was not found, or usage is over limit,
+	// re-query for the usage
+	if !isCacheFound || usageCache.Is24HrOld() || usageCache.ResourceMemory > limit.ResourceMemory || usageCache.ResourceCPU > limit.ResourceCPU {
+		cpu, memory, err := getResourceUsage(config, clusters)
+
+		if err != nil {
+			return nil, nil, err
+		}
+
+		if !isCacheFound {
+			usageCache, err = config.Repo.ProjectUsage().CreateProjectUsageCache(&models.ProjectUsageCache{
+				ProjectID:      proj.ID,
+				ResourceCPU:    cpu,
+				ResourceMemory: memory,
+			})
+		} else {
+			usageCache.ResourceCPU = cpu
+			usageCache.ResourceMemory = memory
+
+			usageCache, err = config.Repo.ProjectUsage().UpdateProjectUsageCache(usageCache)
+		}
+	}
+
+	return &types.ProjectUsage{
+		ResourceCPU:    usageCache.ResourceCPU,
+		ResourceMemory: usageCache.ResourceMemory,
+		Clusters:       uint(len(clusters)),
+		Users:          uint(len(roles)),
+	}, limit, nil
+}
+
+// gets the total resource usage across all nodes in all clusters
+func getResourceUsage(config *config.Config, clusters []*models.Cluster) (uint, uint, error) {
+	// TODO; pass this in?
+	var totCPU, totMem uint = 0, 0
+	getter := authz.NewOutOfClusterAgentGetter(config)
+
+	for _, cluster := range clusters {
+		ooc := getter.GetOutOfClusterConfig(cluster)
+
+		agent, err := kubernetes.GetAgentOutOfClusterConfig(ooc)
+
+		if err != nil {
+			return 0, 0, fmt.Errorf("failed to get agent: %s", err.Error())
+		}
+
+		totAlloc, err := nodes.GetAllocatableResources(agent.Clientset)
+
+		if err != nil {
+			return 0, 0, fmt.Errorf("failed to get alloc: %s", err.Error())
+		}
+
+		totCPU += totAlloc.CPU
+		totMem += totAlloc.Memory
+	}
+
+	return totCPU / 1000, totMem / (1000 * 1000), nil
+}

+ 2 - 0
api/server/router/cluster.go

@@ -153,6 +153,8 @@ func getClusterRoutes(
 				types.UserScope,
 				types.ProjectScope,
 			},
+			CheckUsage:  true,
+			UsageMetric: types.Clusters,
 		},
 	)
 

+ 8 - 0
api/server/router/invite.go

@@ -64,6 +64,7 @@ func getInviteRoutes(
 			Scopes: []types.PermissionScope{
 				types.UserScope,
 				types.ProjectScope,
+				types.SettingsScope,
 			},
 		},
 	)
@@ -91,7 +92,10 @@ func getInviteRoutes(
 			Scopes: []types.PermissionScope{
 				types.UserScope,
 				types.ProjectScope,
+				types.SettingsScope,
 			},
+			CheckUsage:  true,
+			UsageMetric: types.Users,
 		},
 	)
 
@@ -122,6 +126,8 @@ func getInviteRoutes(
 				types.UserScope,
 			},
 			ShouldRedirect: true,
+			CheckUsage:     true,
+			UsageMetric:    types.Users,
 		},
 	)
 
@@ -145,6 +151,7 @@ func getInviteRoutes(
 			Scopes: []types.PermissionScope{
 				types.UserScope,
 				types.ProjectScope,
+				types.SettingsScope,
 				types.InviteScope,
 			},
 		},
@@ -173,6 +180,7 @@ func getInviteRoutes(
 			Scopes: []types.PermissionScope{
 				types.UserScope,
 				types.ProjectScope,
+				types.SettingsScope,
 				types.InviteScope,
 			},
 		},

+ 93 - 0
api/server/router/middleware/usage.go

@@ -0,0 +1,93 @@
+package middleware
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/porter-dev/porter/api/server/handlers/project"
+	"github.com/porter-dev/porter/api/server/shared/apierrors"
+	"github.com/porter-dev/porter/api/server/shared/config"
+	"github.com/porter-dev/porter/api/types"
+)
+
+type UsageMiddleware struct {
+	config *config.Config
+	metric types.UsageMetric
+}
+
+func NewUsageMiddleware(config *config.Config, metric types.UsageMetric) *UsageMiddleware {
+	return &UsageMiddleware{config, metric}
+}
+
+var UsageErrFmt = "usage limit reached for metric %s: limit %d, requested %d"
+
+func (b *UsageMiddleware) Middleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// get the project usage limits
+		currentUsage, limit, err := project.GetUsage(b.config, r)
+
+		if err != nil {
+			apierrors.HandleAPIError(
+				b.config,
+				w, r,
+				apierrors.NewErrInternal(err),
+				true,
+			)
+
+			return
+		}
+
+		// check the usage limits
+		allowed := allowUsage(limit, currentUsage, b.metric)
+
+		if allowed {
+			next.ServeHTTP(w, r)
+		} else {
+			limit, curr := getMetricUsage(limit, currentUsage, b.metric)
+
+			apierrors.HandleAPIError(
+				b.config,
+				w, r,
+				apierrors.NewErrPassThroughToClient(
+					fmt.Errorf(UsageErrFmt, b.metric, limit, curr),
+					http.StatusBadRequest,
+				),
+				true,
+			)
+		}
+	})
+}
+
+// checkUsage returns true if the increase in usage is allowed for the given metric,
+// false otherwise. We only assume increments of 1 in usage for now.
+func allowUsage(
+	plan, current *types.ProjectUsage,
+	metric types.UsageMetric,
+) bool {
+	switch metric {
+	case types.Users:
+		return plan.Users > current.Users+1
+	case types.Clusters:
+		return plan.Clusters > current.Clusters+1
+	default:
+		return false
+	}
+}
+
+func getMetricUsage(
+	plan, current *types.ProjectUsage,
+	metric types.UsageMetric,
+) (limit uint, curr uint) {
+	switch metric {
+	case types.CPU:
+		return plan.ResourceCPU, current.ResourceCPU
+	case types.Memory:
+		return plan.ResourceMemory, current.ResourceMemory
+	case types.Users:
+		return plan.Users, current.Users
+	case types.Clusters:
+		return plan.Users, current.Users
+	default:
+		return 0, 0
+	}
+}

+ 33 - 0
api/server/router/project.go

@@ -137,6 +137,33 @@ func getProjectRoutes(
 		Router:   r,
 	})
 
+	// GET /api/projects/{project_id}/usage -> project.NewProjectGetUsageHandler
+	getUsageEndpoint := factory.NewAPIEndpoint(
+		&types.APIRequestMetadata{
+			Verb:   types.APIVerbGet,
+			Method: types.HTTPVerbGet,
+			Path: &types.Path{
+				Parent:       basePath,
+				RelativePath: relPath + "/usage",
+			},
+			Scopes: []types.PermissionScope{
+				types.UserScope,
+				types.ProjectScope,
+			},
+		},
+	)
+
+	getUsageHandler := project.NewProjectGetUsageHandler(
+		config,
+		factory.GetResultWriter(),
+	)
+
+	routes = append(routes, &Route{
+		Endpoint: getUsageEndpoint,
+		Handler:  getUsageHandler,
+		Router:   r,
+	})
+
 	// GET /api/projects/{project_id}/clusters -> cluster.NewClusterListHandler
 	listClusterEndpoint := factory.NewAPIEndpoint(
 		&types.APIRequestMetadata{
@@ -509,6 +536,8 @@ func getProjectRoutes(
 				types.UserScope,
 				types.ProjectScope,
 			},
+			CheckUsage:  true,
+			UsageMetric: types.Clusters,
 		},
 	)
 
@@ -565,6 +594,8 @@ func getProjectRoutes(
 				types.UserScope,
 				types.ProjectScope,
 			},
+			CheckUsage:  true,
+			UsageMetric: types.Clusters,
 		},
 	)
 
@@ -621,6 +652,8 @@ func getProjectRoutes(
 				types.UserScope,
 				types.ProjectScope,
 			},
+			CheckUsage:  true,
+			UsageMetric: types.Clusters,
 		},
 	)
 

+ 6 - 0
api/server/router/router.go

@@ -233,6 +233,12 @@ func registerRoutes(config *config.Config, routes []*Route) {
 			atomicGroup.Use(websocketMw.Middleware)
 		}
 
+		if route.Endpoint.Metadata.CheckUsage {
+			usageMW := middleware.NewUsageMiddleware(config, route.Endpoint.Metadata.UsageMetric)
+
+			atomicGroup.Use(usageMW.Middleware)
+		}
+
 		atomicGroup.Method(
 			string(route.Endpoint.Metadata.Method),
 			route.Endpoint.Metadata.Path.RelativePath,

+ 6 - 0
api/types/request.go

@@ -63,6 +63,12 @@ type APIRequestMetadata struct {
 
 	// Whether the endpoint upgrades to a websocket
 	IsWebsocket bool
+
+	// Whether the endpoint should check for a usage limit
+	CheckUsage bool
+
+	// The usage metric that the request should check for, if CheckUsage
+	UsageMetric UsageMetric
 }
 
 const RequestScopeCtxKey = "requestscopes"

+ 61 - 0
api/types/usage.go

@@ -0,0 +1,61 @@
+package types
+
+type UsageMetric string
+
+const (
+	CPU      UsageMetric = "cpu"
+	Memory   UsageMetric = "memory"
+	Clusters UsageMetric = "clusters"
+	Users    UsageMetric = "users"
+)
+
+type ProjectUsage struct {
+	// The CPU usage, in vCPUs
+	ResourceCPU uint `json:"resource_cpu"`
+
+	// The memory usage, in mibibytes (?)
+	ResourceMemory uint `json:"resource_memory"`
+
+	// The number of clusters
+	Clusters uint `json:"clusters"`
+
+	// The number of users
+	Users uint `json:"users"`
+}
+
+var BasicPlan = ProjectUsage{
+	ResourceCPU: 10,
+	// 20 GB converted to Mebibytes
+	ResourceMemory: 19074,
+	Clusters:       1,
+	Users:          1,
+}
+
+var TeamPlan = ProjectUsage{
+	ResourceCPU: 20,
+	// 40 GB converted to Mebibytes
+	ResourceMemory: 38148,
+	Clusters:       3,
+	Users:          3,
+}
+
+var GrowthPlan = ProjectUsage{
+	ResourceCPU: 80,
+	// 160 GB converted to Mebibytes
+	ResourceMemory: 152592,
+	Clusters:       0,
+	Users:          5,
+}
+
+// all unlimited
+var EnterprisePlan = ProjectUsage{
+	ResourceCPU:    0,
+	ResourceMemory: 0,
+	Clusters:       0,
+	Users:          0,
+}
+
+type GetProjectUsageResponse struct {
+	Current ProjectUsage `json:"current"`
+	Limit   ProjectUsage `json:"limit"`
+}

+ 15 - 3
dashboard/src/main/home/project-settings/InviteList.tsx

@@ -25,7 +25,7 @@ export type Collaborator = {
 };
 
 const InvitePage: React.FunctionComponent<Props> = ({}) => {
-  const { currentProject, setCurrentModal, user } = useContext(Context);
+  const { currentProject, setCurrentModal, setCurrentError, user } = useContext(Context);
   const [isLoading, setIsLoading] = useState(true);
   const [invites, setInvites] = useState<Array<InviteType>>([]);
   const [email, setEmail] = useState("");
@@ -115,7 +115,13 @@ const InvitePage: React.FunctionComponent<Props> = ({}) => {
         getData();
         setEmail("");
       })
-      .catch((err) => console.log(err));
+      .catch((err) => {
+        if (err.response.data?.error) {
+          setCurrentError(err.response.data?.error)
+        }
+
+        console.log(err)
+      });
   };
 
   const deleteInvite = (inviteId: number) => {
@@ -154,7 +160,13 @@ const InvitePage: React.FunctionComponent<Props> = ({}) => {
         )
       )
       .then(getData)
-      .catch((err) => console.log(err));
+      .catch((err) => {
+        if (err.response.data?.error) {
+          setCurrentError(err.response.data?.error)
+        }
+
+        console.log(err)
+      });
   };
 
   const validateEmail = () => {

+ 28 - 0
internal/kubernetes/nodes/nodes.go

@@ -10,6 +10,34 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
+type TotalAllocatable struct {
+	CPU    uint
+	Memory uint
+}
+
+func GetAllocatableResources(clientset kubernetes.Interface) (*TotalAllocatable, error) {
+	nodeList, err := clientset.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
+
+	if err != nil {
+		return nil, err
+	}
+
+	var totCPU uint64 = 0
+	var totMem uint64 = 0
+
+	for _, node := range nodeList.Items {
+		capac := node.Status.Allocatable
+
+		totCPU += uint64(capac.Cpu().MilliValue())
+		totMem += capac.Memory().AsDec().UnscaledBig().Uint64()
+	}
+
+	return &TotalAllocatable{
+		CPU:    uint(totCPU),
+		Memory: uint(totMem),
+	}, nil
+}
+
 type NodeUsage struct {
 	cpuReqs                        string
 	memoryReqs                     string

+ 12 - 0
internal/models/project.go

@@ -7,6 +7,15 @@ import (
 	ints "github.com/porter-dev/porter/internal/models/integrations"
 )
 
+type ProjectPlan string
+
+const (
+	ProjectPlanBasic      ProjectPlan = "basic"
+	ProjectPlanTeam       ProjectPlan = "team"
+	ProjectPlanGrowth     ProjectPlan = "growth"
+	ProjectPlanEnterprise ProjectPlan = "enterprise"
+)
+
 // Project type that extends gorm.Model
 type Project struct {
 	gorm.Model
@@ -14,6 +23,9 @@ type Project struct {
 	Name  string `json:"name"`
 	Roles []Role `json:"roles"`
 
+	ProjectUsageID      uint
+	ProjectUsageCacheID uint
+
 	// linked repos
 	GitRepos []GitRepo `json:"git_repos,omitempty"`
 

+ 58 - 0
internal/models/usage.go

@@ -0,0 +1,58 @@
+package models
+
+import (
+	"time"
+
+	"github.com/porter-dev/porter/api/types"
+	"gorm.io/gorm"
+)
+
+// ProjectUsage keeps track of the usage limits for a project
+type ProjectUsage struct {
+	gorm.Model
+
+	// The project ID that this model refers to
+	ProjectID uint
+
+	// The CPU usage, in vCPUs
+	ResourceCPU uint
+
+	// The memory usage, in bytes
+	ResourceMemory uint
+
+	// The number of clusters
+	Clusters uint
+
+	// The number of users
+	Users uint
+}
+
+// ToProjectUsageType converts the project usage model to a project usage API type
+func (p *ProjectUsage) ToProjectUsageType() *types.ProjectUsage {
+	return &types.ProjectUsage{
+		ResourceCPU:    p.ResourceCPU,
+		ResourceMemory: p.ResourceMemory,
+		Clusters:       p.Clusters,
+		Users:          p.Users,
+	}
+}
+
+// ProjectUsageCache stores the latest cache of the resource usage for a project,
+// for fields that are expensive to compute
+type ProjectUsageCache struct {
+	gorm.Model
+
+	// The project ID that this model refers to
+	ProjectID uint
+
+	// The CPU usage, in vCPUs
+	ResourceCPU uint
+
+	// The memory usage, in bytes
+	ResourceMemory uint
+}
+
+func (p *ProjectUsageCache) Is24HrOld() bool {
+	timeSince := time.Now().Sub(p.UpdatedAt)
+	return timeSince > 24*time.Hour
+}

+ 2 - 0
internal/repository/gorm/migrate.go

@@ -29,6 +29,8 @@ func AutoMigrate(db *gorm.DB) error {
 		&models.NotificationConfig{},
 		&models.EventContainer{},
 		&models.SubEvent{},
+		&models.ProjectUsage{},
+		&models.ProjectUsageCache{},
 		&ints.KubeIntegration{},
 		&ints.BasicIntegration{},
 		&ints.OIDCIntegration{},

+ 6 - 0
internal/repository/gorm/repository.go

@@ -31,6 +31,7 @@ type GormRepository struct {
 	slackIntegration          repository.SlackIntegrationRepository
 	notificationConfig        repository.NotificationConfigRepository
 	event                     repository.EventRepository
+	projectUsage              repository.ProjectUsageRepository
 }
 
 func (t *GormRepository) User() repository.UserRepository {
@@ -133,6 +134,10 @@ func (t *GormRepository) Event() repository.EventRepository {
 	return t.event
 }
 
+func (t *GormRepository) ProjectUsage() repository.ProjectUsageRepository {
+	return t.projectUsage
+}
+
 // NewRepository returns a Repository which persists users in memory
 // and accepts a parameter that can trigger read/write errors
 func NewRepository(db *gorm.DB, key *[32]byte) repository.Repository {
@@ -162,5 +167,6 @@ func NewRepository(db *gorm.DB, key *[32]byte) repository.Repository {
 		slackIntegration:          NewSlackIntegrationRepository(db, key),
 		notificationConfig:        NewNotificationConfigRepository(db),
 		event:                     NewEventRepository(db),
+		projectUsage:              NewProjectUsageRepository(db),
 	}
 }

+ 87 - 0
internal/repository/gorm/usage.go

@@ -0,0 +1,87 @@
+package gorm
+
+import (
+	"github.com/porter-dev/porter/internal/models"
+	"github.com/porter-dev/porter/internal/repository"
+	"gorm.io/gorm"
+)
+
+// ProjectUsageRepository implements repository.ProjectUsageRepository
+type ProjectUsageRepository struct {
+	db *gorm.DB
+}
+
+// NewProjectUsageRepository will return errors if canQuery is false
+func NewProjectUsageRepository(db *gorm.DB) repository.ProjectUsageRepository {
+	return &ProjectUsageRepository{db}
+}
+
+// CreateProjectUsage creates a new project usage limit
+func (repo *ProjectUsageRepository) CreateProjectUsage(
+	usage *models.ProjectUsage,
+) (*models.ProjectUsage, error) {
+	if err := repo.db.Create(usage).Error; err != nil {
+		return nil, err
+	}
+
+	return usage, nil
+}
+
+// ReadProjectUsage finds the project usage matching a project ID
+func (repo *ProjectUsageRepository) ReadProjectUsage(
+	projID uint,
+) (*models.ProjectUsage, error) {
+	res := &models.ProjectUsage{}
+
+	if err := repo.db.Where("project_id = ?", projID).First(res).Error; err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// UpdateProjectUsage modifies an existing ProjectUsage in the database
+func (repo *ProjectUsageRepository) UpdateProjectUsage(
+	usage *models.ProjectUsage,
+) (*models.ProjectUsage, error) {
+	if err := repo.db.Save(usage).Error; err != nil {
+		return nil, err
+	}
+
+	return usage, nil
+}
+
+// CreateProjectUsageCache creates a new project usage cache
+func (repo *ProjectUsageRepository) CreateProjectUsageCache(
+	cache *models.ProjectUsageCache,
+) (*models.ProjectUsageCache, error) {
+	if err := repo.db.Create(cache).Error; err != nil {
+		return nil, err
+	}
+
+	return cache, nil
+}
+
+// ReadProjectUsageCache finds the project usage cache matching a project ID
+func (repo *ProjectUsageRepository) ReadProjectUsageCache(
+	projID uint,
+) (*models.ProjectUsageCache, error) {
+	res := &models.ProjectUsageCache{}
+
+	if err := repo.db.Where("project_id = ?", projID).First(res).Error; err != nil {
+		return nil, err
+	}
+
+	return res, nil
+}
+
+// UpdateProjectUsageCache modifies an existing ProjectUsageCache in the database
+func (repo *ProjectUsageRepository) UpdateProjectUsageCache(
+	cache *models.ProjectUsageCache,
+) (*models.ProjectUsageCache, error) {
+	if err := repo.db.Save(cache).Error; err != nil {
+		return nil, err
+	}
+
+	return cache, nil
+}

+ 2 - 1
internal/repository/repository.go

@@ -25,5 +25,6 @@ type Repository interface {
 	GithubAppOAuthIntegration() GithubAppOAuthIntegrationRepository
 	SlackIntegration() SlackIntegrationRepository
 	NotificationConfig() NotificationConfigRepository
-	Event()                     EventRepository
+	Event() EventRepository
+	ProjectUsage() ProjectUsageRepository
 }

+ 5 - 0
internal/repository/test/repository.go

@@ -30,6 +30,7 @@ type TestRepository struct {
 	slackIntegration          repository.SlackIntegrationRepository
 	notificationConfig        repository.NotificationConfigRepository
 	event                     repository.EventRepository
+	projectUsage              repository.ProjectUsageRepository
 }
 
 func (t *TestRepository) User() repository.UserRepository {
@@ -132,6 +133,10 @@ func (t *TestRepository) Event() repository.EventRepository {
 	return t.event
 }
 
+func (t *TestRepository) ProjectUsage() repository.ProjectUsageRepository {
+	return t.projectUsage
+}
+
 // NewRepository returns a Repository which persists users in memory
 // and accepts a parameter that can trigger read/write errors
 func NewRepository(canQuery bool, failingMethods ...string) repository.Repository {

+ 129 - 0
internal/repository/test/usage.go

@@ -0,0 +1,129 @@
+package test
+
+import (
+	"errors"
+
+	"github.com/porter-dev/porter/internal/models"
+	"github.com/porter-dev/porter/internal/repository"
+	"gorm.io/gorm"
+)
+
+// ProjectUsageRepository implements repository.ProjectUsageRepository
+type ProjectUsageRepository struct {
+	canQuery bool
+	usages   []*models.ProjectUsage
+	caches   []*models.ProjectUsageCache
+}
+
+// NewProjectUsageRepository will return errors if canQuery is false
+func NewProjectUsageRepository(canQuery bool) repository.ProjectUsageRepository {
+	return &ProjectUsageRepository{
+		canQuery,
+		[]*models.ProjectUsage{},
+		[]*models.ProjectUsageCache{},
+	}
+}
+
+// CreateProjectUsage creates a new project usage limit
+func (repo *ProjectUsageRepository) CreateProjectUsage(
+	usage *models.ProjectUsage,
+) (*models.ProjectUsage, error) {
+	if !repo.canQuery {
+		return nil, errors.New("Cannot write database")
+	}
+
+	if usage == nil {
+		return nil, nil
+	}
+
+	repo.usages = append(repo.usages, usage)
+
+	return usage, nil
+}
+
+// CreateProjectUsage reads a project usage by project id
+func (repo *ProjectUsageRepository) ReadProjectUsage(
+	projID uint,
+) (*models.ProjectUsage, error) {
+	if !repo.canQuery {
+		return nil, errors.New("Cannot read from database")
+	}
+
+	for _, pu := range repo.usages {
+		if pu != nil && pu.ProjectID == projID {
+			return pu, nil
+		}
+	}
+
+	return nil, gorm.ErrRecordNotFound
+}
+
+// UpdateProjectUsage modifies an existing ProjectUsage in the database
+func (repo *ProjectUsageRepository) UpdateProjectUsage(
+	usage *models.ProjectUsage,
+) (*models.ProjectUsage, error) {
+	if !repo.canQuery {
+		return nil, errors.New("Cannot write database")
+	}
+
+	if int(usage.ID-1) >= len(repo.usages) || repo.usages[usage.ID-1] == nil {
+		return nil, gorm.ErrRecordNotFound
+	}
+
+	index := int(usage.ID - 1)
+	repo.usages[index] = usage
+
+	return usage, nil
+}
+
+// CreateProjectUsageCache creates a new project usage cache
+func (repo *ProjectUsageRepository) CreateProjectUsageCache(
+	cache *models.ProjectUsageCache,
+) (*models.ProjectUsageCache, error) {
+	if !repo.canQuery {
+		return nil, errors.New("Cannot write database")
+	}
+
+	if cache == nil {
+		return nil, nil
+	}
+
+	repo.caches = append(repo.caches, cache)
+
+	return cache, nil
+}
+
+// CreateProjectUsageCache reads a project usage by project id
+func (repo *ProjectUsageRepository) ReadProjectUsageCache(
+	projID uint,
+) (*models.ProjectUsageCache, error) {
+	if !repo.canQuery {
+		return nil, errors.New("Cannot read from database")
+	}
+
+	for _, puc := range repo.caches {
+		if puc != nil && puc.ProjectID == projID {
+			return puc, nil
+		}
+	}
+
+	return nil, gorm.ErrRecordNotFound
+}
+
+// UpdateProjectUsageCache modifies an existing ProjectUsageCache in the database
+func (repo *ProjectUsageRepository) UpdateProjectUsageCache(
+	cache *models.ProjectUsageCache,
+) (*models.ProjectUsageCache, error) {
+	if !repo.canQuery {
+		return nil, errors.New("Cannot write database")
+	}
+
+	if int(cache.ID-1) >= len(repo.caches) || repo.usages[cache.ID-1] == nil {
+		return nil, gorm.ErrRecordNotFound
+	}
+
+	index := int(cache.ID - 1)
+	repo.caches[index] = cache
+
+	return cache, nil
+}

+ 13 - 0
internal/repository/usage.go

@@ -0,0 +1,13 @@
+package repository
+
+import "github.com/porter-dev/porter/internal/models"
+
+// ProjectUsageRepository represents the set of queries on the ProjectUsage model
+type ProjectUsageRepository interface {
+	CreateProjectUsage(usage *models.ProjectUsage) (*models.ProjectUsage, error)
+	ReadProjectUsage(projID uint) (*models.ProjectUsage, error)
+	UpdateProjectUsage(cache *models.ProjectUsage) (*models.ProjectUsage, error)
+	CreateProjectUsageCache(cache *models.ProjectUsageCache) (*models.ProjectUsageCache, error)
+	ReadProjectUsageCache(projID uint) (*models.ProjectUsageCache, error)
+	UpdateProjectUsageCache(cache *models.ProjectUsageCache) (*models.ProjectUsageCache, error)
+}