Skip to content

Commit dd8cc31

Browse files
jakubnojakubno
committed
Add addons table
1 parent dfcfcb6 commit dd8cc31

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

67 files changed

+633
-4592
lines changed

‎packages/api/internal/auth/middleware.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@ import (
1515
"go.uber.org/zap"
1616

1717
"github.com/e2b-dev/infra/packages/api/internal/api"
18-
authcache "github.com/e2b-dev/infra/packages/api/internal/cache/auth"
1918
"github.com/e2b-dev/infra/packages/api/internal/cfg"
2019
"github.com/e2b-dev/infra/packages/api/internal/db"
20+
"github.com/e2b-dev/infra/packages/api/internal/db/types"
2121
"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
2222
)
2323

@@ -134,13 +134,13 @@ func adminValidationFunction(adminToken string) func(context.Context, string) (s
134134

135135
func CreateAuthenticationFunc(
136136
config cfg.Config,
137-
teamValidationFunction func(context.Context, string) (authcache.AuthTeamInfo, *api.APIError),
137+
teamValidationFunction func(context.Context, string) (*types.Team, *api.APIError),
138138
userValidationFunction func(context.Context, string) (uuid.UUID, *api.APIError),
139139
supabaseTokenValidationFunction func(context.Context, string) (uuid.UUID, *api.APIError),
140-
supabaseTeamValidationFunction func(context.Context, string) (authcache.AuthTeamInfo, *api.APIError),
140+
supabaseTeamValidationFunction func(context.Context, string) (*types.Team, *api.APIError),
141141
) openapi3filter.AuthenticationFunc {
142142
authenticators := []authenticator{
143-
&commonAuthenticator[authcache.AuthTeamInfo]{
143+
&commonAuthenticator[*types.Team]{
144144
securitySchemeName: "ApiKeyAuth",
145145
headerKey: headerKey{
146146
name: "X-API-Key",
@@ -173,7 +173,7 @@ func CreateAuthenticationFunc(
173173
contextKey: UserIDContextKey,
174174
errorMessage: "Invalid Supabase token.",
175175
},
176-
&commonAuthenticator[authcache.AuthTeamInfo]{
176+
&commonAuthenticator[*types.Team]{
177177
securitySchemeName: "Supabase2TeamAuth",
178178
headerKey: headerKey{
179179
name: "X-Supabase-Team",

‎packages/api/internal/cache/auth/cache.go‎

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,28 +8,22 @@ import (
88
"github.com/jellydator/ttlcache/v3"
99
"golang.org/x/sync/singleflight"
1010

11-
"github.com/e2b-dev/infra/packages/db/queries"
11+
"github.com/e2b-dev/infra/packages/api/internal/db/types"
1212
)
1313

1414
const (
1515
authInfoExpiration = 5 * time.Minute
1616
refreshInterval = 1 * time.Minute
1717
)
1818

19-
type AuthTeamInfo struct {
20-
Team *queries.Team
21-
Tier *queries.Tier
22-
}
23-
2419
type TeamInfo struct {
25-
team *queries.Team
26-
tier *queries.Tier
20+
team *types.Team
2721

2822
lastRefresh time.Time
2923
once singleflight.Group
3024
}
3125

32-
type DataCallback = func(ctx context.Context, key string) (*queries.Team, *queries.Tier, error)
26+
type DataCallback = func(ctx context.Context, key string) (*types.Team, error)
3327

3428
type TeamAuthCache struct {
3529
cache *ttlcache.Cache[string, *TeamInfo]
@@ -45,21 +39,21 @@ func NewTeamAuthCache() *TeamAuthCache {
4539
}
4640

4741
// TODO: save blocked teams to cache as well, handle the condition in the GetOrSet method
48-
func (c *TeamAuthCache) GetOrSet(ctx context.Context, key string, dataCallback DataCallback) (team *queries.Team, tier *queries.Tier, err error) {
42+
func (c *TeamAuthCache) GetOrSet(ctx context.Context, key string, dataCallback DataCallback) (team *types.Team, err error) {
4943
var item *ttlcache.Item[string, *TeamInfo]
5044
var templateInfo *TeamInfo
5145

5246
item = c.cache.Get(key)
5347
if item == nil {
54-
team, tier, err = dataCallback(ctx, key)
48+
team, err = dataCallback(ctx, key)
5549
if err != nil {
56-
return nil, nil, fmt.Errorf("error while getting the team: %w", err)
50+
return nil, fmt.Errorf("error while getting the team: %w", err)
5751
}
5852

59-
templateInfo = &TeamInfo{team: team, tier: tier, lastRefresh: time.Now()}
53+
templateInfo = &TeamInfo{team: team, lastRefresh: time.Now()}
6054
c.cache.Set(key, templateInfo, authInfoExpiration)
6155

62-
return team, tier, nil
56+
return team, nil
6357
}
6458

6559
templateInfo = item.Value()
@@ -70,20 +64,20 @@ func (c *TeamAuthCache) GetOrSet(ctx context.Context, key string, dataCallback D
7064
})
7165
}
7266

73-
return templateInfo.team, templateInfo.tier, nil
67+
return templateInfo.team, nil
7468
}
7569

7670
// Refresh refreshes the cache for the given team ID.
7771
func (c *TeamAuthCache) Refresh(key string, dataCallback DataCallback) {
7872
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
7973
defer cancel()
8074

81-
team, tier, err := dataCallback(ctx, key)
75+
team, err := dataCallback(ctx, key)
8276
if err != nil {
8377
c.cache.Delete(key)
8478

8579
return
8680
}
8781

88-
c.cache.Set(key, &TeamInfo{team: team, tier: tier, lastRefresh: time.Now()}, authInfoExpiration)
82+
c.cache.Set(key, &TeamInfo{team: team, lastRefresh: time.Now()}, authInfoExpiration)
8983
}

‎packages/api/internal/db/apikeys.go‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"fmt"
66

7+
"github.com/e2b-dev/infra/packages/api/internal/db/types"
78
sqlcdb "github.com/e2b-dev/infra/packages/db/client"
89
"github.com/e2b-dev/infra/packages/db/queries"
910
)
@@ -36,18 +37,27 @@ func validateTeamUsage(team queries.Team) error {
3637
return nil
3738
}
3839

39-
func GetTeamAuth(ctx context.Context, db *sqlcdb.Client, apiKey string) (*queries.Team, *queries.Tier, error) {
40+
func GetTeamAuth(ctx context.Context, db *sqlcdb.Client, apiKey string) (*types.Team, error) {
4041
result, err := db.GetTeamWithTierByAPIKeyWithUpdateLastUsed(ctx, apiKey)
4142
if err != nil {
4243
errMsg := fmt.Errorf("failed to get team from API key: %w", err)
4344

44-
return nil, nil, errMsg
45+
return nil, errMsg
4546
}
4647

4748
err = validateTeamUsage(result.Team)
4849
if err != nil {
49-
return nil, nil, err
50+
return nil, err
5051
}
5152

52-
return &result.Team, &result.Tier, nil
53+
team := types.NewTeam(
54+
&result.Team,
55+
&result.Tier,
56+
result.ExtraConcurrentSandboxes,
57+
result.ExtraConcurrentTemplateBuilds,
58+
result.ExtraMaxVcpu,
59+
result.ExtraMaxRamMb,
60+
result.ExtraDiskMb,
61+
)
62+
return team, nil
5363
}

‎packages/api/internal/db/teams.go‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
package db
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
"github.com/google/uuid"
8+
9+
"github.com/e2b-dev/infra/packages/api/internal/db/types"
10+
sqlcdb "github.com/e2b-dev/infra/packages/db/client"
11+
)
12+
13+
func GetTeamByUser(ctx context.Context, db *sqlcdb.Client, userID uuid.UUID) ([]*types.TeamWithDefault, error) {
14+
teams, err := db.GetTeamsWithUsersTeamsWithTier(ctx, userID)
15+
if err != nil {
16+
return nil, fmt.Errorf("error when getting default team: %w", err)
17+
}
18+
19+
teamsWithLimits := make([]*types.TeamWithDefault, 0, len(teams))
20+
for _, team := range teams {
21+
t := types.NewTeam(
22+
&team.Team,
23+
&team.Tier,
24+
team.ExtraConcurrentSandboxes,
25+
team.ExtraConcurrentTemplateBuilds,
26+
team.ExtraMaxVcpu,
27+
team.ExtraMaxRamMb,
28+
team.ExtraDiskMb,
29+
)
30+
teamsWithLimits = append(teamsWithLimits, &types.TeamWithDefault{
31+
Team: t,
32+
IsDefault: team.UsersTeam.IsDefault,
33+
})
34+
}
35+
36+
return teamsWithLimits, nil
37+
}

‎packages/api/internal/db/types/limits.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package types
2+
3+
type TeamLimits struct {
4+
SandboxConcurrency int64
5+
BuildConcurrency int64
6+
MaxLengthHours int64
7+
8+
MaxVcpu int64
9+
MaxRamMb int64
10+
DiskMb int64
11+
}

‎packages/api/internal/db/types/teams.go‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package types
2+
3+
import (
4+
"github.com/e2b-dev/infra/packages/db/queries"
5+
)
6+
7+
type Team struct {
8+
*queries.Team
9+
10+
Limits *TeamLimits
11+
}
12+
13+
func newTeamLimits(
14+
tier *queries.Tier,
15+
extraMaxSandboxes int64,
16+
extraMaxConcurrentBuilds int64,
17+
extraMaxVcpu int64,
18+
extraMaxRamMb int64,
19+
addonsDiskMb int64,
20+
) *TeamLimits {
21+
return &TeamLimits{
22+
SandboxConcurrency: tier.ConcurrentInstances + extraMaxSandboxes,
23+
BuildConcurrency: tier.ConcurrentTemplateBuilds + extraMaxConcurrentBuilds,
24+
MaxLengthHours: tier.MaxLengthHours,
25+
26+
MaxVcpu: tier.MaxVcpu + extraMaxVcpu,
27+
MaxRamMb: tier.MaxRamMb + extraMaxRamMb,
28+
DiskMb: tier.DiskMb + addonsDiskMb,
29+
}
30+
}
31+
32+
func NewTeam(
33+
team *queries.Team,
34+
tier *queries.Tier,
35+
extraMaxSandboxes int64,
36+
extraMaxConcurrentBuilds int64,
37+
extraMaxVcpu int64,
38+
extraMaxRamMb int64,
39+
extraDiskMb int64,
40+
) *Team {
41+
return &Team{
42+
Team: team,
43+
Limits: newTeamLimits(
44+
tier,
45+
extraMaxSandboxes,
46+
extraMaxConcurrentBuilds,
47+
extraMaxVcpu,
48+
extraMaxRamMb,
49+
extraDiskMb,
50+
),
51+
}
52+
}
53+
54+
type TeamWithDefault struct {
55+
*Team
56+
57+
IsDefault bool
58+
}

‎packages/api/internal/db/user.go‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,17 @@ import (
66

77
"github.com/google/uuid"
88

9+
"github.com/e2b-dev/infra/packages/api/internal/db/types"
910
sqlcdb "github.com/e2b-dev/infra/packages/db/client"
1011
"github.com/e2b-dev/infra/packages/db/queries"
1112
)
1213

13-
func GetTeamByIDAndUserIDAuth(ctx context.Context, db *sqlcdb.Client, teamID string, userID uuid.UUID) (*queries.Team, *queries.Tier, error) {
14+
func GetTeamByIDAndUserIDAuth(ctx context.Context, db *sqlcdb.Client, teamID string, userID uuid.UUID) (*types.Team, error) {
1415
teamIDParsed, err := uuid.Parse(teamID)
1516
if err != nil {
1617
errMsg := fmt.Errorf("failed to parse team ID: %w", err)
1718

18-
return nil, nil, errMsg
19+
return nil, errMsg
1920
}
2021

2122
result, err := db.GetTeamWithTierByTeamAndUser(ctx, queries.GetTeamWithTierByTeamAndUserParams{
@@ -25,13 +26,22 @@ func GetTeamByIDAndUserIDAuth(ctx context.Context, db *sqlcdb.Client, teamID str
2526
if err != nil {
2627
errMsg := fmt.Errorf("failed to get team from teamID and userID key: %w", err)
2728

28-
return nil, nil, errMsg
29+
return nil, errMsg
2930
}
3031

3132
err = validateTeamUsage(result.Team)
3233
if err != nil {
33-
return nil, nil, err
34+
return nil, err
3435
}
3536

36-
return &result.Team, &result.Tier, nil
37+
team := types.NewTeam(
38+
&result.Team,
39+
&result.Tier,
40+
result.ExtraConcurrentSandboxes,
41+
result.ExtraConcurrentTemplateBuilds,
42+
result.ExtraMaxVcpu,
43+
result.ExtraMaxRamMb,
44+
result.ExtraDiskMb,
45+
)
46+
return team, nil
3747
}

0 commit comments

Comments
 (0)