Skip to content

Commit 4621212

Browse files
gm-e2bgm-e2b
committed
[EN-1038] feat(db): add template disk entitlement schema
Separate the default template free-space target from the maximum total root filesystem size that a team may build. Add nullable default_free_disk_size_mb and max_disk_size_mb columns to tiers, plus extra_max_disk_size_mb to add-ons. Keep extra_disk_mb as the add-on contribution to default free space and add constraints preventing invalid negative or inverted tier values. Keep the existing team_limits view unchanged for current consumers and introduce team_limits_v2 with the new effective default and maximum fields. Preserve compatibility with legacy add-on writers by treating a null extra maximum as equal to extra_disk_mb. Leave data population and runtime adoption to later rollout steps so this commit remains an expand-only schema change. Regenerate the database models and add focused migration coverage for the new columns, the unchanged V1 view, the V2 contract, security invoker behavior, and effective tier-plus-add-on calculations. Part of: EN-1038
1 parent 9174104 commit 4621212

5 files changed

Lines changed: 220 additions & 0 deletions

File tree

‎packages/db/migrations/20260714091414_add_template_disk_entitlements.sql‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
-- +goose Up
2+
-- +goose StatementBegin
3+
4+
ALTER TABLE "public"."tiers"
5+
-- Free-rootfs target in the existing MiB convention.
6+
ADD COLUMN "default_free_disk_size_mb" bigint,
7+
-- Total logical-rootfs ceiling before active add-ons.
8+
ADD COLUMN "max_disk_size_mb" bigint;
9+
10+
ALTER TABLE "public"."addons"
11+
-- Total-ceiling increment in the existing MiB convention.
12+
ADD COLUMN "extra_max_disk_size_mb" bigint;
13+
14+
ALTER TABLE "public"."tiers"
15+
ADD CONSTRAINT "tiers_default_free_disk_size_mb_check"
16+
CHECK (default_free_disk_size_mb >= 0),
17+
ADD CONSTRAINT "tiers_max_disk_size_mb_check"
18+
CHECK (max_disk_size_mb > 0),
19+
ADD CONSTRAINT "tiers_default_free_disk_size_lte_max_check"
20+
CHECK (default_free_disk_size_mb <= max_disk_size_mb);
21+
22+
ALTER TABLE "public"."addons"
23+
ADD CONSTRAINT "addons_extra_max_disk_size_mb_check"
24+
CHECK (extra_max_disk_size_mb >= 0);
25+
26+
-- +goose StatementEnd
27+
28+
-- +goose Down
29+
-- +goose StatementBegin
30+
31+
ALTER TABLE "public"."tiers"
32+
DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_lte_max_check",
33+
DROP CONSTRAINT IF EXISTS "tiers_max_disk_size_mb_check",
34+
DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_mb_check",
35+
DROP COLUMN IF EXISTS "max_disk_size_mb",
36+
DROP COLUMN IF EXISTS "default_free_disk_size_mb";
37+
38+
ALTER TABLE "public"."addons"
39+
DROP CONSTRAINT IF EXISTS "addons_extra_max_disk_size_mb_check",
40+
DROP COLUMN IF EXISTS "extra_max_disk_size_mb";
41+
42+
-- +goose StatementEnd

‎packages/db/migrations/20260714091415_create_team_limits_v2.sql‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
-- +goose Up
2+
-- +goose StatementBegin
3+
4+
CREATE VIEW "public"."team_limits_v2"
5+
WITH (security_invoker=on) AS
6+
SELECT
7+
t.id,
8+
tier.max_length_hours,
9+
(tier.concurrent_instances + a.extra_concurrent_sandboxes) AS concurrent_sandboxes,
10+
(tier.concurrent_template_builds + a.extra_concurrent_template_builds) AS concurrent_template_builds,
11+
(tier.max_vcpu + a.extra_max_vcpu) AS max_vcpu,
12+
(tier.max_ram_mb + a.extra_max_ram_mb) AS max_ram_mb,
13+
(tier.disk_mb + a.extra_disk_mb) AS disk_mb,
14+
(tier.events_ttl_days + a.extra_events_ttl_days) AS events_ttl_days,
15+
(tier.default_free_disk_size_mb + a.extra_disk_mb)::bigint AS default_free_disk_size_mb,
16+
(tier.max_disk_size_mb + a.extra_max_disk_size_mb)::bigint AS max_disk_size_mb
17+
FROM "public"."teams" t
18+
JOIN "public"."tiers" tier ON t.tier = tier.id
19+
LEFT JOIN LATERAL (
20+
SELECT COALESCE(SUM(extra_concurrent_sandboxes), 0)::bigint AS extra_concurrent_sandboxes,
21+
COALESCE(SUM(extra_concurrent_template_builds), 0)::bigint AS extra_concurrent_template_builds,
22+
COALESCE(SUM(extra_max_vcpu), 0)::bigint AS extra_max_vcpu,
23+
COALESCE(SUM(extra_max_ram_mb), 0)::bigint AS extra_max_ram_mb,
24+
COALESCE(SUM(extra_disk_mb), 0)::bigint AS extra_disk_mb,
25+
COALESCE(SUM(extra_events_ttl_days), 0)::bigint AS extra_events_ttl_days,
26+
COALESCE(SUM(COALESCE(extra_max_disk_size_mb, extra_disk_mb)), 0)::bigint AS extra_max_disk_size_mb
27+
FROM "public"."addons" addon
28+
WHERE addon.team_id = t.id
29+
AND addon.valid_from <= now()
30+
AND (addon.valid_to IS NULL OR addon.valid_to > now())
31+
) a ON true;
32+
33+
-- +goose StatementEnd
34+
35+
-- +goose Down
36+
-- +goose StatementBegin
37+
38+
DROP VIEW IF EXISTS "public"."team_limits_v2";
39+
40+
-- +goose StatementEnd

‎packages/db/pkg/tests/disk_entitlements_migration_test.go‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
package tests
2+
3+
import (
4+
"database/sql"
5+
"testing"
6+
7+
"github.com/google/uuid"
8+
"github.com/stretchr/testify/require"
9+
10+
"github.com/e2b-dev/infra/packages/db/pkg/testutils"
11+
testqueries "github.com/e2b-dev/infra/packages/db/pkg/testutils/queries"
12+
)
13+
14+
func TestDiskEntitlementsMigration(t *testing.T) {
15+
t.Parallel()
16+
17+
db := testutils.SetupDatabase(t)
18+
ctx := t.Context()
19+
20+
sqlDB, err := sql.Open("pgx", db.ConnStr())
21+
require.NoError(t, err)
22+
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
23+
24+
var columnCount int64
25+
var nullableWithoutDefaults bool
26+
err = sqlDB.QueryRowContext(ctx, `
27+
SELECT COUNT(*), BOOL_AND(is_nullable = 'YES' AND column_default IS NULL)
28+
FROM information_schema.columns
29+
WHERE table_schema = 'public'
30+
AND (table_name, column_name) IN (
31+
('tiers', 'default_free_disk_size_mb'),
32+
('tiers', 'max_disk_size_mb'),
33+
('addons', 'extra_max_disk_size_mb')
34+
)
35+
`).Scan(&columnCount, &nullableWithoutDefaults)
36+
require.NoError(t, err)
37+
require.Equal(t, int64(3), columnCount)
38+
require.True(t, nullableWithoutDefaults)
39+
40+
var v1Columns, v2Columns string
41+
err = sqlDB.QueryRowContext(ctx, `
42+
SELECT string_agg(column_name, ',' ORDER BY ordinal_position)
43+
FROM information_schema.columns
44+
WHERE table_schema = 'public' AND table_name = 'team_limits'
45+
`).Scan(&v1Columns)
46+
require.NoError(t, err)
47+
err = sqlDB.QueryRowContext(ctx, `
48+
SELECT string_agg(column_name, ',' ORDER BY ordinal_position)
49+
FROM information_schema.columns
50+
WHERE table_schema = 'public' AND table_name = 'team_limits_v2'
51+
`).Scan(&v2Columns)
52+
require.NoError(t, err)
53+
54+
const legacyColumns = "id,max_length_hours,concurrent_sandboxes,concurrent_template_builds," +
55+
"max_vcpu,max_ram_mb,disk_mb,events_ttl_days"
56+
require.Equal(t, legacyColumns, v1Columns)
57+
require.Equal(t, legacyColumns+",default_free_disk_size_mb,max_disk_size_mb", v2Columns)
58+
59+
var securityInvoker bool
60+
err = sqlDB.QueryRowContext(ctx, `
61+
SELECT COALESCE(reloptions @> ARRAY['security_invoker=on']::text[], false)
62+
FROM pg_class
63+
WHERE oid = 'public.team_limits_v2'::regclass
64+
`).Scan(&securityInvoker)
65+
require.NoError(t, err)
66+
require.True(t, securityInvoker)
67+
68+
_, err = sqlDB.ExecContext(ctx, `
69+
INSERT INTO public.tiers (
70+
id, name, disk_mb, concurrent_instances, max_length_hours,
71+
default_free_disk_size_mb, max_disk_size_mb
72+
)
73+
VALUES ('en-1038-test', 'EN-1038 test', 10240, 1, 24, 8000, 30000)
74+
`)
75+
require.NoError(t, err)
76+
77+
teamID := uuid.New()
78+
err = db.TestQueries.InsertTestTeam(ctx, testqueries.InsertTestTeamParams{
79+
ID: teamID,
80+
Name: "EN-1038 migration test",
81+
Tier: "en-1038-test",
82+
Email: "en-1038-migration@example.com",
83+
Slug: "en-1038-migration",
84+
})
85+
require.NoError(t, err)
86+
87+
_, err = sqlDB.ExecContext(ctx, `
88+
INSERT INTO public.addons (
89+
team_id, name, extra_disk_mb, extra_max_disk_size_mb, added_by
90+
)
91+
VALUES ($1, 'EN-1038 test add-on', 1000, 3000,
92+
'00000000-0000-0000-0000-000000000000')
93+
`, teamID)
94+
require.NoError(t, err)
95+
96+
var legacyDisk int64
97+
err = sqlDB.QueryRowContext(ctx, `
98+
SELECT disk_mb FROM public.team_limits WHERE id = $1
99+
`, teamID).Scan(&legacyDisk)
100+
require.NoError(t, err)
101+
require.Equal(t, int64(11240), legacyDisk)
102+
103+
var v2Disk, defaultFree, maximum int64
104+
err = sqlDB.QueryRowContext(ctx, `
105+
SELECT disk_mb, default_free_disk_size_mb, max_disk_size_mb
106+
FROM public.team_limits_v2
107+
WHERE id = $1
108+
`, teamID).Scan(&v2Disk, &defaultFree, &maximum)
109+
require.NoError(t, err)
110+
require.Equal(t, int64(11240), v2Disk)
111+
require.Equal(t, int64(9000), defaultFree)
112+
require.Equal(t, int64(33000), maximum)
113+
}

‎packages/db/pkg/testutils/queries/models.go‎

Lines changed: 16 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/sqlc.yaml‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,15 @@ overrides:
2121
import: "time"
2222
type: "Time"
2323
pointer: true
24+
# sqlc does not infer nullability for these view expressions.
25+
- column: "public.team_limits_v2.default_free_disk_size_mb"
26+
go_type:
27+
type: "int64"
28+
pointer: true
29+
- column: "public.team_limits_v2.max_disk_size_mb"
30+
go_type:
31+
type: "int64"
32+
pointer: true
2433

2534
sql:
2635
- engine: "postgresql"

0 commit comments

Comments
 (0)